Make gloo context setTimeout getTimeout callable from Python

For now, if we directly call `setTimeout` or `getTimeout` from Python, it will throw the parameter unmatch error, this PR addresses this by use a wrapper around the original function. Signed-off-by: Hollow Man <[email protected]>
ray-project · Feb 4, 2025 · 742c8d6 · 742c8d6
1 parent 49723ee
commit 742c8d6
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/pygloo/main.cc b/pygloo/main.cc
@@ -8,11 +8,20 @@
 #include <collective.h>
 #include <rendezvous.h>
 #include <sstream>
+#include <chrono>
 
 namespace pygloo {
 bool transport_tcp_available() { return GLOO_HAVE_TRANSPORT_TCP; }
 
 bool transport_uv_available() { return GLOO_HAVE_TRANSPORT_UV; }
+
+void setTimeout_wrapper(gloo::Context* context, int timeout_ms) {
+    context->setTimeout(std::chrono::milliseconds(timeout_ms));
+}
+
+int64_t getTimeout_wrapper(gloo::Context* context) {
+    return context->getTimeout().count();
+}
 } // namespace pygloo
 
 PYBIND11_MODULE(pygloo, m) {
@@ -124,8 +133,8 @@ PYBIND11_MODULE(pygloo, m) {
       .def("createUnboundBuffer", &gloo::Context::createUnboundBuffer)
       .def("nextSlot", &gloo::Context::nextSlot)
       .def("closeConnections", &gloo::Context::closeConnections)
-      .def("setTimeout", &gloo::Context::setTimeout)
-      .def("getTimeout", &gloo::Context::getTimeout);
+      .def("setTimeout", &pygloo::setTimeout_wrapper)
+      .def("getTimeout", &pygloo::getTimeout_wrapper);
 
   pygloo::transport::def_transport_module(m);
   pygloo::rendezvous::def_rendezvous_module(m);