add enable_async_update into tbe signature (pytorch#518)

Summary:
X-link: pytorch#3461

X-link: pytorch/torchrec#2599

X-link: pytorch#3431

Pull Request resolved: facebookresearch/FBGEMM#518

ad eanble_async_update into tbe config

Reviewed By: chrisxcai

Differential Revision: D66802199

fbshipit-source-id: 56d83586396f293b5932f5b6ad7661493278dd8c

Author: duduyi2013

fbgemm_gpu/fbgemm_gpu/tbe/ssd/training.py

@@ -145,6 +145,7 @@ def __init__(
         # Set to True to alloc a UVM tensor using malloc+cudaHostRegister.
         # Set to False to use cudaMallocManaged
         uvm_host_mapped: bool = False,
+        enable_async_update: bool = True,  # whether enable L2/rocksdb write to async background thread
     ) -> None:
         super(SSDTableBatchedEmbeddingBags, self).__init__()
 
@@ -427,7 +428,7 @@ def __init__(
         logging.info(f"tbe_unique_id: {tbe_unique_id}")
         if not ps_hosts:
             logging.info(
-                f"Logging SSD offloading setup, tbe_unique_id:{tbe_unique_id}, l2_cache_size:{l2_cache_size}GB, "
+                f"Logging SSD offloading setup, tbe_unique_id:{tbe_unique_id}, l2_cache_size:{l2_cache_size}GB, enable_async_update:{enable_async_update}"
                 f"passed_in_path={ssd_directory}, num_shards={ssd_rocksdb_shards},num_threads={ssd_rocksdb_shards},"
                 f"memtable_flush_period={ssd_memtable_flush_period},memtable_flush_offset={ssd_memtable_flush_offset},"
                 f"l0_files_per_compact={ssd_l0_files_per_compact},max_D={self.max_D},rate_limit_mbps={ssd_rate_limit_mbps},"
@@ -459,6 +460,7 @@ def __init__(
                 use_passed_in_path,
                 tbe_unique_id,
                 l2_cache_size,
+                enable_async_update,
             )
         else:
             # pyre-fixme[4]: Attribute must be annotated.

fbgemm_gpu/src/ssd_split_embeddings_cache/ssd_split_table_batched_embeddings.cpp

@@ -279,7 +279,8 @@ class EmbeddingRocksDBWrapper : public torch::jit::CustomClassHolder {
       int64_t cache_size = 0,
       bool use_passed_in_path = false,
       int64_t tbe_unique_id = 0,
-      int64_t l2_cache_size_gb = 0)
+      int64_t l2_cache_size_gb = 0,
+      bool enable_async_update = false)
       : impl_(std::make_shared<ssd::EmbeddingRocksDB>(
             path,
             num_shards,
@@ -299,7 +300,8 @@ class EmbeddingRocksDBWrapper : public torch::jit::CustomClassHolder {
             cache_size,
             use_passed_in_path,
             tbe_unique_id,
-            l2_cache_size_gb)) {}
+            l2_cache_size_gb,
+            enable_async_update)) {}
 
   void set_cuda(
       Tensor indices,
@@ -481,7 +483,8 @@ static auto embedding_rocks_db_wrapper =
                 int64_t,
                 bool,
                 int64_t,
-                int64_t>(),
+                int64_t,
+                bool>(),
             "",
             {
                 torch::arg("path"),
@@ -503,6 +506,7 @@ static auto embedding_rocks_db_wrapper =
                 torch::arg("use_passed_in_path") = true,
                 torch::arg("tbe_unique_id") = 0,
                 torch::arg("l2_cache_size_gb") = 0,
+                torch::arg("enable_async_update") = true,
             })
         .def(
             "set_cuda",

fbgemm_gpu/src/ssd_split_embeddings_cache/ssd_table_batched_embeddings.h

@@ -192,13 +192,15 @@ class EmbeddingRocksDB : public kv_db::EmbeddingKVDB {
       int64_t cache_size = 0,
       bool use_passed_in_path = false,
       int64_t tbe_unqiue_id = 0,
-      int64_t l2_cache_size_gb = 0)
+      int64_t l2_cache_size_gb = 0,
+      bool enable_async_update = false)
       : kv_db::EmbeddingKVDB(
             num_shards,
             max_D,
             l2_cache_size_gb,
             tbe_unqiue_id,
-            row_storage_bitwidth / 8),
+            row_storage_bitwidth / 8,
+            enable_async_update),
         max_D_(max_D) {
     // TODO: lots of tunables. NNI or something for this?
     rocksdb::Options options;