[perf] Reduce directory lock conflicts during batch dumps in PosixStore (#711)

…ref (#707)

Reduce directory lock conflicts during batch dumps in PosixStore,
improve its dump performance.

- Split the unified concurrency configuration item (`stream_number`)
into two independent concurrency configurations
(`posix_lookup_concurrency` and `posix_data_trans_concurrency`),
allowing for individual adjustment of the `Lookup` concurrency.
- The directory creation action during `Dump` is moved to the `Setup`
stage, reducing metadata operations during `Dump`.
- Write multiple files during the `Dump` process to different shard
directories to reduce directory conflicts during concurrent file writes.

Author: mag1c-h

ucm/store/posix/cc/global_config.h

@@ -36,9 +36,10 @@ struct Config {
     size_t shardSize{0};
     size_t blockSize{0};
     bool ioDirect{false};
-    size_t streamNumber{8};
+    size_t dataTransConcurrency{8};
+    size_t lookupConcurrency{8};
     size_t timeoutMs{30000};
-    size_t dataDirShardBytes{4};
+    size_t dataDirShardBytes{3};
 };
 
 }  // namespace UC::PosixStore

ucm/store/posix/cc/posix_store.cc

@@ -63,8 +63,12 @@ class PosixStoreImpl {
         if (config.deviceId < -1) {
             return Status::InvalidParam("invalid device({})", config.deviceId);
         }
-        if (config.streamNumber == 0) {
-            return Status::InvalidParam("invalid stream number({})", config.streamNumber);
+        if (config.dataTransConcurrency == 0 || config.lookupConcurrency == 0) {
+            return Status::InvalidParam("invalid concurrency({},{})", config.dataTransConcurrency,
+                                        config.lookupConcurrency);
+        }
+        if (config.dataDirShardBytes > 5) {
+            return Status::InvalidParam("invalid shard bytes({})", config.dataDirShardBytes);
         }
         if (config.deviceId == -1) { return Status::OK(); }
         if (config.tensorSize == 0 || config.shardSize < config.tensorSize ||
@@ -87,7 +91,8 @@ class PosixStoreImpl {
         UC_INFO("Set {}::ShardSize to {}.", ns, config.shardSize);
         UC_INFO("Set {}::BlockSize to {}.", ns, config.blockSize);
         UC_INFO("Set {}::IoDirect to {}.", ns, config.ioDirect);
-        UC_INFO("Set {}::StreamNumber to {}.", ns, config.streamNumber);
+        UC_INFO("Set {}::DataTransConcurrency to {}.", ns, config.dataTransConcurrency);
+        UC_INFO("Set {}::LookupConcurrency to {}.", ns, config.lookupConcurrency);
         UC_INFO("Set {}::TimeoutMs to {}.", ns, config.timeoutMs);
         UC_INFO("Set {}::DataDirShardBytes to {}.", ns, config.dataDirShardBytes);
     }

ucm/store/posix/cc/space_layout.cc

@@ -30,13 +30,32 @@
 namespace UC::PosixStore {
 
 static const std::string DATA_ROOT = "data";
-static const std::string TEMP_ROOT = ".temp";
+static const std::string ACTIVATED_FILE_EXTENSION = ".tmp";
 
 inline std::string DataFileName(const Detail::BlockId& blockId)
 {
     return fmt::format("{:02x}", fmt::join(blockId, ""));
 }
 
+std::vector<std::string> GenerateHexStrings(const size_t n)
+{
+    if (n == 0) [[unlikely]] { return {}; }
+    size_t nCombinations = 1ULL << (n * 4);
+    std::vector<std::string> result;
+    result.reserve(nCombinations);
+    constexpr char hexChars[] = "0123456789abcdef";
+    for (size_t i = 0; i < nCombinations; ++i) {
+        std::string s(n, '0');
+        auto temp = i;
+        for (int j = n - 1; j >= 0; --j) {
+            s[j] = hexChars[temp & 0xF];
+            temp >>= 4;
+        }
+        result.push_back(s);
+    }
+    return result;
+}
+
 Status SpaceLayout::Setup(const Config& config)
 {
     dataDirShardBytes_ = config.dataDirShardBytes;
@@ -52,33 +71,27 @@ std::string SpaceLayout::DataFilePath(const Detail::BlockId& blockId, bool activ
 {
     const auto& backend = StorageBackend(blockId);
     const auto& file = DataFileName(blockId);
-    const auto& shard = activated ? TEMP_ROOT : (dataDirShard_ ? FileShardName(file) : DATA_ROOT);
-    return fmt::format("{}{}/{}", backend, shard, file);
+    const auto& shard = dataDirShard_ ? FileShardName(file) : DATA_ROOT;
+    if (!activated) { return fmt::format("{}{}/{}", backend, shard, file); }
+    return fmt::format("{}{}/{}{}", backend, shard, file, ACTIVATED_FILE_EXTENSION);
 }
 
 Status SpaceLayout::CommitFile(const Detail::BlockId& blockId, bool success) const
 {
-    const auto& backend = StorageBackend(blockId);
-    const auto& file = DataFileName(blockId);
-    const auto& activated = fmt::format("{}{}/{}", backend, TEMP_ROOT, file);
+    const auto& activated = DataFilePath(blockId, true);
     auto s = Status::OK();
     if (success) {
-        const auto& shard = dataDirShard_ ? FileShardName(file) : DATA_ROOT;
-        const auto& archived = fmt::format("{}{}/{}", backend, shard, file);
-        if (dataDirShard_) { s = PosixFile{backend + shard}.MkDir(); }
-        if (s == Status::OK() || s == Status::DuplicateKey()) {
-            s = PosixFile{activated}.Rename(archived);
-        }
+        const auto& archived = DataFilePath(blockId, false);
+        s = PosixFile{activated}.Rename(archived);
     }
     if (!success || s.Failure()) { PosixFile{activated}.Remove(); }
     return s;
 }
 
 std::vector<std::string> SpaceLayout::RelativeRoots() const
 {
-    std::vector<std::string> roots{TEMP_ROOT};
-    if (!dataDirShard_) { roots.push_back(DATA_ROOT); }
-    return roots;
+    if (dataDirShard_) { return GenerateHexStrings(dataDirShardBytes_); }
+    return {DATA_ROOT};
 }
 
 Status SpaceLayout::AddStorageBackend(const std::string& path)

ucm/store/posix/cc/space_manager.cc

@@ -35,7 +35,7 @@ Status SpaceManager::Setup(const Config& config)
         lookupSrv_.SetWorkerFn([this](LookupContext& ctx, auto&) { OnLookup(ctx); })
             .SetWorkerTimeoutFn([this](LookupContext& ctx, auto) { OnLookupTimeout(ctx); },
                                 config.timeoutMs)
-            .SetNWorker(config.streamNumber)
+            .SetNWorker(config.lookupConcurrency)
             .Run();
     if (!success) [[unlikely]] { return Status::Error("failed to run lookup service thread pool"); }
     return Status::OK();

ucm/store/posix/cc/trans_queue.cc

@@ -35,11 +35,11 @@ Status TransQueue::Setup(const Config& config, TaskIdSet* failureSet, const Spac
     shardSize_ = config.shardSize;
     nShardPerBlock_ = config.blockSize / config.shardSize;
     ioDirect_ = config.ioDirect;
-    auto success = pool_.SetNWorker(config.streamNumber)
+    auto success = pool_.SetNWorker(config.dataTransConcurrency)
                        .SetWorkerFn([this](auto& ios, auto&) { Worker(ios); })
                        .Run();
     if (!success) [[unlikely]] {
-        return Status::Error(fmt::format("workers({}) start failed", config.streamNumber));
+        return Status::Error(fmt::format("workers({}) start failed", config.dataTransConcurrency));
     }
     return Status::OK();
 }

ucm/store/posix/connector.py

@@ -48,7 +48,8 @@ def __init__(self, config: Dict[str, object]) -> None:
             "shard_size": "shardSize",
             "block_size": "blockSize",
             "io_direct": "ioDirect",
-            "stream_number": "streamNumber",
+            "posix_data_trans_concurrency": "dataTransConcurrency",
+            "posix_lookup_concurrency": "lookupConcurrency",
             "timeout_ms": "timeoutMs",
             "data_dir_shard_bytes": "dataDirShardBytes",
         }
@@ -154,7 +155,7 @@ def check(self, task: Task) -> bool:
     config["shard_size"] = block_size
     config["block_size"] = block_size
     config["io_direct"] = True
-    config["stream_number"] = 16
+    config["posix_data_trans_concurrency"] = 16
     store = UcmPosixStore(config)
     block_num = 1024
     block_ids = [secrets.token_bytes(16) for _ in range(block_num)]

ucm/store/posix/cpy/posix_store.py.cc

@@ -42,7 +42,8 @@ PYBIND11_MODULE(ucmposixstore, module)
     config.def_readwrite("shardSize", &Config::shardSize);
     config.def_readwrite("blockSize", &Config::blockSize);
     config.def_readwrite("ioDirect", &Config::ioDirect);
-    config.def_readwrite("streamNumber", &Config::streamNumber);
+    config.def_readwrite("dataTransConcurrency", &Config::dataTransConcurrency);
+    config.def_readwrite("lookupConcurrency", &Config::lookupConcurrency);
     config.def_readwrite("timeoutMs", &Config::timeoutMs);
     config.def_readwrite("dataDirShardBytes", &Config::dataDirShardBytes);
     store.def(py::init<>());

ucm/store/test/case/posix/posix_space_manager_test.cc

@@ -72,7 +72,7 @@ TEST_F(UCPosixSpaceManagerTest, DataFilePath)
     ASSERT_EQ(s, UC::Status::OK());
     auto blockId = UC::Test::Detail::TypesHelper::MakeBlockId("a1b2c3d4e5f6789012345678901234ab");
     auto activated = spaceMgr.GetLayout()->DataFilePath(blockId, true);
-    ASSERT_EQ(activated, fmt::format("{}.temp/{:02x}", this->Path(), fmt::join(blockId, "")));
+    ASSERT_EQ(activated, fmt::format("{}data/{:02x}.tmp", this->Path(), fmt::join(blockId, "")));
     ASSERT_EQ(PosixFile{activated}.Access(PosixFile::AccessMode::EXIST), UC::Status::NotFound());
     ASSERT_EQ(PosixFile{activated}.Open(PosixFile::OpenFlag::CREATE), UC::Status::OK());
     ASSERT_EQ(PosixFile{activated}.Access(PosixFile::AccessMode::EXIST), UC::Status::OK());
@@ -95,8 +95,10 @@ TEST_F(UCPosixSpaceManagerTest, ShardFilePath)
     auto s = spaceMgr.Setup(config);
     ASSERT_EQ(s, UC::Status::OK());
     auto blockId = UC::Test::Detail::TypesHelper::MakeBlockIdRandomly();
+    const auto& file = fmt::format("{:02x}", fmt::join(blockId, ""));
+    const auto& shard = file.substr(0, config.dataDirShardBytes);
     auto activated = spaceMgr.GetLayout()->DataFilePath(blockId, true);
-    ASSERT_EQ(activated, fmt::format("{}.temp/{:02x}", this->Path(), fmt::join(blockId, "")));
+    ASSERT_EQ(activated, fmt::format("{}{}/{}.tmp", this->Path(), shard, file));
     ASSERT_EQ(PosixFile{activated}.Access(PosixFile::AccessMode::EXIST), UC::Status::NotFound());
     ASSERT_EQ(PosixFile{activated}.Open(PosixFile::OpenFlag::CREATE), UC::Status::OK());
     ASSERT_EQ(PosixFile{activated}.Access(PosixFile::AccessMode::EXIST), UC::Status::OK());
@@ -105,8 +107,6 @@ TEST_F(UCPosixSpaceManagerTest, ShardFilePath)
     ASSERT_EQ(spaceMgr.Lookup(&blockId, 1).Value(), std::vector<uint8_t>{true});
     ASSERT_EQ(PosixFile{activated}.Access(PosixFile::AccessMode::EXIST), UC::Status::NotFound());
     auto archived = spaceMgr.GetLayout()->DataFilePath(blockId, false);
-    const auto& file = fmt::format("{:02x}", fmt::join(blockId, ""));
-    const auto& shard = file.substr(0, config.dataDirShardBytes);
     ASSERT_EQ(archived, fmt::format("{}{}/{}", this->Path(), shard, file));
     ASSERT_EQ(PosixFile{archived}.Access(PosixFile::AccessMode::EXIST), UC::Status::OK());
 }

ucm/store/test/case/posix/posix_store_test.cc

@@ -50,7 +50,7 @@ TEST_F(UCPosixStoreTest, SetupWithInvalidParam)
         config.tensorSize = 4096;
         config.shardSize = config.tensorSize;
         config.blockSize = config.shardSize;
-        config.streamNumber = 0;
+        config.dataTransConcurrency = 0;
         PosixStore store;
         ASSERT_EQ(store.Setup(config), UC::Status::InvalidParam());
     }

ucm/store/test/e2e/cache_on_posix_test.py

@@ -53,10 +53,19 @@ def e2e_test(
     request_size: int,
     device_id: int,
 ):
+    # make block id randomly
     chunk_block_ids = [secrets.token_bytes(16) for _ in range(request_size)]
+    # fully lookup at 0% hit
+    tp = time.perf_counter()
     founds = scheduler.lookup(chunk_block_ids)
+    cost_fully_lookup1 = time.perf_counter() - tp
     assert not any(founds)
-    assert scheduler.lookup_on_prefix(chunk_block_ids) == -1
+    # prefix lookup at 0% hit
+    tp = time.perf_counter()
+    found_idx = scheduler.lookup_on_prefix(chunk_block_ids)
+    cost_prefix_lookup1 = time.perf_counter() - tp
+    assert found_idx == -1
+    # make tensor randomly
     shard_indexes = [0 for _ in range(request_size)]
     src_tensors = [
         [
@@ -69,15 +78,43 @@ def e2e_test(
         ]
         for _ in range(request_size)
     ]
+    # dump data to store
+    tp = time.perf_counter()
     task = worker.dump(chunk_block_ids, shard_indexes, src_tensors)
     worker.wait(task)
+    cost_dump = time.perf_counter() - tp
+    # fully lookup at 100% hit
+    tp = time.perf_counter()
     founds = scheduler.lookup(chunk_block_ids)
+    cost_fully_lookup2 = time.perf_counter() - tp
     assert all(founds)
-    assert scheduler.lookup_on_prefix(chunk_block_ids) + 1 == request_size
+    # prefix lookup at 100% hit
+    tp = time.perf_counter()
+    found_idx = scheduler.lookup_on_prefix(chunk_block_ids)
+    cost_prefix_lookup2 = time.perf_counter() - tp
+    assert found_idx + 1 == request_size
+    # make tensor buffer for fetching
     dst_tensors = [[torch.empty_like(t) for t in row] for row in src_tensors]
+    # fetch data from store
+    tp = time.perf_counter()
     task = worker.load(chunk_block_ids, shard_indexes, dst_tensors)
     worker.wait(task)
+    cost_load = time.perf_counter() - tp
+    # compare data
     cmp_and_print_diff(src_tensors, dst_tensors)
+    # show cost
+    data_size = tensor_size * layer_size * chunk_size * request_size
+    bw_dump = data_size / cost_dump
+    bw_load = data_size / cost_load
+    print(
+        f"[{tensor_size}-{layer_size}-{chunk_size}-{request_size}] "
+        f"fully_lookup1={cost_fully_lookup1 * 1e3:.3f}ms, "
+        f"prefix_lookup1={cost_prefix_lookup1 * 1e3:.3f}ms, "
+        f"fully_lookup2={cost_fully_lookup2 * 1e3:.3f}ms, "
+        f"prefix_lookup2={cost_prefix_lookup2 * 1e3:.3f}ms, "
+        f"dump={cost_dump * 1e3:.3f}ms, load={cost_load * 1e3:.3f}ms, "
+        f"bw_dump={bw_dump / 1e9:.3f}GB/s, bw_load={bw_load / 1e9:.3f}GB/s."
+    )
 
 
 def main():
@@ -100,7 +137,8 @@ def main():
     config["waiting_queue_depth"] = 16
     config["running_queue_depth"] = 1024
     config["io_direct"] = True
-    config["stream_number"] = 16
+    config["posix_data_trans_concurrency"] = 32
+    config["posix_lookup_concurrency"] = 32
     worker = UcmPipelineStore(config | {"device_id": device_id})
     scheduler = UcmPipelineStore(config)
     test_batch_number = 512
@@ -118,5 +156,5 @@ def main():
 
 
 if __name__ == "__main__":
-    os.environ["UC_LOGGER_LEVEL"] = "debug"
+    os.environ["UC_LOGGER_LEVEL"] = "info"
     main()