Creating Checkpointhandle and connecting to EmbeddingRocksDB

Summary:
X-link: facebookresearch/FBGEMM#1300

Design doc: https://docs.google.com/document/d/149LdAEHOLP7ei4hwVVkAFXGa4N9uLs1J7efxfBZp3dY/edit?tab=t.0#heading=h.49t3yfaqmt54

Context:
We are enabling the usage of rocksDB checkpoint feature in KVTensorWrapper. This allows us to create checkpoints of the embedding tables in SSD. Later, these checkpoints are used by the checkpointing component to create a checkpoint and upload it it to the manifold 

In this diff:
CheckpointHandle: 
It is an entity is responsible for storing the details of the rocksDB Checkpoint. It consists of the file paths to the checkpoint of all shards. When creating KVTensorWrappers, we use the same checkpointHandle object

1. Creating CheckpointHandle class 
2. Adding the definition for CheckpointHandle constructor 
3. Adding CRUD functions for CheckpointHandle 
4. Connecting CheckpointHandle to EmbeddingRocksDB entity

Differential Revision: D75489827

Author: Raahul Kalyaan Jakka

fbgemm_gpu/src/ssd_split_embeddings_cache/ssd_split_table_batched_embeddings.cpp

@@ -18,7 +18,7 @@
 #include "embedding_rocksdb_wrapper.h"
 #include "fbgemm_gpu/split_embeddings_cache/kv_db_cpp_utils.h"
 #include "fbgemm_gpu/utils/ops_utils.h"
-
+#include "rocksdb/utilities/checkpoint.h"
 using namespace at;
 using namespace ssd;
 using namespace kv_mem;
@@ -293,6 +293,37 @@ snapshot_ptr_t SnapshotHandle::get_snapshot_for_shard(size_t shard) const {
   return shard_snapshots_[shard];
 }
 
+CheckpointHandle::CheckpointHandle(
+    EmbeddingRocksDB* db,
+    const std::string& tbe_uuid,
+    const std::string& ckpt_uuid,
+    const std::string& base_path,
+    bool use_default_ssd_path)
+    : db_(db), ckpt_uuid_(ckpt_uuid) {
+  auto num_shards = db->num_shards();
+  CHECK_GT(num_shards, 0);
+  shard_checkpoints_.reserve(num_shards);
+  for (auto shard = 0; shard < num_shards; ++shard) {
+    auto rocksdb_path = kv_db_utils::get_rocksdb_path(
+        base_path, shard, tbe_uuid, use_default_ssd_path);
+    auto checkpoint_shard_dir =
+        kv_db_utils::get_rocksdb_checkpoint_dir(shard, rocksdb_path);
+    kv_db_utils::create_dir(checkpoint_shard_dir);
+    rocksdb::Checkpoint* checkpoint = nullptr;
+    rocksdb::Status s =
+        rocksdb::Checkpoint::Create(db->dbs_[shard].get(), &checkpoint);
+    CHECK(s.ok()) << "ERROR: Checkpoint init for tbe_uuid " << tbe_uuid
+                  << ", db shard " << shard << " failed, " << s.code() << ", "
+                  << s.ToString();
+    std::string checkpoint_shard_path = checkpoint_shard_dir + "/" + ckpt_uuid_;
+    s = checkpoint->CreateCheckpoint(checkpoint_shard_path);
+    CHECK(s.ok()) << "ERROR: Checkpoint creation for tbe_uuid " << tbe_uuid
+                  << ", db shard " << shard << " failed, " << s.code() << ", "
+                  << s.ToString();
+    shard_checkpoints_.push_back(checkpoint_shard_path);
+  }
+}
+
 EmbeddingSnapshotHandleWrapper::EmbeddingSnapshotHandleWrapper(
     const SnapshotHandle* handle,
     std::shared_ptr<EmbeddingRocksDB> db)

fbgemm_gpu/src/ssd_split_embeddings_cache/ssd_table_batched_embeddings.h

@@ -65,6 +65,25 @@ class SnapshotHandle {
   std::vector<snapshot_ptr_t> shard_snapshots_;
 }; // class SnapshotHandle
 
+using checkpoint_path = std::string;
+// @lint-ignore CLANGTIDY cppcoreguidelines-special-member-functions
+class CheckpointHandle {
+ public:
+  explicit CheckpointHandle(
+      EmbeddingRocksDB* db,
+      const std::string& tbe_uuid,
+      const std::string& ckpt_uuid,
+      const std::string& base_path,
+      bool use_default_ssd_path);
+
+ private:
+  friend class EmbeddingRocksDB;
+
+  EmbeddingRocksDB* db_;
+  std::string ckpt_uuid_;
+  std::vector<checkpoint_path> shard_checkpoints_;
+}; // class CheckpointHandle
+
 /// @ingroup embedding-ssd
 ///
 /// @brief An implementation of EmbeddingKVDB for RocksDB
@@ -488,6 +507,10 @@ class EmbeddingRocksDB : public kv_db::EmbeddingKVDB {
     return snapshots_.find(snapshot_handle) != snapshots_.end();
   }
 
+  bool is_valid_checkpoint(const std::string ckpt_uuid) const {
+    return checkpoints_.find(ckpt_uuid) != checkpoints_.end();
+  }
+
   int64_t get_snapshot_count() const {
     return snapshots_.size();
   }
@@ -505,6 +528,58 @@ class EmbeddingRocksDB : public kv_db::EmbeddingKVDB {
     return handlePtr;
   }
 
+  std::string create_checkpoint(int64_t global_step) {
+    const auto num_ckpts = checkpoints_.size();
+    if (num_ckpts > 0) {
+      std::cerr << "WARNING: rocksdb create_checkpoint found " << num_ckpts
+                << " other checkpoints_" << std::endl;
+    }
+
+    // Removing the CHECK needs extra caution, basically we need to find a way
+    // to consistently link a KVTensor to a checkpoint. Ideally by uuid, but
+    // create_checkpoint is called earlier than state_dict and KVTensor is
+    // created inside state_dict, additionally we don't want to pass the uuid
+    // into state_dict, unless we link it manually on the caller side inside the
+    // state_dict loop, it would be random linking. For example, Case1 and Case2
+    // could call state_dict in the same batch at any order with random times,
+    // it is impossible to correctly attach each checkpoint to each KVTensor.
+    CHECK(
+        global_step_to_ckpt_uuid_.find(global_step) ==
+        global_step_to_ckpt_uuid_.end())
+        << "multiple rdb checkpoint in one global step isn't supported right now";
+    auto ckpt_uuid = facebook::strings::generateUUID();
+    auto handle = std::make_unique<CheckpointHandle>(
+        this, tbe_uuid_, ckpt_uuid, path_, use_default_ssd_path_);
+    checkpoints_[ckpt_uuid] = std::move(handle);
+    global_step_to_ckpt_uuid_[global_step] = ckpt_uuid;
+    return ckpt_uuid;
+  }
+
+  std::optional<std::string> get_active_checkpoint_uuid(int64_t global_step) {
+    if (global_step_to_ckpt_uuid_.find(global_step) !=
+        global_step_to_ckpt_uuid_.end()) {
+      return std::make_optional<std::string>(
+          global_step_to_ckpt_uuid_[global_step]);
+    }
+    return std::nullopt;
+  }
+
+  void release_checkpoint(const std::string ckpt_uuid) {
+    CHECK(is_valid_checkpoint(ckpt_uuid));
+    LOG(INFO) << "Checkpoint " << ckpt_uuid << " released";
+    checkpoints_.erase(ckpt_uuid);
+    // sweep through global_step_to_ckpt_uuid_, it should be small
+    int64_t glb_step_to_purge = -1;
+    for (const auto& [global_step, uuid] : global_step_to_ckpt_uuid_) {
+      if (ckpt_uuid == uuid) {
+        glb_step_to_purge = global_step;
+        break;
+      }
+    }
+    CHECK(glb_step_to_purge != -1) << "There must be a rdb ckpt uuid to purge";
+    global_step_to_ckpt_uuid_.erase(glb_step_to_purge);
+  }
+
   void release_snapshot(const SnapshotHandle* snapshot_handle) {
     CHECK(is_valid_snapshot(snapshot_handle));
     LOG(INFO) << "Snapshot " << snapshot_handle << " released";
@@ -1117,6 +1192,7 @@ class EmbeddingRocksDB : public kv_db::EmbeddingKVDB {
   }
 
   friend class SnapshotHandle;
+  friend class CheckpointHandle;
 
   std::vector<std::unique_ptr<rocksdb::DB>> dbs_;
   std::vector<std::unique_ptr<Initializer>> initializers_;
@@ -1149,6 +1225,22 @@ class EmbeddingRocksDB : public kv_db::EmbeddingKVDB {
   std::string tbe_uuid_;
   std::string path_;
   bool use_default_ssd_path_;
+
+  // rocksdb checkpoint is used to create an on disk database to support
+  // cross process read-only access, currently this is only needed for publish
+  // in async checkpoint, SSD read is handled in the same proc
+  std::unordered_map<std::string, std::unique_ptr<CheckpointHandle>>
+      checkpoints_;
+  // this is used to link KVTensor to the corresponding rdb checkpoint by global
+  // step reasons are shown below
+  // 1. rdb checkpoint is only needed during publish not checkpoint, therefore
+  // only some state_dict() calls need it
+  // 2. we dont' want to add new argument in state_dict to add complexity
+  // 3. therefore create_checkpoint is exposed up to the publish component
+  // 4. publish will call create_checkpoint() then state_dict()
+  // 5. need a way to link rdb ckpt and KVTensor(from state_dict()) together so
+  // that caller doesn't need to do the link themselves
+  std::unordered_map<int64_t, std::string> global_step_to_ckpt_uuid_;
 }; // class EmbeddingRocksDB
 
 } // namespace ssd