bugfix in cache dump queue

Author: mag1c-h

ucm/store/cache/cc/dump_queue.cc

@@ -83,7 +83,7 @@ void DumpQueue::DispatchOneTask(Trans::Stream* stream, size_t tensorSize, TaskPa
     auto& task = pair.first;
     auto& waiter = pair.second;
     auto wait = NowTime::Now() - waiter->startTp;
-    UC_DEBUG("Cache task({}) start running, wait {:06f}s.", task->id, wait);
+    UC_DEBUG("Cache task({}) start running, wait {:.3f}ms.", task->id, wait * 1e3);
     if (!failureSet_->Contains(task->id)) {
         auto s = DumpOneTask(stream, tensorSize, task);
         if (s.Failure()) [[unlikely]] { failureSet_->Insert(task->id); }
@@ -97,49 +97,43 @@ Status DumpQueue::DumpOneTask(Trans::Stream* stream, size_t tensorSize, TaskPtr
     Detail::TaskDesc backendTaskDesc;
     backendTaskDesc.brief = "Cache2Backend";
     const auto nShard = task->desc.size();
-    std::vector<size_t> backendTaskIndex;
-    backendTaskIndex.reserve(nShard);
-    std::vector<ShardTask> shardTasks(nShard);
-    for (size_t i = 0; i < task->desc.size(); i++) {
+    DumpCtx dumpCtx;
+    for (size_t i = 0; i < nShard; i++) {
         auto& shard = task->desc[i];
-        auto& shardTask = shardTasks[i];
-        shardTask.bufferHandle = buffer_->Get(shard.owner, shard.index);
-        if (!shardTask.bufferHandle.Owner()) { continue; }
-        if (!shardTask.bufferHandle.Ready()) {
-            auto s = stream->DeviceToHostAsync(shard.addrs.data(), shardTask.bufferHandle.Data(),
-                                               tensorSize, shard.addrs.size());
+        auto handle = buffer_->Get(shard.owner, shard.index);
+        if (!handle.Owner()) { continue; }
+        if (!handle.Ready()) {
+            auto s = stream->DeviceToHostAsync(shard.addrs.data(), handle.Data(), tensorSize,
+                                               shard.addrs.size());
             if (s.Failure()) [[unlikely]] {
                 UC_ERROR("Failed({}) to do D2H({}) batch({}) async.", s, tensorSize,
                          shard.addrs.size());
                 return s;
             }
         }
-        backendTaskDesc.push_back(
-            Detail::Shard{shard.owner, shard.index, {shardTask.bufferHandle.Data()}});
-        backendTaskIndex.emplace_back(i);
+        backendTaskDesc.push_back(Detail::Shard{shard.owner, shard.index, {handle.Data()}});
+        dumpCtx.bufferHandles.push_back(std::move(handle));
     }
     auto tpMakeBuffer = NowTime::Now();
-    if (backendTaskIndex.empty()) { return Status::OK(); }
+    if (backendTaskDesc.empty()) { return Status::OK(); }
     auto s = stream->Synchronized();
     if (s.Failure()) [[unlikely]] {
         UC_ERROR("Failed({}) to sync on stream.", s);
         return s;
     }
     auto tpSyncStream = NowTime::Now();
-    for (const auto& i : backendTaskIndex) { shardTasks[i].bufferHandle.MarkReady(); }
+    for (auto& handle : dumpCtx.bufferHandles) { handle.MarkReady(); }
     auto res = backend_->Dump(std::move(backendTaskDesc));
     if (!res) [[unlikely]] {
         UC_ERROR("Failed({}) to submit dump task to backend.", res.Error());
         return res.Error();
     }
-    for (const auto& i : backendTaskIndex) {
-        auto& shardTask = shardTasks[i];
-        shardTask.backendTaskHandle = res.Value();
-        dumping_.Push(std::move(shardTask));
-    }
+    dumpCtx.backendTaskHandle = res.Value();
+    dumping_.Push(std::move(dumpCtx));
     auto tpEnd = NowTime::Now();
-    UC_DEBUG("Cache task({}) mk_buf={:06f}s, sync={:06f}s, back={:06f}s.", task->id,
-             tpMakeBuffer - tp, tpSyncStream - tpMakeBuffer, tpEnd - tpSyncStream);
+    UC_DEBUG("Cache task({}) mk_buf={:.3f}ms, sync={:.3f}ms, back={:.3f}ms.", task->id,
+             (tpMakeBuffer - tp) * 1e3, (tpSyncStream - tpMakeBuffer) * 1e3,
+             (tpEnd - tpSyncStream) * 1e3);
     return Status::OK();
 }
 

ucm/store/cache/cc/dump_queue.h

@@ -41,9 +41,9 @@ class DumpQueue {
     using WaiterPtr = std::shared_ptr<Latch>;
     using TaskPair = std::pair<TaskPtr, WaiterPtr>;
     using TaskIdSet = HashSet<Detail::TaskHandle>;
-    struct ShardTask {
+    struct DumpCtx {
         Detail::TaskHandle backendTaskHandle;
-        TransBuffer::Handle bufferHandle;
+        std::vector<TransBuffer::Handle> bufferHandles;
     };
 
 private:
@@ -53,7 +53,7 @@ class DumpQueue {
     TransBuffer* buffer_{nullptr};
     Store* backend_{nullptr};
     SpscRingQueue<TaskPair> waiting_;
-    SpscRingQueue<ShardTask> dumping_;
+    SpscRingQueue<DumpCtx> dumping_;
     std::thread dispatcher_;
     std::thread dumper_;