ensure raft-dask wheel tests install pylibraft wheel from the same CI…

… run, fix wheel dependencies (#2349)

Fixes #2348 

#2331 introduced `rapids-build-backend` (https://github.com/rapidsai/rapids-build-backend) as the build backend for `pylibraft`, `raft-dask`, and `raft-ann-bench`.

That library handles automatically modifying a wheel's dependencies based on the target CUDA version. Unfortunately, we missed a few cases in #2331, and as a result the last few days of nightly `raft-dask` wheels had the following issues:

* depending on `pylibraft`
  - *(does not exist, it's called `pylibraft-cu12`)*
* depending on `ucx-py==0.39.*`
  - *(does not exist, it's called `ucx-py-cu12`)*
* depending on `distributed-ucxx-cu11==0.39.*` instead of `distributed-ucxx-cu11==0.39.*,>=0.0.0a0`
   - *(without that alpha specifier, `pip install --pre` is required to install pre-release versions)*

This wasn't caught in `raft`'s CI, but in downstream CI like `cuml` and `cugraph`, with errors like this:

```text
ERROR: ResolutionImpossible:
  
The conflict is caused by:
    raft-dask-cu12 24.8.0a20 depends on pylibraft==24.8.* and >=0.0.0a0
    raft-dask-cu12 24.8.0a19 depends on pylibraft==24.8.* and >=0.0.0a0
```

([example cugraph build](https://github.com/rapidsai/cugraph/actions/runs/9315062495/job/25656684762?pr=4454#step:7:1811))

This PR:

* fixes those dependency issues
* modifies `raft`'s CI so that similar issues would be caught here in the future, before publishing wheels

## Notes for Reviewers

### What was the root cause of CI missing this, and how does this PR fix it?

The `raft-dask` test CI jobs use this pattern to install the `raft-dask` wheel built earlier in the CI pipeline.

```shell
pip install "raft_dask-cu12[test]>=0.0.0a0" --find-links dist/
```

As described in the `pip` docs ([link](https://pip.pypa.io/en/stable/cli/pip_install/#finding-packages)), `--find-links` just adds a directory to the list of other places `pip` searches for packages. Because the wheel there had unsatisfiable constraints (e.g. `pylibraft==24.8.*` does not exist anywhere), `pip install` silently disregarded that locally-downloaded `raft_dask` wheel and backtracked (i.e. downloaded older and older wheels from https://pypi.anaconda.org/rapidsai-wheels-nightly/simple/) until it found one that wasn't problematic.

This PR ensures that won't happen by telling `pip` to install **exactly that locally-downloaded file**, like this

```shell
pip install "$(echo ./dist/raft_dask_cu12*.whl)[test]"
```

If that file is uninstallable, `pip install` fails and you find out via a CI failure.

### How I tested this

Initially pushed a commit with just the changes to the test script. Saw the `wheel-tests-raft-dask` CI jobs fail in the expected way, instead of silently falling back to an older wheel and passing 🎉 .

```text
ERROR: Could not find a version that satisfies the requirement ucx-py-cu12==0.39.* (from raft-dask-cu12[test]) (from versions: 0.32.0, 0.33.0, 0.34.0, 0.35.0, 0.36.0, 0.37.0, 0.38.0a4, 0.38.0a5, 0.38.0a6, 0.39.0a0)
ERROR: No matching distribution found for ucx-py-cu12==0.39.*
```

([build link](https://github.com/rapidsai/raft/actions/runs/9323598882/job/25668146747?pr=2349))

Authors:
  - James Lamb (https://github.com/jameslamb)

Approvers:
  - Bradley Dice (https://github.com/bdice)
  - Dante Gama Dessavre (https://github.com/dantegd)

URL: #2349

ci/test_wheel_raft_dask.sh

@@ -11,7 +11,8 @@ RAPIDS_PY_WHEEL_NAME="raft_dask_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels
 RAPIDS_PY_WHEEL_NAME="pylibraft_${RAPIDS_PY_CUDA_SUFFIX}" rapids-download-wheels-from-s3 ./local-pylibraft-dep
 python -m pip install --no-deps ./local-pylibraft-dep/pylibraft*.whl
 
-python -m pip install "raft_dask-${RAPIDS_PY_CUDA_SUFFIX}[test]>=0.0.0a0" --find-links dist/
+# echo to expand wildcard before adding `[extra]` requires for pip
+python -m pip install -v "$(echo ./dist/raft_dask_${RAPIDS_PY_CUDA_SUFFIX}*.whl)[test]"
 
 test_dir="python/raft-dask/raft_dask/test"
 

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

@@ -21,7 +21,7 @@ dependencies:
 - cxx-compiler
 - cython>=3.0.0
 - dask-cuda==24.8.*,>=0.0.0a0
-- distributed-ucxx==0.39.*
+- distributed-ucxx==0.39.*,>=0.0.0a0
 - doxygen>=1.8.20
 - gcc_linux-aarch64=11.*
 - graphviz
@@ -44,6 +44,7 @@ dependencies:
 - nvcc_linux-aarch64=11.8
 - pre-commit
 - pydata-sphinx-theme
+- pylibraft==24.8.*,>=0.0.0a0
 - pytest-cov
 - pytest==7.*
 - rapids-build-backend>=0.3.0,<0.4.0.dev0
@@ -56,5 +57,5 @@ dependencies:
 - sphinx-copybutton
 - sphinx-markdown-tables
 - sysroot_linux-aarch64==2.17
-- ucx-py==0.39.*
+- ucx-py==0.39.*,>=0.0.0a0
 name: all_cuda-118_arch-aarch64

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

@@ -21,7 +21,7 @@ dependencies:
 - cxx-compiler
 - cython>=3.0.0
 - dask-cuda==24.8.*,>=0.0.0a0
-- distributed-ucxx==0.39.*
+- distributed-ucxx==0.39.*,>=0.0.0a0
 - doxygen>=1.8.20
 - gcc_linux-64=11.*
 - graphviz
@@ -44,6 +44,7 @@ dependencies:
 - nvcc_linux-64=11.8
 - pre-commit
 - pydata-sphinx-theme
+- pylibraft==24.8.*,>=0.0.0a0
 - pytest-cov
 - pytest==7.*
 - rapids-build-backend>=0.3.0,<0.4.0.dev0
@@ -56,5 +57,5 @@ dependencies:
 - sphinx-copybutton
 - sphinx-markdown-tables
 - sysroot_linux-64==2.17
-- ucx-py==0.39.*
+- ucx-py==0.39.*,>=0.0.0a0
 name: all_cuda-118_arch-x86_64

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

@@ -22,7 +22,7 @@ dependencies:
 - cxx-compiler
 - cython>=3.0.0
 - dask-cuda==24.8.*,>=0.0.0a0
-- distributed-ucxx==0.39.*
+- distributed-ucxx==0.39.*,>=0.0.0a0
 - doxygen>=1.8.20
 - gcc_linux-aarch64=11.*
 - graphviz
@@ -40,6 +40,7 @@ dependencies:
 - numpydoc
 - pre-commit
 - pydata-sphinx-theme
+- pylibraft==24.8.*,>=0.0.0a0
 - pytest-cov
 - pytest==7.*
 - rapids-build-backend>=0.3.0,<0.4.0.dev0
@@ -52,5 +53,5 @@ dependencies:
 - sphinx-copybutton
 - sphinx-markdown-tables
 - sysroot_linux-aarch64==2.17
-- ucx-py==0.39.*
+- ucx-py==0.39.*,>=0.0.0a0
 name: all_cuda-122_arch-aarch64

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

@@ -22,7 +22,7 @@ dependencies:
 - cxx-compiler
 - cython>=3.0.0
 - dask-cuda==24.8.*,>=0.0.0a0
-- distributed-ucxx==0.39.*
+- distributed-ucxx==0.39.*,>=0.0.0a0
 - doxygen>=1.8.20
 - gcc_linux-64=11.*
 - graphviz
@@ -40,6 +40,7 @@ dependencies:
 - numpydoc
 - pre-commit
 - pydata-sphinx-theme
+- pylibraft==24.8.*,>=0.0.0a0
 - pytest-cov
 - pytest==7.*
 - rapids-build-backend>=0.3.0,<0.4.0.dev0
@@ -52,5 +53,5 @@ dependencies:
 - sphinx-copybutton
 - sphinx-markdown-tables
 - sysroot_linux-64==2.17
-- ucx-py==0.39.*
+- ucx-py==0.39.*,>=0.0.0a0
 name: all_cuda-122_arch-x86_64

dependencies.yaml

@@ -475,13 +475,10 @@ dependencies:
           - numba>=0.57
           - *numpy
           - rapids-dask-dependency==24.8.*,>=0.0.0a0
-          - ucx-py==0.39.*
       - output_types: conda
-        packages:
-          - &ucx_py_conda ucx-py==0.39.*
-      - output_types: pyproject
         packages:
           - &pylibraft_conda pylibraft==24.8.*,>=0.0.0a0
+          - &ucx_py_conda ucx-py==0.39.*,>=0.0.0a0
       - output_types: requirements
         packages:
           # pip recognizes the index as a global option for the requirements.txt file
@@ -494,11 +491,11 @@ dependencies:
           - matrix: {cuda: "12.*"}
             packages:
               - &pylibraft_cu12 pylibraft-cu12==24.8.*,>=0.0.0a0
-              - &ucx_py_cu12 ucx-py-cu12==0.39.*
+              - &ucx_py_cu12 ucx-py-cu12==0.39.*,>=0.0.0a0
           - matrix: {cuda: "11.*"}
             packages:
               - &pylibraft_cu11 pylibraft-cu11==24.8.*,>=0.0.0a0
-              - &ucx_py_cu11 ucx-py-cu11==0.39.*
+              - &ucx_py_cu11 ucx-py-cu11==0.39.*,>=0.0.0a0
           - {matrix: null, packages: [*pylibraft_conda, *ucx_py_conda]}
   test_python_common:
     common:
@@ -518,7 +515,7 @@ dependencies:
         packages:
           # UCXX is not currently a hard-dependency thus only installed during tests,
           # this will change in the future.
-          - &distributed_ucxx_conda distributed-ucxx==0.39.*
+          - &distributed_ucxx_conda distributed-ucxx==0.39.*,>=0.0.0a0
       - output_types: requirements
         packages:
           # pip recognizes the index as a global option for the requirements.txt file
@@ -529,10 +526,10 @@ dependencies:
         matrices:
           - matrix: {cuda: "12.*"}
             packages:
-              - distributed-ucxx-cu12==0.39.*
+              - distributed-ucxx-cu12==0.39.*,>=0.0.0a0
           - matrix: {cuda: "11.*"}
             packages:
-              - distributed-ucxx-cu11==0.39.*
+              - distributed-ucxx-cu11==0.39.*,>=0.0.0a0
           - {matrix: null, packages: [*distributed_ucxx_conda]}
   depends_on_ucx_build:
     common:

python/raft-dask/pyproject.toml

@@ -32,13 +32,13 @@ license = { text = "Apache 2.0" }
 requires-python = ">=3.9"
 dependencies = [
     "dask-cuda==24.8.*,>=0.0.0a0",
-    "distributed-ucxx==0.39.*",
+    "distributed-ucxx==0.39.*,>=0.0.0a0",
     "joblib>=0.11",
     "numba>=0.57",
     "numpy>=1.23,<2.0a0",
     "pylibraft==24.8.*,>=0.0.0a0",
     "rapids-dask-dependency==24.8.*,>=0.0.0a0",
-    "ucx-py==0.39.*",
+    "ucx-py==0.39.*,>=0.0.0a0",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 classifiers = [
     "Intended Audience :: Developers",