introduce rapids-init-pip

[jameslamb]

Splitting this off from #173

I've been working on updating the documentation for how to accomplish something like "use packages from a cuvs PR in the CI on a cuml PR" (rapidsai/docs#601).

Downloading the wheel artifacts from one PR in another PR's CI is pretty much the same as it was before with downloads.rapids.ai, thanks to @VenkateshJaya 's work on rapids-get-pr-artifact-github.

But once you have those downloaded wheels, it can be challenging to get all of the wheel-related things to use them. We have, across RAPIDS CI:

• pip wheel builds with build isolation
• pip wheel builds WITHOUT build isolation
• pip install calls with --constraint, pointing to custom constraints files (e.g. those created with rapids-generate-pip-constraints)
• all of the above, constrained by the PIP_CONSTRAINT environment variable

This PR proposes introducing a small new tool called rapids-pip-init, which does the following:

• ensures that environment variable PIP_CONSTRAINT is set
• ensures that the file it points to exists
• does not modify PIP_CONSTRAINT if it is already set

Proposing that in every CI script that installs wheels (like ci/build_*_wheel.sh and ci/test_wheel_*.sh) source this near the top:

source rapids-pip-init

And that every CI script passing --constraint to pip explicitly add --constraint "${PIP_CONSTRAINT}" (even if during normal operation, the file PIP_CONSTRAINT points to is empty).

With that in place, it'd then be possible to reliably constrain all pip invocations to use a certain set of wheels (e.g. those from upstream projects' CI) via something like the following:

source rapids-init-pip

RAPIDS_PY_CUDA_SUFFIX=$(rapids-wheel-ctk-name-gen "${RAPIDS_CUDA_VERSION}")

# download wheels, store the directories holding them in variables
LIBRMM_WHEELHOUSE=$(
  RAPIDS_PY_WHEEL_NAME="librmm_${RAPIDS_PY_CUDA_SUFFIX}" rapids-get-pr-wheel-artifact rmm 1909 cpp
)

# write a pip constraints file saying e.g. "whenever you encounter a requirement for 'librmm-cu12', use this wheel"
cat > "${PIP_CONSTRAINT}" <<EOF
librmm-${RAPIDS_PY_CUDA_SUFFIX} @ file://$(echo ${LIBRMM_WHEELHOUSE}/librmm_*.whl)
EOF

Notes for Reviewers

Why call this rapids-init-pip?

I went with a generic-ish name instead of something like rapids-init-pip-constraint-config so that it's extensible in the future to other initialization we might want to do in wheel jobs.

How I tested this

On a cudf PR: NVIDIA/cudf#18747 (comment)

Next steps

If folks agree with doing this, I'll put up an issue in https://github.com/rapidsai/build-planning about adding calls like this across RAPIDS wheel CI scripts.

[bdice]

Design seems solid to me.

[jameslamb]

Alright thanks! I'm gonna merge this, we can pick up other design work and deployment over in this tracking issue: rapidsai/build-planning#179