Fix UCX wheels dependencies

rapidsai · May 3, 2024 · aedf931 · aedf931
1 parent 068cb91
commit aedf931
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 5 deletions.
diff --git a/dependencies.yaml b/dependencies.yaml
@@ -92,7 +92,7 @@ files:
       table: build-system
     includes:
       - build
-      - ucx_wheel
+      - depends_on_ucx_build
   py_run_raft_dask:
     output: pyproject
     pyproject_dir: python/raft-dask
@@ -108,8 +108,8 @@ files:
       key: test
     includes:
       - test_python_common
-      - ucx_wheel
       - depends_on_distributed_ucxx
+      - depends_on_ucx_run
   py_build_raft_ann_bench:
     output: pyproject
     pyproject_dir: python/raft-ann-bench
@@ -505,8 +505,47 @@ dependencies:
             packages:
               - distributed-ucxx-cu11==0.38.*
           - {matrix: null, packages: [*distributed_ucxx_conda]}
-  ucx_wheel:
+  depends_on_ucx_build:
     common:
+      - output_types: conda
+        packages:
+          - &ucx_conda_build ucx==1.15.0
+      - output_types: requirements
+        packages:
+          # pip recognizes the index as a global option for the requirements.txt file
+          - --extra-index-url=https://pypi.nvidia.com
+          - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple
+    specific:
       - output_types: [requirements, pyproject]
+        matrices:
+          - matrix: {cuda: "12.*"}
+            packages:
+              - libucx-cu12==1.15.0
+          - matrix: {cuda: "11.*"}
+            packages:
+              - libucx-cu11==1.15.0
+          - matrix: null
+            packages:
+              - libucx==1.15.0
+  depends_on_ucx_run:
+    common:
+      - output_types: conda
+        packages:
+          - &ucx_conda_run ucx>=1.15.0
+      - output_types: requirements
         packages:
-          - libucx==1.15.0
+          # pip recognizes the index as a global option for the requirements.txt file
+          - --extra-index-url=https://pypi.nvidia.com
+          - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple
+    specific:
+      - output_types: [requirements, pyproject]
+        matrices:
+          - matrix: {cuda: "12.*"}
+            packages:
+              - libucx-cu12>=1.15.0
+          - matrix: {cuda: "11.*"}
+            packages:
+              - libucx-cu11>=1.15.0
+          - matrix: null
+            packages:
+              - libucx>=1.15.0
diff --git a/python/raft-dask/pyproject.toml b/python/raft-dask/pyproject.toml
@@ -53,7 +53,7 @@ classifiers = [
 [project.optional-dependencies]
 test = [
     "distributed-ucxx==0.38.*",
-    "libucx==1.15.0",
+    "libucx>=1.15.0",
     "pytest-cov",
     "pytest==7.*",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.