|
| 1 | +#!/bin/bash |
| 2 | +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 3 | +# SPDX-License-Identifier: Apache-2.0 |
| 4 | + |
| 5 | +# Install the built cuVS Conda packages with only their runtime dependencies |
| 6 | +# and import them on a CPU-only runner. This replaces the import tests that |
| 7 | +# used to run inside each recipe build, taking them off the critical path of |
| 8 | +# the GPU test jobs, and checks that nothing links against the CUDA driver or |
| 9 | +# initializes CUDA at import time. |
| 10 | + |
| 11 | +set -euo pipefail |
| 12 | + |
| 13 | +source /opt/conda/etc/profile.d/conda.sh |
| 14 | + |
| 15 | +rapids-logger "Downloading built Conda packages" |
| 16 | +CPP_CHANNEL=$(rapids-download-from-github "$(rapids-artifact-name conda_cpp libcuvs cuvs --cuda "$RAPIDS_CUDA_VERSION")") |
| 17 | +PYTHON_CHANNEL=$(rapids-download-from-github "$(rapids-artifact-name conda_python cuvs cuvs --stable --cuda "$RAPIDS_CUDA_VERSION")") |
| 18 | + |
| 19 | +# Create an environment holding the given built packages and their runtime |
| 20 | +# dependencies, then import the given Python modules in it. |
| 21 | +# |
| 22 | +# usage: import_packages <env name> <package glob> [<package glob> ...] -- <module> [<module> ...] |
| 23 | +import_packages() { |
| 24 | + local env_name="$1" |
| 25 | + shift |
| 26 | + local packages=() |
| 27 | + while [[ "$1" != "--" ]]; do |
| 28 | + packages+=("$1") |
| 29 | + shift |
| 30 | + done |
| 31 | + shift |
| 32 | + |
| 33 | + rapids-logger "Install ${packages[*]} with only their runtime dependencies" |
| 34 | + rapids-mamba-retry create --yes --name "${env_name}" --strict-channel-priority \ |
| 35 | + --override-channels \ |
| 36 | + --channel "${PYTHON_CHANNEL}" --channel "${CPP_CHANNEL}" \ |
| 37 | + --channel rapidsai-nightly --channel conda-forge \ |
| 38 | + "python=${RAPIDS_PY_VERSION}" "${packages[@]}" |
| 39 | + |
| 40 | + set +u |
| 41 | + conda activate "${env_name}" |
| 42 | + set -u |
| 43 | + rapids-print-env |
| 44 | + |
| 45 | + for module in "$@"; do |
| 46 | + rapids-logger "Import ${module} on a CPU-only runner" |
| 47 | + python -c "import ${module}" |
| 48 | + done |
| 49 | + |
| 50 | + rapids-logger "Check that cuvs_bench ships its dataset configuration" |
| 51 | + python -c "import cuvs_bench, pathlib; path = pathlib.Path(cuvs_bench.__file__).parent / 'config' / 'datasets' / 'datasets.yaml'; assert path.is_file(), path" |
| 52 | + |
| 53 | + set +u |
| 54 | + conda deactivate |
| 55 | + set -u |
| 56 | +} |
| 57 | + |
| 58 | +# cuvs-bench and cuvs-bench-cpu both provide the cuvs_bench module, so they |
| 59 | +# are checked in separate environments. |
| 60 | +import_packages python_imports \ |
| 61 | + "${PYTHON_CHANNEL}"/linux-64/cuvs-[0-9]*.conda \ |
| 62 | + "${PYTHON_CHANNEL}"/linux-64/cuvs-bench-[0-9]*.conda \ |
| 63 | + -- \ |
| 64 | + cuvs \ |
| 65 | + cuvs.cluster \ |
| 66 | + cuvs.common \ |
| 67 | + cuvs.distance \ |
| 68 | + cuvs.neighbors \ |
| 69 | + cuvs.preprocessing \ |
| 70 | + cuvs_bench \ |
| 71 | + cuvs_bench.generate_groundtruth \ |
| 72 | + cuvs_bench.get_dataset \ |
| 73 | + cuvs_bench.plot \ |
| 74 | + cuvs_bench.run \ |
| 75 | + cuvs_bench.split_groundtruth |
| 76 | + |
| 77 | +# cuvs-bench-cpu is only built for one CUDA major version (see ci/build_python.sh). |
| 78 | +if compgen -G "${PYTHON_CHANNEL}/linux-64/cuvs-bench-cpu-[0-9]*.conda" > /dev/null; then |
| 79 | + import_packages python_imports_cpu \ |
| 80 | + "${PYTHON_CHANNEL}"/linux-64/cuvs-bench-cpu-[0-9]*.conda \ |
| 81 | + -- \ |
| 82 | + cuvs_bench \ |
| 83 | + cuvs_bench.generate_groundtruth \ |
| 84 | + cuvs_bench.get_dataset \ |
| 85 | + cuvs_bench.plot \ |
| 86 | + cuvs_bench.run \ |
| 87 | + cuvs_bench.split_groundtruth |
| 88 | +else |
| 89 | + rapids-logger "No cuvs-bench-cpu package was built for CUDA ${RAPIDS_CUDA_VERSION}; skipping" |
| 90 | +fi |
0 commit comments