Add UCX-Py connect timeout patch

rapidsai · Jan 21, 2025 · 35a3aa6 · 35a3aa6
1 parent 29af6cc
commit 35a3aa6
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/rapids_dask_dependency/patches/distributed/ucx_connect_timeout.py b/rapids_dask_dependency/patches/distributed/ucx_connect_timeout.py
@@ -0,0 +1,10 @@
+import distributed.comm.ucx
+
+init_once_original = distributed.comm.ucx.init_once
+
+def init_once_patched():
+    print("PATCHED!", flush=True)
+    init_once_original()
+
+print("Patching...", flush=True)
+distributed.comm.ucx.init_once = init_once_patched