Adding python api to support sync trigger evict (#4984)

Summary:
Pull Request resolved: #4984

X-link: https://github.com/facebookresearch/FBGEMM/pull/1997

As title, `has_running_evict` and `trigger_feature_evict` are needed to support sync trigger eviction

Reviewed By: kathyxuyy

Differential Revision: D83896308

fbshipit-source-id: 2c68a691ff66ca68c225528cdc7a8c7d50aab516

Author: EddyLXJ

fbgemm_gpu/src/dram_kv_embedding_cache/dram_kv_embedding_cache.h

@@ -1175,17 +1175,8 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
 
   void compact() override {}
 
-  void trigger_feature_evict(
-      std::optional<uint32_t> inplace_update_ts = std::nullopt) {
+  void trigger_feature_evict() {
     if (feature_evict_) {
-      if (inplace_update_ts.has_value() &&
-          feature_evict_config_.value()->trigger_strategy_ ==
-              EvictTriggerStrategy::BY_TIMESTAMP_THRESHOLD) {
-        auto* tt_evict = dynamic_cast<TimeThresholdBasedEvict<weight_type>*>(
-            feature_evict_.get());
-        CHECK(tt_evict != nullptr);
-        tt_evict->set_eviction_timestamp_threshold(inplace_update_ts.value());
-      }
       feature_evict_->trigger_evict();
     }
   }
@@ -1269,6 +1260,13 @@ class DramKVEmbeddingCache : public kv_db::EmbeddingKVDB {
     }
   }
 
+  bool is_evicting() override {
+    if (feature_evict_) {
+      return feature_evict_->is_evicting();
+    }
+    return false;
+  }
+
   // for inference only, this logs the total hit/miss count
   // this should be called at the end of full/delta snapshot chunk by chunk
   // update

fbgemm_gpu/src/dram_kv_embedding_cache/dram_kv_embedding_cache_wrapper.h

@@ -179,6 +179,14 @@ class DramKVEmbeddingCacheWrapper : public torch::jit::CustomClassHolder {
     impl_->set_backend_return_whole_row(backend_return_whole_row);
   }
 
+  void trigger_feature_evict() {
+    impl_->trigger_feature_evict();
+  }
+
+  bool is_evicting() {
+    return impl_->is_evicting();
+  }
+
   void set_feature_score_metadata_cuda(
       at::Tensor indices,
       at::Tensor count,

fbgemm_gpu/src/ssd_split_embeddings_cache/embedding_rocksdb_wrapper.h

@@ -236,6 +236,14 @@ class EmbeddingRocksDBWrapper : public torch::jit::CustomClassHolder {
     impl_->set_backend_return_whole_row(backend_return_whole_row);
   }
 
+  void trigger_feature_evict() {
+    impl_->trigger_feature_evict();
+  }
+
+  bool is_evicting() {
+    return impl_->is_evicting();
+  }
+
  private:
   friend class KVTensorWrapper;
 

fbgemm_gpu/src/ssd_split_embeddings_cache/kv_db_table_batched_embeddings.cpp

@@ -383,6 +383,14 @@ void EmbeddingKVDB::set_backend_return_whole_row(
   return;
 }
 
+void EmbeddingKVDB::trigger_feature_evict() {
+  return;
+}
+
+bool EmbeddingKVDB::is_evicting() {
+  return false;
+}
+
 void EmbeddingKVDB::set(
     const at::Tensor& indices,
     const at::Tensor& weights,

fbgemm_gpu/src/ssd_split_embeddings_cache/kv_db_table_batched_embeddings.h

@@ -301,6 +301,10 @@ class EmbeddingKVDB : public std::enable_shared_from_this<EmbeddingKVDB> {
     FBEXCEPTION("Not implemented");
   }
 
+  virtual void trigger_feature_evict();
+
+  virtual bool is_evicting();
+
   /**
    * @brief need to support set backend_return_whole_row from frontend
    * if one model changed from SSD to DRAM, or vice versa we need to

fbgemm_gpu/src/ssd_split_embeddings_cache/ssd_split_table_batched_embeddings.cpp

@@ -880,6 +880,10 @@ static auto embedding_rocks_db_wrapper =
             {
                 torch::arg("backend_return_whole_row"),
             })
+        .def(
+            "trigger_feature_evict",
+            &EmbeddingRocksDBWrapper::trigger_feature_evict)
+        .def("is_evicting", &EmbeddingRocksDBWrapper::is_evicting)
         .def("stream_sync_cuda", &EmbeddingRocksDBWrapper::stream_sync_cuda)
         .def("get_cuda", &EmbeddingRocksDBWrapper::get_cuda)
         .def("compact", &EmbeddingRocksDBWrapper::compact)
@@ -980,6 +984,10 @@ static auto dram_kv_embedding_cache_wrapper =
             {
                 torch::arg("backend_return_whole_row"),
             })
+        .def(
+            "trigger_feature_evict",
+            &DramKVEmbeddingCacheWrapper::trigger_feature_evict)
+        .def("is_evicting", &DramKVEmbeddingCacheWrapper::is_evicting)
         .def("set", &DramKVEmbeddingCacheWrapper::set)
         .def(
             "set_range_to_storage",