MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / RestoreSSD

Method RestoreSSD

tensorflow/core/framework/embedding/storage.h:214–252  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

212
213 protected:
214 virtual Status RestoreSSD(int64 emb_index, int64 emb_slot_num,
215 int64 value_len,
216 const std::string& ssd_emb_file_name,
217 EmbeddingVar<K, V>* ev,
218 RestoreSSDBuffer<K>& restore_buff) {
219 for (int64 i = 0; i < restore_buff.num_of_keys; i++) {
220 int64 file_id = restore_buff.key_file_id_list_buf[i];
221 int64 key_offset = restore_buff.key_offset_list_buf[i];
222 // Read data from embedding files on SSD. Data are stored in
223 // NormalContiguousValuePtr temporarily.
224 std::stringstream ss;
225 ss << ssd_emb_file_name << "/" << file_id << ".emb";
226 int fd = open(ss.str().data(), O_RDONLY);
227 EmbeddingConfig& emb_config = storage_config_.embedding_config;
228 FeatureDescriptor<V> normal_feat_desc(
229 emb_config.block_num, emb_config.slot_num + 1,
230 ev_allocator(), StorageType::DRAM, true,
231 true, {false, 0});
232 void* value_ptr = normal_feat_desc.Allocate();
233 char* file_addr = (char*)mmap(nullptr,
234 normal_feat_desc.data_bytes() +
235 key_offset,
236 PROT_READ, MAP_PRIVATE, fd, 0);
237 memcpy(value_ptr, file_addr + key_offset,
238 normal_feat_desc.data_bytes());
239 munmap(file_addr,
240 normal_feat_desc.data_bytes() +
241 key_offset);
242 close(fd);
243 // Copy Data to ValuePtr, data of slots are set by primary here.
244 int64 import_freq = normal_feat_desc.GetFreq(value_ptr);
245 int64 import_version = normal_feat_desc.GetVersion(value_ptr);
246 V* value = normal_feat_desc.GetEmbedding(value_ptr, emb_index);
247 Import(restore_buff.key_list_buf[i], value,
248 import_freq, import_version, emb_index);
249 normal_feat_desc.Deallocate(value_ptr);
250 }
251 return Status::OK();
252 }
253
254 private:
255 void GeneratePartitionedCkptData(

Callers

nothing calls this directly

Calls 9

ev_allocatorFunction · 0.85
ImportFunction · 0.50
dataMethod · 0.45
AllocateMethod · 0.45
data_bytesMethod · 0.45
GetFreqMethod · 0.45
GetVersionMethod · 0.45
GetEmbeddingMethod · 0.45
DeallocateMethod · 0.45

Tested by

no test coverage detected