refactor: reusable transport abstraction & optimized NSFStore pipeline (#296)

Author: mag1c-h

ucm/store/CMakeLists.txt

@@ -9,4 +9,5 @@ add_subdirectory(nfsstore)
 add_subdirectory(dramstore)
 add_subdirectory(localstore)
 add_subdirectory(mooncakestore)
+add_subdirectory(task)
 add_subdirectory(test)

ucm/store/device/ascend/ascend_device.cc

@@ -91,6 +91,14 @@ class AscendDevice : public IBufferedDevice {
         }
         return Status::OK();
     }
+    Status H2DSync(std::byte* dst, const std::byte* src, const size_t count) override
+    {
+        return ASCEND_API(aclrtMemcpy, dst, count, src, count, ACL_MEMCPY_HOST_TO_DEVICE);
+    }
+    Status D2HSync(std::byte* dst, const std::byte* src, const size_t count) override
+    {
+        return ASCEND_API(aclrtMemcpy, dst, count, src, count, ACL_MEMCPY_DEVICE_TO_HOST);
+    }
     Status H2DAsync(std::byte* dst, const std::byte* src, const size_t count) override
     {
         return ASCEND_API(aclrtMemcpyAsync, dst, count, src, count, ACL_MEMCPY_HOST_TO_DEVICE,
@@ -111,6 +119,7 @@ class AscendDevice : public IBufferedDevice {
         return ASCEND_API(aclrtLaunchCallback, Trampoline, (void*)c, ACL_CALLBACK_NO_BLOCK,
                           this->stream_);
     }
+    Status Synchronized() override { return ASCEND_API(aclrtSynchronizeStream, this->stream_); }
 
 protected:
     std::shared_ptr<std::byte> MakeBuffer(const size_t size) override

ucm/store/device/cuda/cuda_device.cc

@@ -78,6 +78,14 @@ class CudaDevice : public IBufferedDevice {
         }
         return status;
     }
+    virtual Status H2DSync(std::byte* dst, const std::byte* src, const size_t count) override
+    {
+        return CUDA_API(cudaMemcpy, dst, src, count, cudaMemcpyHostToDevice);
+    }
+    virtual Status D2HSync(std::byte* dst, const std::byte* src, const size_t count) override
+    {
+        return CUDA_API(cudaMemcpy, dst, src, count, cudaMemcpyDeviceToHost);
+    }
     Status H2DAsync(std::byte* dst, const std::byte* src, const size_t count) override
     {
         return CUDA_API(cudaMemcpyAsync, dst, src, count, cudaMemcpyHostToDevice,
@@ -100,6 +108,10 @@ class CudaDevice : public IBufferedDevice {
         if (status.Failure()) { delete c; }
         return status;
     }
+    Status Synchronized() override
+    {
+        return CUDA_API(cudaStreamSynchronize, (cudaStream_t)this->stream_);
+    }
 
 protected:
     std::shared_ptr<std::byte> MakeBuffer(const size_t size) override

ucm/store/device/idevice.h

@@ -39,9 +39,12 @@ class IDevice {
     virtual ~IDevice() = default;
     virtual Status Setup() = 0;
     virtual std::shared_ptr<std::byte> GetBuffer(const size_t size) = 0;
+    virtual Status H2DSync(std::byte* dst, const std::byte* src, const size_t count) = 0;
+    virtual Status D2HSync(std::byte* dst, const std::byte* src, const size_t count) = 0;
     virtual Status H2DAsync(std::byte* dst, const std::byte* src, const size_t count) = 0;
     virtual Status D2HAsync(std::byte* dst, const std::byte* src, const size_t count) = 0;
     virtual Status AppendCallback(std::function<void(bool)> cb) = 0;
+    virtual Status Synchronized() = 0;
 
 protected:
     virtual std::shared_ptr<std::byte> MakeBuffer(const size_t size) = 0;

ucm/store/device/simu/simu_device.cc

@@ -23,6 +23,7 @@
  * */
 #include "ibuffered_device.h"
 #include "logger/logger.h"
+#include "thread/latch.h"
 #include "thread/thread_pool.h"
 
 namespace UC {
@@ -38,7 +39,19 @@ class SimuDevice : public IBufferedDevice {
     {
         auto status = IBufferedDevice::Setup();
         if (status.Failure()) { return status; }
-        if (!this->backend_.Setup([](auto& task) { task(); })) { return Status::Error(); }
+        if (!this->backend_.SetWorkerFn([](auto& task, const auto&) { task(); }).Run()) {
+            return Status::Error();
+        }
+        return Status::OK();
+    }
+    Status H2DSync(std::byte* dst, const std::byte* src, const size_t count) override
+    {
+        std::copy(src, src + count, dst);
+        return Status::OK();
+    }
+    Status D2HSync(std::byte* dst, const std::byte* src, const size_t count) override
+    {
+        std::copy(src, src + count, dst);
         return Status::OK();
     }
     Status H2DAsync(std::byte* dst, const std::byte* src, const size_t count) override
@@ -64,6 +77,13 @@ class SimuDevice : public IBufferedDevice {
         this->backend_.Push([=] { cb(true); });
         return Status::OK();
     }
+    virtual Status Synchronized()
+    {
+        Latch waiter{1};
+        this->backend_.Push([&] { waiter.Done(nullptr); });
+        waiter.Wait();
+        return Status::OK();
+    }
 
 protected:
     std::shared_ptr<std::byte> MakeBuffer(const size_t size) override

ucm/store/dramstore/CMakeLists.txt

@@ -1,7 +1,7 @@
 file(GLOB_RECURSE UCMSTORE_DRAM_CC_SOURCE_FILES "./cc/*.cc")
 add_library(dramstore STATIC ${UCMSTORE_DRAM_CC_SOURCE_FILES})
 target_include_directories(dramstore PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/cc)
-target_link_libraries(dramstore PUBLIC storeinfra)
+target_link_libraries(dramstore PUBLIC storeinfra storetask)
 
 file(GLOB_RECURSE UCMSTORE_DRAM_CPY_SOURCE_FILES "./cpy/*.cc")
 pybind11_add_module(ucmdramstore ${UCMSTORE_DRAM_CPY_SOURCE_FILES})

ucm/store/dramstore/cpy/dramstore.py.cc

@@ -56,31 +56,30 @@ class DRAMStorePy : public DRAMStore {
     size_t Load(const py::list& blockIds, const py::list& offsets, const py::list& addresses,
                 const py::list& lengths)
     {
-        return this->SubmitPy(blockIds, offsets, addresses, lengths, CCStore::Task::Type::LOAD,
-                              CCStore::Task::Location::DEVICE, "DRAM::H2D");
+        return this->SubmitPy(blockIds, offsets, addresses, lengths, Task::Type::LOAD,
+                              Task::Location::DEVICE, "DRAM::H2D");
     }
     size_t Dump(const py::list& blockIds, const py::list& offsets, const py::list& addresses,
                 const py::list& lengths)
     {
-        return this->SubmitPy(blockIds, offsets, addresses, lengths, CCStore::Task::Type::DUMP,
-                              CCStore::Task::Location::DEVICE, "DRAM::D2H");
+        return this->SubmitPy(blockIds, offsets, addresses, lengths, Task::Type::DUMP,
+                              Task::Location::DEVICE, "DRAM::D2H");
     }
 
 private:
     size_t SubmitPy(const py::list& blockIds, const py::list& offsets, const py::list& addresses,
-                    const py::list& lengths, const CCStore::Task::Type type,
-                    const CCStore::Task::Location location, const std::string& brief)
+                    const py::list& lengths, Task::Type&& type, Task::Location&& location,
+                    std::string&& brief)
     {
-        CCStore::Task task{type, location, brief};
+        Task task{std::move(type), std::move(location), std::move(brief)};
         auto blockId = blockIds.begin();
         auto offset = offsets.begin();
         auto address = addresses.begin();
         auto length = lengths.begin();
         while ((blockId != blockIds.end()) && (offset != offsets.end()) &&
                (address != addresses.end()) && (length != lengths.end())) {
-            auto ret = task.Append(blockId->cast<std::string>(), offset->cast<size_t>(),
-                                   address->cast<uintptr_t>(), length->cast<size_t>());
-            if (ret != 0) { return CCStore::invalidTaskId; }
+            task.Append(blockId->cast<std::string>(), offset->cast<size_t>(),
+                        address->cast<uintptr_t>(), length->cast<size_t>());
             blockId++;
             offset++;
             address++;

ucm/store/infra/file/file.cc

@@ -54,21 +54,25 @@ Status File::Access(const std::string& path, const int32_t mode)
 }
 
 Status File::Read(const std::string& path, const size_t offset, const size_t length,
-                  uintptr_t address)
+                  uintptr_t address, const bool directIo)
 {
     FileImpl file{path};
     Status status = Status::OK();
-    if ((status = file.Open(IFile::OpenFlag::READ_ONLY)).Failure()) { return status; }
+    auto flags = directIo ? IFile::OpenFlag::READ_ONLY | IFile::OpenFlag::DIRECT
+                          : IFile::OpenFlag::READ_ONLY;
+    if ((status = file.Open(flags)).Failure()) { return status; }
     if ((status = file.Read((void*)address, length, offset)).Failure()) { return status; }
     return status;
 }
 
 Status File::Write(const std::string& path, const size_t offset, const size_t length,
-                   const uintptr_t address)
+                   const uintptr_t address, const bool directIo)
 {
     FileImpl file{path};
     Status status = Status::OK();
-    if ((status = file.Open(IFile::OpenFlag::WRITE_ONLY)).Failure()) { return status; }
+    auto flags = directIo ? IFile::OpenFlag::WRITE_ONLY | IFile::OpenFlag::DIRECT
+                          : IFile::OpenFlag::WRITE_ONLY;
+    if ((status = file.Open(flags)).Failure()) { return status; }
     if ((status = file.Write((const void*)address, length, offset)).Failure()) { return status; }
     return status;
 }

ucm/store/infra/file/file.h

@@ -37,9 +37,9 @@ class File {
     static Status Rename(const std::string& path, const std::string& newName);
     static Status Access(const std::string& path, const int32_t mode);
     static Status Read(const std::string& path, const size_t offset, const size_t length,
-                       uintptr_t address);
+                       uintptr_t address, const bool directIo = false);
     static Status Write(const std::string& path, const size_t offset, const size_t length,
-                        const uintptr_t address);
+                        const uintptr_t address, const bool directIo = false);
     static void MUnmap(void* addr, size_t size);
     static void ShmUnlink(const std::string& path);
 };

ucm/store/infra/thread/latch.h

@@ -26,6 +26,7 @@
 
 #include <atomic>
 #include <condition_variable>
+#include <functional>
 #include <mutex>
 
 namespace UC {
@@ -34,8 +35,22 @@ class Latch {
 public:
     explicit Latch(const size_t expected = 0) : counter_{expected} {}
     void Up() { ++this->counter_; }
-    size_t Done() { return --this->counter_; }
-    void Notify() { this->cv_.notify_all(); }
+    void Done(std::function<void(void)> finish) noexcept
+    {
+        auto counter = this->counter_.load(std::memory_order_acquire);
+        while (counter > 0) {
+            auto desired = counter - 1;
+            if (this->counter_.compare_exchange_weak(counter, desired, std::memory_order_acq_rel)) {
+                if (desired == 0) {
+                    if (finish) { finish(); }
+                    std::lock_guard<std::mutex> lg(this->mutex_);
+                    this->cv_.notify_all();
+                }
+                return;
+            }
+            counter = this->counter_.load(std::memory_order_acquire);
+        }
+    }
     void Wait()
     {
         std::unique_lock<std::mutex> lk(this->mutex_);