Adding function to create a snapshot and exposing it from EmbeddingRocksDBWrapper (pytorch#4223)

Summary:
X-link: pytorch/torchrec#3024

Pull Request resolved: pytorch#4223

X-link: facebookresearch/FBGEMM#1299

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:

Creating a function to create a checkpoint and exposing it to EmbeddingRocksDBWrapper

Reviewed By: duduyi2013

Differential Revision: D75489841

Author: Raahul Kalyaan Jakka

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -2490,6 +2490,12 @@ def flush(self, force: bool = False) -> None:
         self.ssd_db.flush()
         self.last_flush_step = self.step
 
+    def create_rocksdb_hard_link_snapshot(self) -> None:
+        """
+        Create a rocksdb hard link snapshot to provide cross procs access to the underlying data
+        """
+        self.ssd_db.create_rocksdb_hard_link_snapshot(self.step)
+
     def prepare_inputs(
         self,
         indices: Tensor,

fbgemm_gpu/src/ssd_split_embeddings_cache/embedding_rocksdb_wrapper.h

@@ -176,6 +176,10 @@ class EmbeddingRocksDBWrapper : public torch::jit::CustomClassHolder {
     return impl_->get_snapshot_count();
   }
 
+  void create_rocksdb_hard_link_snapshot(int64_t global_step) {
+    impl_->create_checkpoint(global_step);
+  }
+
  private:
   friend class KVTensorWrapper;
 

fbgemm_gpu/src/ssd_split_embeddings_cache/ssd_split_table_batched_embeddings.cpp

@@ -611,7 +611,10 @@ static auto embedding_rocks_db_wrapper =
         .def("get_snapshot_count", &EmbeddingRocksDBWrapper::get_snapshot_count)
         .def(
             "get_keys_in_range_by_snapshot",
-            &EmbeddingRocksDBWrapper::get_keys_in_range_by_snapshot);
+            &EmbeddingRocksDBWrapper::get_keys_in_range_by_snapshot)
+        .def(
+            "create_rocksdb_hard_link_snapshot",
+            &EmbeddingRocksDBWrapper::create_rocksdb_hard_link_snapshot);
 
 static auto dram_kv_embedding_cache_wrapper =
     torch::class_<DramKVEmbeddingCacheWrapper>(