ensure raft-dask wheel tests install pylibraft wheel from the same CI run, fix wheel dependencies

[jameslamb]

Description

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:

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)

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.

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

As described in the pip docs (link), --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

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 🎉 .

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)

[dantegd]

@jameslamb @bdice This PR should unblock CI for downstream repos, and ready to merge, right?

[bdice]

@dantegd @jameslamb Yes, I think so. I’ll go ahead and merge so we can unblock downstream work.

[bdice]

/merge

[jameslamb]

Yes it should, thanks for merging it @bdice

[nv-rliu]

Thanks everyone! 👍