Skip to content

Commit 7d36fea

Browse files
leofangclaude
andcommitted
Consolidate core and nightly-* install logic in run-tests
Merge the duplicated variable setup (TEST_CUDA_MAJOR, CUDA_VER_MINOR, FREE_THREADING, BINDINGS_ARGS, CORE_WHL) into a shared block for both core and nightly-* modes. Use nightly-* glob pattern for validation, pathfinder guard, and branch dispatch for future extensibility. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 4d1c871 commit 7d36fea

1 file changed

Lines changed: 66 additions & 88 deletions

File tree

ci/tools/run-tests

Lines changed: 66 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@ if [[ ${#} -ne 1 ]]; then
1313
echo "Error: This script requires exactly 1 argument. You provided ${#}"
1414
exit 1
1515
fi
16-
if [[ "${1}" != "bindings" && "${1}" != "core" && "${1}" != "pathfinder" && "${1}" != "nightly-pytorch" && "${1}" != "nightly-numba-cuda" ]]; then
17-
echo "Error: Invalid test module '${1}'. Must be 'bindings', 'core', 'pathfinder', 'nightly-pytorch', or 'nightly-numba-cuda'"
16+
if [[ "${1}" != "bindings" && "${1}" != "core" && "${1}" != "pathfinder" && "${1}" != nightly-* ]]; then
17+
echo "Error: Invalid test module '${1}'. Must be 'bindings', 'core', 'pathfinder', or 'nightly-*'"
1818
exit 1
1919
fi
2020

@@ -23,7 +23,7 @@ test_module=${1}
2323
# For standard modes, install pathfinder up front (it is a direct dependency
2424
# of bindings, and a transitive dependency of core). Nightly modes install
2525
# all wheels together in a single pip call further below.
26-
if [[ "${test_module}" != "nightly-pytorch" && "${test_module}" != "nightly-numba-cuda" ]]; then
26+
if [[ "${test_module}" != nightly-* ]]; then
2727
pushd ./cuda_pathfinder
2828
echo "Installing pathfinder wheel"
2929
pip install ./*.whl --group test
@@ -56,54 +56,8 @@ elif [[ "${test_module}" == "bindings" ]]; then
5656
${SANITIZER_CMD} pytest -rxXs -v --durations=0 --randomly-dont-reorganize tests/cython
5757
fi
5858
popd
59-
elif [[ "${test_module}" == "core" ]]; then
60-
# Install cuda.bindings for core tests based on BINDINGS_SOURCE.
61-
if [[ "${BINDINGS_SOURCE}" == "published" ]]; then
62-
echo "Installing published cuda-bindings==${TEST_CUDA_MAJOR}.${TEST_CUDA_MINOR}.* from PyPI"
63-
pip install "cuda-bindings==${TEST_CUDA_MAJOR}.${TEST_CUDA_MINOR}.*"
64-
elif [[ "${BINDINGS_SOURCE}" == "backport" || "${BINDINGS_SOURCE}" == "main" ]]; then
65-
echo "Installing bindings wheel (source: ${BINDINGS_SOURCE})"
66-
if [[ "${LOCAL_CTK}" == 1 ]]; then
67-
pip install "${CUDA_BINDINGS_ARTIFACTS_DIR}"/*.whl
68-
else
69-
pip install $(ls "${CUDA_BINDINGS_ARTIFACTS_DIR}"/*.whl)[all]
70-
fi
71-
fi
72-
TEST_CUDA_MAJOR="$(cut -d '.' -f 1 <<< ${CUDA_VER})"
73-
echo "Installing core wheel"
74-
75-
FREE_THREADING=""
76-
if python -c 'import sys; assert not sys._is_gil_enabled()' 2> /dev/null; then
77-
FREE_THREADING+="-ft"
78-
fi
79-
80-
pushd ./cuda_core
81-
CUDA_VER_MINOR="$(cut -d '.' -f 1-2 <<< "${CUDA_VER}")"
82-
# Start from the built wheel path, then add the published cuda.bindings extra
83-
# when this job is resolving against wheel-installed CTK packages.
84-
WHL_EXTRA=("${CUDA_CORE_ARTIFACTS_DIR}"/*.whl)
85-
if [[ "${LOCAL_CTK}" != 1 ]]; then
86-
WHL_EXTRA=("${WHL_EXTRA[0]}[cu${TEST_CUDA_MAJOR}]")
87-
fi
88-
# Constrain cuda-toolkit to the requested CTK version to avoid
89-
# pip pulling in a newer nvidia-cuda-runtime that conflicts with it.
90-
pip install "${WHL_EXTRA[@]}" --group "test-cu${TEST_CUDA_MAJOR}${FREE_THREADING}" "cuda-toolkit==${CUDA_VER_MINOR}.*"
91-
echo "Installed packages before core tests:"
92-
pip list
93-
echo "Running core tests"
94-
${SANITIZER_CMD} pytest -rxXs -v --durations=0 --randomly-dont-reorganize tests/
95-
# Currently our CI always installs the latest bindings (from either major version).
96-
# This is not compatible with the test requirements.
97-
if [[ "${SKIP_CYTHON_TEST}" == 0 ]]; then
98-
${SANITIZER_CMD} pytest -rxXs -v --durations=0 --randomly-dont-reorganize tests/cython
99-
fi
100-
popd
101-
elif [[ "${test_module}" == "nightly-pytorch" || "${test_module}" == "nightly-numba-cuda" ]]; then
102-
# Nightly optional-dependency testing.
103-
# Install ALL wheels (pathfinder + bindings + core) and the optional dep
104-
# in a single pip call so pip resolves version constraints in one shot
105-
# and avoids costly uninstall/reinstall cycles.
106-
59+
elif [[ "${test_module}" == "core" || "${test_module}" == nightly-* ]]; then
60+
# Shared setup for core and nightly modes.
10761
TEST_CUDA_MAJOR="$(cut -d '.' -f 1 <<< ${CUDA_VER})"
10862
CUDA_VER_MINOR="$(cut -d '.' -f 1-2 <<< "${CUDA_VER}")"
10963

@@ -112,12 +66,7 @@ elif [[ "${test_module}" == "nightly-pytorch" || "${test_module}" == "nightly-nu
11266
FREE_THREADING+="-ft"
11367
fi
11468

115-
# Resolve pathfinder wheel to absolute path before pushd.
116-
# CUDA_BINDINGS_ARTIFACTS_DIR and CUDA_CORE_ARTIFACTS_DIR are already
117-
# absolute (set via realpath in env-vars).
118-
PATHFINDER_WHL=($(realpath ./cuda_pathfinder/*.whl))
119-
120-
# Build bindings spec based on BINDINGS_SOURCE (set by env-vars):
69+
# Resolve bindings based on BINDINGS_SOURCE (set by env-vars):
12170
# main/backport → local wheel from artifacts dir
12271
# published → install from PyPI by version
12372
BINDINGS_ARGS=()
@@ -130,48 +79,77 @@ elif [[ "${test_module}" == "nightly-pytorch" || "${test_module}" == "nightly-nu
13079
fi
13180
fi
13281

82+
# Resolve core wheel, adding the published cuda.bindings extra
83+
# when this job is resolving against wheel-installed CTK packages.
13384
CORE_WHL=("${CUDA_CORE_ARTIFACTS_DIR}"/*.whl)
13485
if [[ "${LOCAL_CTK}" != 1 ]]; then
13586
CORE_WHL=("${CORE_WHL[0]}[cu${TEST_CUDA_MAJOR}]")
13687
fi
13788

89+
if [[ "${test_module}" == nightly-* ]]; then
90+
# Resolve pathfinder wheel to absolute path before pushd.
91+
# CUDA_BINDINGS_ARTIFACTS_DIR and CUDA_CORE_ARTIFACTS_DIR are already
92+
# absolute (set via realpath in env-vars).
93+
PATHFINDER_WHL=($(realpath ./cuda_pathfinder/*.whl))
94+
fi
95+
13896
# pushd so --group reads test dependency groups from cuda_core/pyproject.toml.
139-
# The explicit cuda-toolkit[...]==X.Y.* pin overrides the group's looser ==X.*.
14097
pushd ./cuda_core
14198

142-
PIP_ARGS=(
143-
"${PATHFINDER_WHL[@]}"
144-
"${BINDINGS_ARGS[@]}"
145-
"${CORE_WHL[@]}"
146-
--group "test-cu${TEST_CUDA_MAJOR}${FREE_THREADING}"
147-
)
148-
149-
if [[ "${test_module}" == "nightly-pytorch" ]]; then
150-
# TORCH_VER and TORCH_CUDA must be set by the caller.
151-
# Use cuda-toolkit[cudart] only — torch brings its own nvcc/nvrtc/etc.
152-
# This avoids version conflicts between our nvidia-* pins and torch's.
153-
echo "Installing pathfinder + bindings + core + test deps + PyTorch ${TORCH_VER} (${TORCH_CUDA})"
154-
PIP_ARGS+=("cuda-toolkit[cudart]==${CUDA_VER_MINOR}.*")
155-
if [[ "${TORCH_VER}" == "latest" ]]; then
156-
PIP_ARGS+=(torch)
157-
else
158-
PIP_ARGS+=("torch==${TORCH_VER}")
99+
if [[ "${test_module}" == "core" ]]; then
100+
echo "Installing bindings (source: ${BINDINGS_SOURCE})"
101+
pip install "${BINDINGS_ARGS[@]}"
102+
echo "Installing core wheel"
103+
# Constrain cuda-toolkit to the requested CTK version to avoid
104+
# pip pulling in a newer nvidia-cuda-runtime that conflicts with it.
105+
pip install "${CORE_WHL[@]}" --group "test-cu${TEST_CUDA_MAJOR}${FREE_THREADING}" "cuda-toolkit==${CUDA_VER_MINOR}.*"
106+
echo "Installed packages before core tests:"
107+
pip list
108+
echo "Running core tests"
109+
${SANITIZER_CMD} pytest -rxXs -v --durations=0 --randomly-dont-reorganize tests/
110+
# Currently our CI always installs the latest bindings (from either major version).
111+
# This is not compatible with the test requirements.
112+
if [[ "${SKIP_CYTHON_TEST}" == 0 ]]; then
113+
${SANITIZER_CMD} pytest -rxXs -v --durations=0 --randomly-dont-reorganize tests/cython
159114
fi
160-
PIP_ARGS+=(--extra-index-url "https://download.pytorch.org/whl/${TORCH_CUDA}")
161-
elif [[ "${test_module}" == "nightly-numba-cuda" ]]; then
162-
echo "Installing pathfinder + bindings + core + test deps + numba-cuda"
163-
# numba-cuda's test-cuXX group deps (can't use --group for a wheel install):
164-
PIP_ARGS+=(
165-
"cuda-toolkit[curand,cublas]==${CUDA_VER_MINOR}.*"
166-
"numba-cuda[cu${TEST_CUDA_MAJOR}]"
167-
"cupy-cuda${TEST_CUDA_MAJOR}x"
168-
psutil cffi pytest-xdist pytest-benchmark filecheck ml_dtypes statistics
115+
else
116+
# Nightly optional-dependency testing.
117+
# Install ALL wheels (pathfinder + bindings + core) and the optional dep
118+
# in a single pip call so pip resolves version constraints in one shot.
119+
PIP_ARGS=(
120+
"${PATHFINDER_WHL[@]}"
121+
"${BINDINGS_ARGS[@]}"
122+
"${CORE_WHL[@]}"
123+
--group "test-cu${TEST_CUDA_MAJOR}${FREE_THREADING}"
169124
)
125+
126+
if [[ "${test_module}" == "nightly-pytorch" ]]; then
127+
# TORCH_VER and TORCH_CUDA must be set by the caller.
128+
# Use cuda-toolkit[cudart] only — torch brings its own nvcc/nvrtc/etc.
129+
# This avoids version conflicts between our nvidia-* pins and torch's.
130+
echo "Installing pathfinder + bindings + core + test deps + PyTorch ${TORCH_VER} (${TORCH_CUDA})"
131+
PIP_ARGS+=("cuda-toolkit[cudart]==${CUDA_VER_MINOR}.*")
132+
if [[ "${TORCH_VER}" == "latest" ]]; then
133+
PIP_ARGS+=(torch)
134+
else
135+
PIP_ARGS+=("torch==${TORCH_VER}")
136+
fi
137+
PIP_ARGS+=(--extra-index-url "https://download.pytorch.org/whl/${TORCH_CUDA}")
138+
elif [[ "${test_module}" == "nightly-numba-cuda" ]]; then
139+
echo "Installing pathfinder + bindings + core + test deps + numba-cuda"
140+
# numba-cuda's test-cuXX group deps (can't use --group for a wheel install):
141+
PIP_ARGS+=(
142+
"cuda-toolkit[curand,cublas]==${CUDA_VER_MINOR}.*"
143+
"numba-cuda[cu${TEST_CUDA_MAJOR}]"
144+
"cupy-cuda${TEST_CUDA_MAJOR}x"
145+
psutil cffi pytest-xdist pytest-benchmark filecheck ml_dtypes statistics
146+
)
147+
fi
148+
149+
pip install "${PIP_ARGS[@]}"
150+
echo "Nightly install complete — installed packages:"
151+
pip list
170152
fi
171153

172-
pip install "${PIP_ARGS[@]}"
173154
popd
174-
175-
echo "Nightly install complete — installed packages:"
176-
pip list
177155
fi

0 commit comments

Comments
 (0)