Adding returnKVTensorMetaData flag to Staging Read Strategy (#5200)

Summary:

X-link: facebookresearch/FBGEMM#2196

**Context:**

When using SSD Offloaded Tensors during Training, we create a Partially Materialized Tensor (i.e., tensor-like) structure to be used for post-training (checkpointing and publishing)

In this diff, we add a new flag "returnKVTensorMetaData" to the Staging Read Strategy, which helps us to read only the metadata of the KVTensor and not the KVTensor itself

Reviewed By: chunzhao

Differential Revision: D86634084

Author: Raahul Kalyaan Jakka

fbgemm_gpu/src/ssd_split_embeddings_cache/kv_tensor_wrapper_cpu.cpp

@@ -108,4 +108,10 @@ std::string KVTensorWrapper::layout_str() {
   oss << options_.layout();
   return oss.str();
 }
+
+std::vector<std::string> KVTensorWrapper::get_kvtensor_serializable_metadata()
+    const {
+  FBEXCEPTION("Not implemented");
+  return std::vector<std::string>{};
+}
 } // namespace ssd

fbgemm_gpu/src/ssd_split_embeddings_cache/ssd_split_table_batched_embeddings.cpp

@@ -415,6 +415,10 @@ std::string KVTensorWrapper::serialize() const {
 std::vector<std::string> KVTensorWrapper::get_kvtensor_serializable_metadata()
     const {
   std::vector<std::string> metadata;
+  // Return empty metadata if checkpoint_handle_ is not initialized yet
+  if (checkpoint_handle_ == nullptr) {
+    return metadata;
+  }
   auto* db = dynamic_cast<EmbeddingRocksDB*>(db_.get());
   auto checkpoint_paths = db->get_checkpoints(checkpoint_handle_->uuid);
   metadata.push_back(std::to_string(checkpoint_paths.size()));