Increase maximum UCX runtime pin

ci/test_python.sh

@@ -56,4 +56,4 @@ for nbuf in 1 8; do
 done
 
 rapids-logger "C++ future -> Python future notifier example"
-python -m ucxx.examples.python_future_task_example
+timeout 1m python -m ucxx.examples.python_future_task_example

conda/environments/all_cuda-118_arch-x86_64.yaml

@@ -41,5 +41,5 @@ dependencies:
 - scikit-build-core>=0.7.0
 - setuptools>=64.0.0
 - spdlog>=1.12.0,<1.13
-- ucx>=1.15.0
+- ucx>=1.15.0,<1.18
 name: all_cuda-118_arch-x86_64

conda/environments/all_cuda-122_arch-x86_64.yaml

@@ -41,5 +41,5 @@ dependencies:
 - scikit-build-core>=0.7.0
 - setuptools>=64.0.0
 - spdlog>=1.12.0,<1.13
-- ucx>=1.15.0
+- ucx>=1.15.0,<1.18
 name: all_cuda-122_arch-x86_64

conda/recipes/ucxx/meta.yaml

@@ -91,7 +91,7 @@ outputs:
         - cudatoolkit
         {% endif %}
         - {{ pin_compatible('cuda-version', max_pin='x', min_pin='x') }}
-        - ucx >=1.15.0,<1.16.0
+        - ucx >=1.15.0,<1.18.0
     test:
       commands:
         - test -f $PREFIX/lib/libucxx.so
@@ -237,7 +237,7 @@ outputs:
         {% else %}
         - cuda-cudart
         {% endif %}
-        - ucx >=1.15.0,<1.16.0
+        - ucx >=1.15.0,<1.18.0
         - {{ pin_subpackage('libucxx', exact=True) }}
         - {{ pin_compatible('rmm', max_pin='x.x') }}
         - {{ pin_compatible('numpy') }}

dependencies.yaml

@@ -374,13 +374,13 @@ dependencies:
             packages:
               - libucx-cu11==1.15.0
           - matrix: null
-            packages: 
+            packages:
               - libucx==1.15.0
   depends_on_ucx_run:
     common:
       - output_types: conda
         packages:
-          - &ucx_conda_run ucx>=1.15.0
+          - &ucx_conda_run ucx>=1.15.0,<1.18
       - output_types: requirements
         packages:
           # pip recognizes the index as a global option for the requirements.txt file
@@ -391,10 +391,10 @@ dependencies:
         matrices:
           - matrix: {cuda: "12.*"}
             packages:
-              - libucx-cu12>=1.15.0
+              - libucx-cu12>=1.15.0,<1.18
           - matrix: {cuda: "11.*"}
             packages:
-              - libucx-cu11>=1.15.0
+              - libucx-cu11>=1.15.0,<1.18
           - matrix: null
-            packages: 
-              - libucx>=1.15.0
+            packages:
+              - libucx>=1.15.0,<1.18

python/pyproject.toml

@@ -19,7 +19,7 @@ authors = [
 license = { text = "BSD-3-Clause" }
 requires-python = ">=3.9"
 dependencies = [
-    "libucx>=1.15.0",
+    "libucx>=1.15.0,<1.18",
     "numpy>=1.23,<2.0a0",
     "pynvml>=11.4.1",
     "rmm==24.8.*,>=0.0.0a0",

python/ucxx/__init__.py

@@ -95,6 +95,11 @@ def _is_mig_device(handle):
     logger.info("Setting UCX_MAX_RNDV_RAILS=1")
     os.environ["UCX_MAX_RNDV_RAILS"] = "1"
 
+if "UCX_PROTO_ENABLE" not in os.environ:
+    # UCX protov2 still doesn't support CUDA async/managed memory
+    logger.info("Setting UCX_PROTO_ENABLE=n")
+    os.environ["UCX_PROTO_ENABLE"] = "n"
+
 
 from ._version import __git_commit__, __version__