From aedf931cdb4d368998572abf34d1ce6438b45852 Mon Sep 17 00:00:00 2001
From: Peter Andreas Entschev <peter@entschev.com>
Date: Fri, 3 May 2024 08:52:02 -0700
Subject: [PATCH] Fix UCX wheels dependencies

---
 dependencies.yaml               | 47 ++++++++++++++++++++++++++++++---
 python/raft-dask/pyproject.toml |  2 +-
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/dependencies.yaml b/dependencies.yaml
index 1a73cd8cd6..a336aa1577 100644
--- 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
index a48f23f364..0181bef4ce 100644
--- 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`.