Adding helper function for enabling RocksDB Checkpoint (#4213)

Summary:

X-link: facebookresearch/FBGEMM#1288

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:

We are adding all the helper functions to create folder structures in SSDs for storing the rocksDB checkpoints

Reviewed By: duduyi2013

Differential Revision: D75489819

Author: Raahul Kalyaan Jakka

fbgemm_gpu/include/fbgemm_gpu/split_embeddings_cache/kv_db_cpp_utils.h

@@ -10,15 +10,20 @@
 
 #include <ATen/ATen.h>
 #include <folly/hash/Hash.h>
+#include <glog/logging.h>
 #include <stddef.h>
 #include <stdint.h>
+#include <filesystem>
 #include <optional>
-
 /// @defgroup embedding-ssd Embedding SSD Operators
 ///
 
 namespace kv_db_utils {
 
+#ifdef FBGEMM_FBCODE
+constexpr size_t num_ssd_drives = 8;
+#endif
+
 /// @ingroup embedding-ssd
 ///
 /// @brief hash function used for SSD L2 cache and rocksdb sharding algorithm
@@ -65,4 +70,94 @@ std::tuple<at::Tensor, at::Tensor> get_bucket_sorted_indices_and_bucket_tensor(
     std::optional<int64_t> bucket_size,
     std::optional<int64_t> total_num_buckets);
 
+/// @ingroup embedding-ssd
+///
+/// @brief default way to generate rocksdb path based on a user provided
+/// base_path the file hierarchy will be
+///   <base_path><ssd_idx>/<tbe_uuid> for default SSD mount
+///   <base_path>/<tbe_uuid> for user provided base path
+///
+/// @param base_path the base path for all the rocksdb shards tied to one
+/// TBE/EmbeddingRocksDB
+/// @param db_shard_id the rocksdb shard index, this is used to determine which
+/// SSD to use
+/// @param tbe_uuid unique identifier per TBE at the lifetime of a training job
+/// @param default_path whether the base_path is default SSD mount or
+/// user-provided
+///
+/// @return the base path to that rocksdb shard
+inline std::string get_rocksdb_path(
+    const std::string& base_path,
+    int db_shard_id,
+    const std::string& tbe_uuid,
+    bool default_path) {
+  if (default_path) {
+    int ssd_drive_idx = db_shard_id % num_ssd_drives;
+    std::string ssd_idx_tbe_id_str =
+        std::to_string(ssd_drive_idx) + std::string("/") + tbe_uuid;
+    return base_path + ssd_idx_tbe_id_str;
+  } else {
+    return base_path + std::string("/") + tbe_uuid;
+  }
+}
+
+/// @ingroup embedding-ssd
+///
+/// @brief generate rocksdb shard path, based on rocksdb_path
+/// the file hierarchy will be
+///   <rocksdb_shard_path>/shard_<db_shard>
+///
+/// @param db_shard_id the rocksdb shard index
+/// @param rocksdb_path the base path for rocksdb shard
+///
+/// @return the rocksdb shard path
+inline std::string get_rocksdb_shard_path(
+    int db_shard_id,
+    const std::string& rocksdb_path) {
+  return rocksdb_path + std::string("/shard_") + std::to_string(db_shard_id);
+}
+
+/// @ingroup embedding-ssd
+///
+/// @brief generate a directory to hold rocksdb checkpoint for a particular
+/// rocksdb shard path the file hierarchy will be
+///   <rocksdb_shard_path>/checkpoint_shard_<db_shard>
+///
+/// @param db_shard_id the rocksdb shard index
+/// @param rocksdb_path the base path for rocksdb shard
+///
+/// @return the directory that holds rocksdb checkpoints for one rocksdb shard
+inline std::string get_rocksdb_checkpoint_dir(
+    int db_shard_id,
+    const std::string& rocksdb_path) {
+  return rocksdb_path + std::string("/checkpoint_shard_") +
+      std::to_string(db_shard_id);
+}
+
+inline void create_dir(const std::string& dir_path) {
+  try {
+    std::filesystem::path fs_path(dir_path);
+    bool res = std::filesystem::create_directories(fs_path);
+    if (!res) {
+      LOG(ERROR) << "dir: " << dir_path << " already exists";
+    }
+  } catch (const std::exception& e) {
+    LOG(ERROR) << "Error creating directory: " << e.what();
+  }
+}
+
+inline void remove_dir(const std::string& path) {
+  if (std::filesystem::exists(path)) {
+    try {
+      if (std::filesystem::is_directory(path)) {
+        std::filesystem::remove_all(path);
+      } else {
+        std::filesystem::remove(path);
+      }
+    } catch (const std::filesystem::filesystem_error& e) {
+      LOG(ERROR) << "Error removing path: " << path
+                 << ", exception:" << e.what();
+    }
+  }
+}
 }; // namespace kv_db_utils

fbgemm_gpu/src/ssd_split_embeddings_cache/ssd_table_batched_embeddings.h

@@ -33,11 +33,7 @@ namespace ssd {
 
 using namespace at;
 
-#ifdef FBGEMM_FBCODE
-constexpr size_t num_ssd_drives = 8;
 const std::string ssd_mount_point = "/data00_nvidia";
-const size_t base_port = 136000;
-#endif
 
 // mem usage propertiese
 // -- block cache usage
@@ -320,23 +316,22 @@ class EmbeddingRocksDB : public kv_db::EmbeddingKVDB {
     auto db_monitor_options = facebook::fb_rocksdb::DBMonitorOptions();
     db_monitor_options.fb303Prefix = "tbe_metrics";
 
-    std::string tbe_uuid = "";
+    tbe_uuid_ = facebook::strings::generateUUID();
+    use_default_ssd_path_ = !use_passed_in_path;
     if (!use_passed_in_path) {
-      path = ssd_mount_point;
-      tbe_uuid = facebook::strings::generateUUID();
+      path_ = std::move(ssd_mount_point);
+    } else {
+      path_ = std::move(path);
     }
+    std::string all_shards_path;
 #endif
     for (auto i = 0; i < num_shards; ++i) {
 #ifdef FBGEMM_FBCODE
-      int ssd_drive_idx = i % num_ssd_drives;
-      std::string ssd_idx_tbe_id_str = "";
-      if (!use_passed_in_path) {
-        ssd_idx_tbe_id_str =
-            std::to_string(ssd_drive_idx) + std::string("/") + tbe_uuid;
-      }
-      auto shard_path =
-          path + ssd_idx_tbe_id_str + std::string("_shard") + std::to_string(i);
-      used_path += shard_path + ", ";
+      auto rocksdb_path = kv_db_utils::get_rocksdb_path(
+          path_, i, tbe_uuid_, !use_passed_in_path);
+      auto shard_path = kv_db_utils::get_rocksdb_shard_path(i, rocksdb_path);
+      kv_db_utils::create_dir(shard_path);
+      all_shards_path += shard_path + ", ";
 #else
       auto shard_path = path + std::string("/shard_") + std::to_string(i);
 #endif
@@ -369,7 +364,8 @@ class EmbeddingRocksDB : public kv_db::EmbeddingKVDB {
       dbs_.emplace_back(db);
     }
 #ifdef FBGEMM_FBCODE
-    LOG(INFO) << "TBE actual used_path: " << used_path;
+    LOG(INFO) << "TBE uuid: " << tbe_uuid_
+              << ", rocksdb shards paths: " << all_shards_path;
 #endif
   }
 
@@ -1150,6 +1146,9 @@ class EmbeddingRocksDB : public kv_db::EmbeddingKVDB {
   int64_t elem_size_;
   std::vector<int64_t> sub_table_dims_;
   std::vector<int64_t> sub_table_hash_cumsum_;
+  std::string tbe_uuid_;
+  std::string path_;
+  bool use_default_ssd_path_;
 }; // class EmbeddingRocksDB
 
 } // namespace ssd

fbgemm_gpu/test/tbe/ssd/embedding_cache/rocksdb_embedding_cache_test.cpp

@@ -38,7 +38,8 @@ TEST(RocksDbEmbeddingCacheTest, TestPutAndGet) {
       -0.01, // uniform_init_lower,
       0.01, // uniform_init_upper,
       32, // row_storage_bitwidth = 32,
-      0 // cache_size = 0
+      0, // cache_size = 0
+      true // use_passed_in_path
   );
 
   auto write_indices =