refactor(store): use dedicated thread for signal handling (#840)

* refactor(store): use dedicated thread for signal handling and atomic cleanup

* refactor(store): switch PyClient to shared_ptr and improve resource management

* refactor(ResourceTracker): use leaked heap object to avoid static destruction order issues

* test: increase unittest verbosity to show running tests

* try fix ci

* fix setup after close

* fix comments

Author: xiaguan

mooncake-integration/integration_utils.h

@@ -11,7 +11,8 @@ namespace py = pybind11;
 
 namespace mooncake {
 
-auto torch = py::module_::import("torch");
+// Avoid global py::module_ objects
+inline py::module_ torch_module() { return py::module_::import("torch"); }
 
 enum class TensorDtype : int32_t {
     FLOAT32 = 0,
@@ -68,6 +69,8 @@ inline TensorDtype get_tensor_dtype(py::object dtype_obj) {
         return TensorDtype::UNKNOWN;
     }
 
+    auto torch = torch_module();
+
     if (dtype_obj.equal(torch.attr("float32"))) return TensorDtype::FLOAT32;
     if (dtype_obj.equal(torch.attr("float64"))) return TensorDtype::FLOAT64;
     if (dtype_obj.equal(torch.attr("int8"))) return TensorDtype::INT8;

mooncake-integration/store/store_py.cpp

@@ -17,10 +17,12 @@ namespace mooncake {
 // Python-specific wrapper functions that handle GIL and return pybind11 types
 class MooncakeStorePyWrapper {
    public:
-    PyClient store_;
+    std::shared_ptr<PyClient> store_{nullptr};
+
+    MooncakeStorePyWrapper() : store_(PyClient::create()) {}
 
     pybind11::bytes get(const std::string &key) {
-        if (!store_.client_) {
+        if (!store_ || !store_->client_) {
             LOG(ERROR) << "Client is not initialized";
             return pybind11::bytes("\\0", 0);
         }
@@ -29,7 +31,7 @@ class MooncakeStorePyWrapper {
 
         {
             py::gil_scoped_release release_gil;
-            auto buffer_handle = store_.get_buffer(key);
+            auto buffer_handle = store_->get_buffer(key);
             if (!buffer_handle) {
                 py::gil_scoped_acquire acquire_gil;
                 return kNullString;
@@ -44,15 +46,15 @@ class MooncakeStorePyWrapper {
     std::vector<pybind11::bytes> get_batch(
         const std::vector<std::string> &keys) {
         const auto kNullString = pybind11::bytes("\\0", 0);
-        if (!store_.client_) {
+        if (!store_ || !store_->client_) {
             LOG(ERROR) << "Client is not initialized";
             py::gil_scoped_acquire acquire_gil;
             return {kNullString};
         }
 
         {
             py::gil_scoped_release release_gil;
-            auto batch_data = store_.batch_get_buffer(keys);
+            auto batch_data = store_->batch_get_buffer(keys);
             if (batch_data.empty()) {
                 py::gil_scoped_acquire acquire_gil;
                 return {kNullString};
@@ -73,15 +75,15 @@ class MooncakeStorePyWrapper {
     }
 
     pybind11::object get_tensor(const std::string &key) {
-        if (!store_.client_) {
+        if (!store_ || !store_->client_) {
             LOG(ERROR) << "Client is not initialized";
             return pybind11::none();
         }
 
         try {
             // Section with GIL released
             py::gil_scoped_release release_gil;
-            auto buffer_handle = store_.get_buffer(key);
+            auto buffer_handle = store_->get_buffer(key);
             if (!buffer_handle) {
                 py::gil_scoped_acquire acquire_gil;
                 return pybind11::none();
@@ -144,7 +146,8 @@ class MooncakeStorePyWrapper {
                 py::tuple shape_tuple = py::cast(shape_vec);
                 np_array = np_array.attr("reshape")(shape_tuple);
             }
-            pybind11::object tensor = torch.attr("from_numpy")(np_array);
+            pybind11::object tensor =
+                torch_module().attr("from_numpy")(np_array);
             return tensor;
 
         } catch (const pybind11::error_already_set &e) {
@@ -154,7 +157,7 @@ class MooncakeStorePyWrapper {
     }
 
     int put_tensor(const std::string &key, pybind11::object tensor) {
-        if (!store_.client_) {
+        if (!store_ || !store_->client_) {
             LOG(ERROR) << "Client is not initialized";
             return -static_cast<int>(ErrorCode::INVALID_PARAMS);
         }
@@ -211,7 +214,7 @@ class MooncakeStorePyWrapper {
             values.emplace_back(std::span<const char>(buffer, tensor_size));
 
             // Use put_parts to put metadata and tensor together
-            auto put_result = store_.put_parts_internal(key, values);
+            auto put_result = store_->put_parts_internal(key, values);
             if (!put_result) {
                 return -static_cast<int>(put_result.error());
             }
@@ -287,77 +290,83 @@ PYBIND11_MODULE(store, m) {
                 const std::string &protocol = "tcp",
                 const std::string &rdma_devices = "",
                 const std::string &master_server_addr = "127.0.0.1:50051") {
-                 return self.store_.setup(local_hostname, metadata_server,
-                                          global_segment_size,
-                                          local_buffer_size, protocol,
-                                          rdma_devices, master_server_addr);
+                 if (!self.store_) {
+                     self.store_ = PyClient::create();
+                 }
+                 return self.store_->setup(local_hostname, metadata_server,
+                                           global_segment_size,
+                                           local_buffer_size, protocol,
+                                           rdma_devices, master_server_addr);
              })
         .def("init_all",
              [](MooncakeStorePyWrapper &self, const std::string &protocol,
                 const std::string &device_name,
                 size_t mount_segment_size = 1024 * 1024 * 16) {
-                 return self.store_.initAll(protocol, device_name,
-                                            mount_segment_size);
+                 return self.store_->initAll(protocol, device_name,
+                                             mount_segment_size);
              })
         .def("get", &MooncakeStorePyWrapper::get)
         .def("get_batch", &MooncakeStorePyWrapper::get_batch)
         .def(
             "get_buffer",
             [](MooncakeStorePyWrapper &self, const std::string &key) {
                 py::gil_scoped_release release;
-                return self.store_.get_buffer(key);
+                return self.store_->get_buffer(key);
             },
             py::return_value_policy::take_ownership)
         .def(
             "batch_get_buffer",
             [](MooncakeStorePyWrapper &self,
                const std::vector<std::string> &keys) {
                 py::gil_scoped_release release;
-                return self.store_.batch_get_buffer(keys);
+                return self.store_->batch_get_buffer(keys);
             },
             py::return_value_policy::take_ownership)
         .def("remove",
              [](MooncakeStorePyWrapper &self, const std::string &key) {
                  py::gil_scoped_release release;
-                 return self.store_.remove(key);
+                 return self.store_->remove(key);
              })
         .def(
             "remove_by_regex",
             [](MooncakeStorePyWrapper &self, const std::string &str) {
                 py::gil_scoped_release release;
-                return self.store_.removeByRegex(str);
+                return self.store_->removeByRegex(str);
             },
             py::arg("regex_pattern"),
             "Removes objects from the store whose keys match the given "
             "regular expression.")
         .def("remove_all",
              [](MooncakeStorePyWrapper &self) {
                  py::gil_scoped_release release;
-                 return self.store_.removeAll();
+                 return self.store_->removeAll();
              })
         .def("is_exist",
              [](MooncakeStorePyWrapper &self, const std::string &key) {
                  py::gil_scoped_release release;
-                 return self.store_.isExist(key);
+                 return self.store_->isExist(key);
              })
         .def(
             "batch_is_exist",
             [](MooncakeStorePyWrapper &self,
                const std::vector<std::string> &keys) {
                 py::gil_scoped_release release;
-                return self.store_.batchIsExist(keys);
+                return self.store_->batchIsExist(keys);
             },
             py::arg("keys"),
             "Check if multiple objects exist. Returns list of results: 1 if "
             "exists, 0 if not exists, -1 if error")
         .def("close",
              [](MooncakeStorePyWrapper &self) {
-                 return self.store_.tearDownAll();
+                 if (!self.store_) return 0;
+                 int rc = self.store_->tearDownAll();
+                 self.store_.reset();
+                 return rc;
              })
         .def("get_size",
              [](MooncakeStorePyWrapper &self, const std::string &key) {
                  py::gil_scoped_release release;
-                 return self.store_.getSize(key);
+                 return self.store_->getSize(key);
              })
         .def("get_tensor", &MooncakeStorePyWrapper::get_tensor, py::arg("key"),
              "Get a PyTorch tensor from the store")
@@ -370,7 +379,7 @@ PYBIND11_MODULE(store, m) {
                 // Register memory buffer for RDMA operations
                 void *buffer = reinterpret_cast<void *>(buffer_ptr);
                 py::gil_scoped_release release;
-                return self.store_.register_buffer(buffer, size);
+                return self.store_->register_buffer(buffer, size);
             },
             py::arg("buffer_ptr"), py::arg("size"),
             "Register a memory buffer for direct access operations")
@@ -380,7 +389,7 @@ PYBIND11_MODULE(store, m) {
                 // Unregister memory buffer
                 void *buffer = reinterpret_cast<void *>(buffer_ptr);
                 py::gil_scoped_release release;
-                return self.store_.unregister_buffer(buffer);
+                return self.store_->unregister_buffer(buffer);
             },
             py::arg("buffer_ptr"),
             "Unregister a previously registered memory "
@@ -392,7 +401,7 @@ PYBIND11_MODULE(store, m) {
                 // Get data directly into user-provided buffer
                 void *buffer = reinterpret_cast<void *>(buffer_ptr);
                 py::gil_scoped_release release;
-                return self.store_.get_into(key, buffer, size);
+                return self.store_->get_into(key, buffer, size);
             },
             py::arg("key"), py::arg("buffer_ptr"), py::arg("size"),
             "Get object data directly into a pre-allocated buffer")
@@ -408,7 +417,7 @@ PYBIND11_MODULE(store, m) {
                     buffers.push_back(reinterpret_cast<void *>(ptr));
                 }
                 py::gil_scoped_release release;
-                return self.store_.batch_get_into(keys, buffers, sizes);
+                return self.store_->batch_get_into(keys, buffers, sizes);
             },
             py::arg("keys"), py::arg("buffer_ptrs"), py::arg("sizes"),
             "Get object data directly into pre-allocated buffers for "
@@ -422,7 +431,7 @@ PYBIND11_MODULE(store, m) {
                 // Put data directly from user-provided buffer
                 void *buffer = reinterpret_cast<void *>(buffer_ptr);
                 py::gil_scoped_release release;
-                return self.store_.put_from(key, buffer, size, config);
+                return self.store_->put_from(key, buffer, size, config);
             },
             py::arg("key"), py::arg("buffer_ptr"), py::arg("size"),
             py::arg("config") = ReplicateConfig{},
@@ -439,7 +448,7 @@ PYBIND11_MODULE(store, m) {
                 void *metadata_buffer =
                     reinterpret_cast<void *>(metadata_buffer_ptr);
                 py::gil_scoped_release release;
-                return self.store_.put_from_with_metadata(
+                return self.store_->put_from_with_metadata(
                     key, buffer, metadata_buffer, size, metadata_size, config);
             },
             py::arg("key"), py::arg("buffer_ptr"),
@@ -460,7 +469,8 @@ PYBIND11_MODULE(store, m) {
                     buffers.push_back(reinterpret_cast<void *>(ptr));
                 }
                 py::gil_scoped_release release;
-                return self.store_.batch_put_from(keys, buffers, sizes, config);
+                return self.store_->batch_put_from(keys, buffers, sizes,
+                                                   config);
             },
             py::arg("keys"), py::arg("buffer_ptrs"), py::arg("sizes"),
             py::arg("config") = ReplicateConfig{},
@@ -474,7 +484,7 @@ PYBIND11_MODULE(store, m) {
                const ReplicateConfig &config = ReplicateConfig{}) {
                 py::buffer_info info = buf.request(/*writable=*/false);
                 py::gil_scoped_release release;
-                return self.store_.put(
+                return self.store_->put(
                     key,
                     std::span<const char>(static_cast<char *>(info.ptr),
                                           static_cast<size_t>(info.size)),
@@ -507,7 +517,7 @@ PYBIND11_MODULE(store, m) {
 
                 // 2) Call C++ function
                 py::gil_scoped_release unlock;
-                return self.store_.put_parts(key, spans, config);
+                return self.store_->put_parts(key, spans, config);
             },
             py::arg("key"), py::arg("config") = ReplicateConfig{})
         .def(
@@ -530,12 +540,12 @@ PYBIND11_MODULE(store, m) {
                 }
 
                 py::gil_scoped_release release;
-                return self.store_.put_batch(keys, spans, config);
+                return self.store_->put_batch(keys, spans, config);
             },
             py::arg("keys"), py::arg("values"),
             py::arg("config") = ReplicateConfig{})
         .def("get_hostname", [](MooncakeStorePyWrapper &self) {
-            return self.store_.get_hostname();
+            return self.store_->get_hostname();
         });
 
     // Expose NUMA binding as a module-level function (no self required)

mooncake-store/include/pybind_client.h

@@ -1,12 +1,15 @@
 #pragma once
 
 #include <csignal>
-#include <mutex>
+#include <atomic>
+#include <thread>
 #include <string>
-#include <unordered_set>
+#include <memory>
+#include <vector>
 
 #include "client.h"
 #include "client_buffer.hpp"
+#include "mutex.h"
 #include "utils.h"
 
 namespace mooncake {
@@ -40,10 +43,7 @@ class ResourceTracker {
     static ResourceTracker &getInstance();
 
     // Register a DistributedObjectStore instance for cleanup
-    void registerInstance(PyClient *instance);
-
-    // Unregister a DistributedObjectStore instance
-    void unregisterInstance(PyClient *instance);
+    void registerInstance(const std::shared_ptr<PyClient> &instance);
 
    private:
     ResourceTracker();
@@ -62,15 +62,26 @@ class ResourceTracker {
     // Exit handler function
     static void exitHandler();
 
-    std::mutex mutex_;
-    std::unordered_set<PyClient *> instances_;
+    Mutex mutex_;
+    std::vector<std::weak_ptr<PyClient>> instances_ GUARDED_BY(mutex_);
+
+    // Ensure cleanup runs at most once
+    std::atomic<bool> cleaned_{false};
+
+    // Dedicated signal handling thread
+    void startSignalThread();
+    std::once_flag signal_once_{};
+    std::jthread signal_thread_{};  // joins on destruction
 };
 
 class PyClient {
    public:
     PyClient();
     ~PyClient();
 
+    // Factory to create shared instances and auto-register to ResourceTracker
+    static std::shared_ptr<PyClient> create();
+
     int setup(const std::string &local_hostname,
               const std::string &metadata_server,
               size_t global_segment_size = 1024 * 1024 * 16,
@@ -308,6 +319,9 @@ class PyClient {
     std::string protocol;
     std::string device_name;
     std::string local_hostname;
+
+    // Ensure cleanup executes at most once across multiple entry points
+    std::atomic<bool> closed_{false};
 };
 
 }  // namespace mooncake