| 344 | |
| 345 | template<typename Key, typename Engine> |
| 346 | PersistentTableImpl<Key, Engine>::PersistentTableImpl(const PersistentTableOptions& options) |
| 347 | : root_dir_(options.path), |
| 348 | key_size_(options.key_size), |
| 349 | value_size_(options.value_size), |
| 350 | physical_block_size_(options.physical_block_size), |
| 351 | logical_block_size_(GetLogicalBlockSize(options.physical_block_size, value_size_)), |
| 352 | blocks_buffer_(options.physical_block_size), |
| 353 | writable_key_file_chunk_id_(-1), |
| 354 | read_only_(options.read_only) { |
| 355 | const uint64_t capacity_hint = ParseIntegerFromEnv( |
| 356 | "ONEFLOW_ONE_EMBEDDING_PERSISTENT_TABLE_CAPACITY_HINT", options.capacity_hint); |
| 357 | if (capacity_hint > 0) { row_id_mapping_.reserve(capacity_hint); } |
| 358 | PosixFile::RecursiveCreateDirectory(options.path, 0755); |
| 359 | const std::string lock_filename = PosixFile::JoinPath(options.path, kLockFileName); |
| 360 | const bool init = !PosixFile::FileExists(lock_filename); |
| 361 | if (read_only_) { |
| 362 | CHECK(!init) << "The table must be initialized in read only mode"; |
| 363 | } else { |
| 364 | lock_ = PosixFileLockGuard(PosixFile(lock_filename, O_CREAT | O_RDWR, 0644)); |
| 365 | } |
| 366 | const uint64_t target_chunk_size = options.target_chunk_size_mb * 1024 * 1024; |
| 367 | CHECK_GE(target_chunk_size, logical_block_size_); |
| 368 | num_logical_blocks_per_chunk_ = target_chunk_size / logical_block_size_, |
| 369 | num_values_per_block_ = logical_block_size_ / value_size_; |
| 370 | num_values_per_chunk_ = num_values_per_block_ * num_logical_blocks_per_chunk_; |
| 371 | InitOrCheckMetaValue(PosixFile::JoinPath(options.path, kKeySizeFileName), key_size_, init); |
| 372 | InitOrCheckMetaValue(PosixFile::JoinPath(options.path, kValueSizeFileName), value_size_, init); |
| 373 | InitOrCheckMetaValue(PosixFile::JoinPath(options.path, kPhysicalBlockSizeFileName), |
| 374 | options.physical_block_size, init); |
| 375 | InitOrCheckMetaValue(PosixFile::JoinPath(options.path, kNumLogicalBlocksPerChunkFileName), |
| 376 | num_logical_blocks_per_chunk_, init); |
| 377 | keys_dir_ = PosixFile::JoinPath(options.path, kKeysDirName); |
| 378 | values_dir_ = PosixFile::JoinPath(options.path, kValuesDirName); |
| 379 | snapshots_dir_ = PosixFile::JoinPath(options.path, kSnapshotsDirName); |
| 380 | if (init) { |
| 381 | PosixFile::RecursiveCreateDirectory(keys_dir_, 0755); |
| 382 | PosixFile::RecursiveCreateDirectory(values_dir_, 0755); |
| 383 | } |
| 384 | const uint32_t num_workers = ParseIntegerFromEnv( |
| 385 | "ONEFLOW_ONE_EMBEDDING_PERSISTENT_TABLE_NUM_WORKERS", kDefaultNumWorkerThreads); |
| 386 | workers_.resize(num_workers); |
| 387 | for (uint32_t tid = 0; tid < workers_.size(); ++tid) { |
| 388 | workers_.at(tid).reset(new Worker<Engine>); |
| 389 | } |
| 390 | std::unordered_map<uint64_t, std::string> chunks; |
| 391 | ListChunkFiles(values_dir_, kValueFileNamePrefix, &chunks); |
| 392 | for (auto& chunk : chunks) { |
| 393 | if (value_files_.size() <= chunk.first) { value_files_.resize(chunk.first + 1); } |
| 394 | CHECK_EQ(value_files_.at(chunk.first).fd(), -1); |
| 395 | const int flags = read_only_ ? (O_RDONLY | O_DIRECT) : (O_RDWR | O_DIRECT); |
| 396 | PosixFile value_file(chunk.second, flags, 0644); |
| 397 | value_files_.at(chunk.first) = std::move(value_file); |
| 398 | } |
| 399 | if (!value_files_.empty()) { |
| 400 | physical_table_size_ = ((value_files_.size() - 1) * num_logical_blocks_per_chunk_ |
| 401 | + value_files_.back().Size() / logical_block_size_) |
| 402 | * num_values_per_block_; |
| 403 | } else { |
nothing calls this directly
no test coverage detected