more gil releases

microsoft · Oct 8, 2023 · 8193386 · 8193386
1 parent 033a302
commit 8193386
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 9 deletions.
diff --git a/python/mscclpp/core_py.cpp b/python/mscclpp/core_py.cpp
@@ -64,9 +64,12 @@ void register_core(nb::module_& m) {
           nb::arg("nRanks"))
       .def("create_unique_id", &TcpBootstrap::createUniqueId)
       .def("get_unique_id", &TcpBootstrap::getUniqueId)
-      .def("initialize", (void (TcpBootstrap::*)(UniqueId, int64_t)) & TcpBootstrap::initialize, nb::arg("uniqueId"),
-           nb::arg("timeoutSec") = 30)
-      .def("initialize", (void (TcpBootstrap::*)(const std::string&, int64_t)) & TcpBootstrap::initialize,
+      .def("initialize",
+           gil_release_wrapper(static_cast<void (TcpBootstrap::*)(UniqueId, int64_t)>(&TcpBootstrap::initialize)),
+           nb::arg("uniqueId"), nb::arg("timeoutSec") = 30)
+      .def("initialize",
+           gil_release_wrapper(
+               static_cast<void (TcpBootstrap::*)(const std::string&, int64_t)>(&TcpBootstrap::initialize)),
            nb::arg("ifIpPortTrio"), nb::arg("timeoutSec") = 30);
 
   nb::enum_<Transport>(m, "Transport")

diff --git a/python/mscclpp/gil_release.hpp b/python/mscclpp/gil_release.hpp
@@ -4,15 +4,16 @@
 #ifndef MSCCLPP_PYTHON_GIL_RELEASE_HPP_
 #define MSCCLPP_PYTHON_GIL_RELEASE_HPP_
 
-#include <functional>
 #include <nanobind/nanobind.h>
 
+#include <functional>
+
 template <typename ReturnType, typename ClassType, typename... Args>
 std::function<ReturnType(ClassType*, Args...)> gil_release_wrapper(ReturnType (ClassType::*method)(Args...)) {
-    return [method](ClassType* instance, Args... args) -> ReturnType {
-        nanobind::gil_scoped_release release;
-        return (instance->*method)(std::forward<Args>(args)...);
-    };
+  return [method](ClassType* instance, Args... args) -> ReturnType {
+    nanobind::gil_scoped_release release;
+    return (instance->*method)(std::forward<Args>(args)...);
+  };
 }
 
 #endif  // MSCCLPP_PYTHON_GIL_RELEASE_HPP_
diff --git a/python/mscclpp/semaphore_py.cpp b/python/mscclpp/semaphore_py.cpp
@@ -6,6 +6,8 @@
 
 #include <mscclpp/semaphore.hpp>
 
+#include "gil_release.hpp"
+
 namespace nb = nanobind;
 using namespace mscclpp;
 
@@ -30,7 +32,7 @@ void register_semaphore(nb::module_& m) {
       .def("connection", &Host2HostSemaphore::connection)
       .def("signal", &Host2HostSemaphore::signal)
       .def("poll", &Host2HostSemaphore::poll)
-      .def("wait", &Host2HostSemaphore::wait, nb::arg("max_spin_count") = 10000000);
+      .def("wait", gil_release_wrapper(&Host2HostSemaphore::wait), nb::arg("max_spin_count") = 10000000);
 
   nb::class_<SmDevice2DeviceSemaphore> smDevice2DeviceSemaphore(m, "SmDevice2DeviceSemaphore");
   smDevice2DeviceSemaphore