GIL release

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

Hide file tree

Showing 2 changed files with 21 additions and 1 deletion.
diff --git a/python/mscclpp/core_py.cpp b/python/mscclpp/core_py.cpp
@@ -9,6 +9,8 @@
 
 #include <mscclpp/core.hpp>
 
+#include "gil_release.hpp"
+
 namespace nb = nanobind;
 using namespace mscclpp;
 
@@ -118,7 +120,7 @@ void register_core(nb::module_& m) {
             self->updateAndSync(dst, dstOffset, (uint64_t*)src, newValue);
           },
           nb::arg("dst"), nb::arg("dstOffset"), nb::arg("src"), nb::arg("newValue"))
-      .def("flush", &Connection::flush, nb::arg("timeoutUsec") = (int64_t)3e7)
+      .def("flush", gil_release_wrapper(&Connection::flush), nb::arg("timeoutUsec") = (int64_t)3e7)
       .def("transport", &Connection::transport)
       .def("remote_transport", &Connection::remoteTransport);
 

diff --git a/python/mscclpp/gil_release.hpp b/python/mscclpp/gil_release.hpp
@@ -0,0 +1,18 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+#ifndef MSCCLPP_PYTHON_GIL_RELEASE_HPP_
+#define MSCCLPP_PYTHON_GIL_RELEASE_HPP_
+
+#include <functional>
+#include <nanobind/nanobind.h>
+
+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)...);
+    };
+}
+
+#endif  // MSCCLPP_PYTHON_GIL_RELEASE_HPP_