[Embedding] Add interface of EmbeddingVar for Elastic Training. (#933)

Signed-off-by: JunqiHu <[email protected]>

Author: Mesilenceki

configure.py

@@ -1434,7 +1434,7 @@ def main():
                 True, 'star')
 
   set_build_var(environ_cp, 'TF_NEED_ELASTIC', 'ELASTIC TRAINING', 'with_elastic_support',
-                True, 'elastic')
+                False, 'elastic')
 
   set_build_var(environ_cp, 'TF_ENABLE_PMEM', 'PMEM', 'with_pmem_support',
                 False, 'pmem')

tensorflow/contrib/elastic_grpc_server/BUILD

@@ -56,7 +56,8 @@ cc_library(
 tf_cc_test(
     name = "elastic_grpc_test",
     size = "small",
-    srcs = ["elastic_grpc_server_lib_test.cc"],
+    srcs = select({"//tensorflow:with_elastic_support": ["elastic_grpc_server_lib_test.cc"],
+		   "//conditions:default": []}),
     deps = [
         ":elastic_grpc_server_lib",
         "//tensorflow/core/distributed_runtime/rpc:grpc_util",

tensorflow/core/BUILD

@@ -128,6 +128,7 @@ load(
     "tf_additional_numa_deps",
     "tf_additional_numa_lib_defines",
     "tf_additional_star_lib_defines",
+    "tf_additional_elastic_server_lib_defines",
     "tf_additional_api_compatible_defines",
     "tf_additional_pmem_lib_defines",
     "tf_additional_test_deps",
@@ -1441,6 +1442,7 @@ tf_cc_test(
 cc_library(
     name = "ops",
     visibility = ["//visibility:public"],
+    defines = tf_additional_elastic_server_lib_defines(),
     deps = [
         ":array_ops_op_lib",
         ":parquet_ops_op_lib",
@@ -2562,7 +2564,8 @@ LIB_INTERNAL_DEFINES = (
     tf_additional_gdr_lib_defines() +
     tf_additional_numa_lib_defines() +
     tf_additional_star_lib_defines() +
-    tf_additional_pmem_lib_defines()
+    tf_additional_pmem_lib_defines() +
+    tf_additional_elastic_server_lib_defines()
 )
 
 cc_library(

tensorflow/core/framework/embedding/bloom_filter_policy.h

@@ -333,7 +333,7 @@ class BloomFilterPolicy : public FilterPolicy<K, V, EV> {
       // this can describe by graph(Mod + DynamicPartition),
       // but memory waste and slow
       if (*(key_buff + i) % bucket_num % partition_num != partition_id) {
-        LOG(INFO) << "skip EV key:" << *(key_buff + i);
+        VLOG(1) << "skip EV key:" << *(key_buff + i);
         continue;
       }
       void* value_ptr = nullptr;

tensorflow/core/framework/embedding/counter_filter_policy.h

@@ -159,7 +159,7 @@ class CounterFilterPolicy : public FilterPolicy<K, V, EV> {
       // this can describe by graph(Mod + DynamicPartition),
       // but memory waste and slow
       if (*(key_buff + i) % bucket_num % partition_num != partition_id) {
-        LOG(INFO) << "skip EV key:" << *(key_buff + i);
+        VLOG(1) << "skip EV key:" << *(key_buff + i);
         continue;
       }
       int64 import_freq = 0;

tensorflow/core/framework/embedding/cpu_hash_map_kv.h

@@ -137,6 +137,28 @@ class LocklessHashMap : public KVInterface<K, V> {
     return Status::OK();
   }
 
+  Status GetShardedSnapshot(
+      std::vector<K>* key_list, std::vector<void*>* value_ptr_list,
+      int partition_id, int partition_nums) override {
+    std::pair<const K, void*> *hash_map_dump;
+    int64 bucket_count;
+    auto it = hash_map_.GetSnapshot();
+    hash_map_dump = it.first;
+    bucket_count = it.second;
+    for (int64 j = 0; j < bucket_count; j++) {
+      if (hash_map_dump[j].first != LocklessHashMap<K, V>::EMPTY_KEY_ 
+          && hash_map_dump[j].first != LocklessHashMap<K, V>::DELETED_KEY_
+          && hash_map_dump[j].first % kSavedPartitionNum 
+              % partition_nums != partition_id) {
+        key_list->emplace_back(hash_map_dump[j].first);
+        value_ptr_list->emplace_back(hash_map_dump[j].second);
+      }
+    }
+
+    free(hash_map_dump);
+    return Status::OK();
+  }
+
   std::string DebugString() const override {
     LOG(INFO) << "map info size:" << Size()
               << "map info bucket_count:" << hash_map_.bucket_count()

tensorflow/core/framework/embedding/dense_hash_map_kv.h

@@ -121,6 +121,25 @@ class DenseHashMap : public KVInterface<K, V> {
     return Status::OK();
   }
 
+  Status GetShardedSnapshot(
+      std::vector<K>* key_list, std::vector<void*>* value_ptr_list,
+      int partition_id, int partition_nums) override {
+    dense_hash_map hash_map_dump[partition_num_];
+    for (int i = 0; i< partition_num_; i++) {
+      spin_rd_lock l(hash_map_[i].mu);
+      hash_map_dump[i].hash_map = hash_map_[i].hash_map;
+    }
+    for (int i = 0; i< partition_num_; i++) {
+      for (const auto it : hash_map_dump[i].hash_map) {
+        if (it.first % kSavedPartitionNum % partition_nums != partition_id) {
+          key_list->push_back(it.first);
+          value_ptr_list->push_back(it.second);
+        }
+      }
+    }
+    return Status::OK();
+  }
+
   std::string DebugString() const override {
     return "";
   }

tensorflow/core/framework/embedding/embedding_var.h

@@ -435,6 +435,10 @@ class EmbeddingVar : public ResourceBase {
     return storage_->CacheSize();
   }
 
+  int64 MemoryUsage() const {
+    return storage_->Size() * (sizeof(K) + feat_desc_->data_bytes());
+  }
+
   int64 MinFreq() {
     return emb_config_.filter_freq;
   }
@@ -516,6 +520,85 @@ class EmbeddingVar : public ResourceBase {
     }
   }
 
+  Status GetShardedSnapshot(std::vector<K>* key_list,
+                            std::vector<void*>* value_ptr_list,
+                            int partition_id, int partition_num) {
+    return storage_->GetShardedSnapshot(key_list, value_ptr_list,
+                                        partition_id, partition_num);
+  }
+
+  void ExportAndRemove(K* key_list, V* value_list,
+                     int64* version_list, int64* freq_list,
+                     std::vector<K>& tot_keys_list,
+                     std::vector<void*>& tot_value_ptr_list) {
+    bool save_unfiltered_features = true;
+    TF_CHECK_OK(ReadBoolFromEnvVar(
+        "TF_EV_SAVE_FILTERED_FEATURES", true, &save_unfiltered_features));
+
+    bool is_save_freq = emb_config_.is_save_freq();
+    bool is_save_version = emb_config_.is_save_version();
+
+    for (int64 i = 0; i < tot_keys_list.size(); ++i) {
+      auto& value_ptr = tot_value_ptr_list[i];
+      if((int64)value_ptr == embedding::ValuePtrStatus::IS_DELETED)
+        continue;
+
+      bool is_admit = feat_desc_->IsAdmit(value_ptr);
+      bool is_in_dram = ((int64)value_ptr >> kDramFlagOffset == 0);
+
+      if (!is_admit) {
+        key_list[i] = tot_keys_list[i];
+        
+        if (!is_in_dram) {
+          auto tmp_value = value_list + i * value_len_;
+          tmp_value = (V*)embedding::ValuePtrStatus::NOT_IN_DRAM;
+          value_ptr = (void*)((int64)value_ptr & ((1L << kDramFlagOffset) - 1));
+        } else if (feat_desc_->GetEmbedding(value_ptr, 0) == nullptr) {
+          memcpy(value_list + i * value_len_, default_value_, sizeof(V) * value_len_);
+        } else {
+          V* val = feat_desc_->GetEmbedding(value_ptr, emb_config_.emb_index);
+          memcpy(value_list + i * value_len_, val, sizeof(V) * value_len_);
+        }
+
+        if(is_save_version) {
+          int64 dump_version = feat_desc_->GetVersion(value_ptr);
+          version_list[i] = dump_version;
+        }
+
+        if(is_save_freq) {
+          int64 dump_freq = feat_desc_->GetFreq(value_ptr);
+          freq_list[i] = dump_freq;
+        }
+      } else {
+        if (!save_unfiltered_features)
+          return;
+        //TODO(JUNQI) : currently not export filtered keys
+      }
+
+      if (emb_config_.is_primary()) {
+        Status s;
+        s = storage_->Remove(tot_keys_list[i]);
+        if (!s.ok()) {
+          LOG(ERROR) << "Remove keys error: " << s.error_message();
+        }
+        feat_desc_->Deallocate(value_ptr);
+      }
+    }
+  }
+
+  Status RestoreFromKeysAndValues(int64 key_num, int partition_id,
+                                  int partition_num, const K* key_list,
+                                  const V* value_list, const int64* version_list,
+                                  const int64* freq_list,
+                                  const Eigen::GpuDevice* device = nullptr) {
+    RestoreBuffer restore_buff((char*)key_list, (char*)value_list,
+                                (char*)version_list, (char*)freq_list);
+    return storage_->RestoreFeatures(key_num, kSavedPartitionNum, 
+                                     partition_id, partition_num,
+                                     value_len_, false/* is_filter*/, false/* is_incr*/,
+                                     emb_config_, device, filter_, restore_buff);
+  }
+
   mutex* mu() {
     return &mu_;
   }
@@ -537,6 +620,8 @@ class EmbeddingVar : public ResourceBase {
     }
   }
 
+  string Name() {return name_; }
+
   V* GetDefaultValuePtr() {
     return default_value_;
   }
@@ -645,7 +730,6 @@ class EmbeddingVar : public ResourceBase {
   GPUHashTable<K, V>* HashTable() {
     return storage_->HashTable();
   }
-
   FilterPolicy<K, V, EmbeddingVar<K, V>>* GetFilter() const {
     return filter_;
   }

tensorflow/core/framework/embedding/embedding_var_ckpt_data.h

@@ -20,7 +20,6 @@ limitations under the License.
 namespace tensorflow {
 class BundleWriter;
 namespace {
-  const int kSavedPartitionNum = 1000;
   const int kDramFlagOffset = 49;
 }
 

tensorflow/core/framework/embedding/filter_policy.h

@@ -27,19 +27,31 @@ struct RestoreBuffer {
   char* value_buffer = nullptr;
   char* version_buffer = nullptr;
   char* freq_buffer = nullptr;
+  bool should_release = false;
 
   explicit RestoreBuffer(size_t buffer_size) {
     key_buffer = new char[buffer_size];
     value_buffer = new char[buffer_size];
     version_buffer = new char[buffer_size];
     freq_buffer = new char[buffer_size];
+    should_release = true;
+  }
+
+  explicit RestoreBuffer(char* i_key_buffer, char* i_value_buffer,
+                         char* i_version_buffer, char* i_freq_buffer) {
+    key_buffer = i_key_buffer;
+    value_buffer = i_value_buffer;
+    version_buffer = i_version_buffer;
+    freq_buffer = i_freq_buffer;
   }
 
   ~RestoreBuffer() {
-    delete []key_buffer;
-    delete []value_buffer;
-    delete []version_buffer;
-    delete []freq_buffer;
+    if (should_release) {
+      delete []key_buffer;
+      delete []value_buffer;
+      delete []version_buffer;
+      delete []freq_buffer;
+    }
   }
 };