ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix ROCm wheel-index unit test: extract the gfx-arch probe helpers get_torch_index_url now calls (#7399)

Browse source 
* Fix ROCm wheel-index test: extract the gfx-arch probe helpers get_torch_index_url now calls

get_torch_index_url gained a gfx-arch probe on the ROCm path (Strix reroute
work) and now calls _ensure_rocm_probe_env, _probe_amd_gfx_arch,
_infer_linux_amd_gfx_arch and friends. The unit test in
tests/sh/test_get_torch_index_url.sh sources a curated subset of install.sh
functions, and that list was never updated, so those helpers were undefined
in the harness. On the ROCm path the gfx probe hit an undefined function,
the branch silently fell through to the CPU wheel index, and every ROCm
assertion failed (9 failures: all ROCm versions resolved to /whl/cpu).

Extract the six missing helpers so the ROCm branch runs end to end. All 49
assertions pass. Adds a comment noting these must stay in sync with
install.sh.

* Keep the ROCm wheel-index test hermetic: redirect the /opt/rocm prefix

Extracting _ensure_rocm_probe_env pulled its absolute-path host probe into the
harness: it appends /opt/rocm/bin to PATH and runs the real host rocminfo, and
version detection reads /opt/rocm/.info/version. On a host with ROCm installed
that leaks the host GPU into the minimal-PATH test, so the no-GPU and
CUDA-visible-device assertions could select a host ROCm wheel index instead of
their expected CPU result, making the test host-dependent.

Redirect the whole /opt/rocm prefix to an empty temp dir in the same sed pass
that stubs /usr/bin/nvidia-smi, so the probes stay hermetic. All 49 assertions
pass and the generated harness contains no real /opt/rocm path.

---------

Co-authored-by: danielhanchen <unslothai@gmail.com>
This commit is contained in:
Daniel Han 2026-07-24 02:12:16 -07:00 • committed by GitHub
commit 418ae14388
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 27 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

28
tests/sh/test_get_torch_index_url.sh
View file

@ -14,6 +14,13 @@ FAIL=0
# controllable path so we can test the "no GPU" scenario on GPU machines.
_FUNC_FILE=$(mktemp)
_FAKE_SMI_DIR=$(mktemp -d)
# The ROCm probe helpers read the real /opt/rocm prefix by ABSOLUTE path
# (_ensure_rocm_probe_env appends /opt/rocm/bin to PATH and runs the host
# rocminfo; version detection reads /opt/rocm/.info/version). On a real ROCm
# host that leaks the host GPU into the minimal-PATH harness and makes the
# no-GPU / CPU assertions host-dependent. Redirect the whole prefix to an empty
# temp dir so the probes stay hermetic (same idea as the nvidia-smi rewrite).
_FAKE_ROCM_DIR=$(mktemp -d)
{
sed -n '/^_run_bounded()/,/^}/p' "$INSTALL_SH"
echo ""
@ -23,10 +30,28 @@ _FAKE_SMI_DIR=$(mktemp -d)
echo ""
sed -n '/^_has_usable_nvidia_gpu()/,/^}/p' "$INSTALL_SH"
echo ""
# ROCm gfx-arch probe helpers that get_torch_index_url / _has_amd_rocm_gpu
# now call. These MUST stay in sync with install.sh: if get_torch_index_url
# references a helper that is not extracted here, the ROCm branch hits an
# undefined function, silently falls through to the CPU wheel index, and the
# ROCm assertions below fail.
sed -n '/^_ensure_rocm_probe_env()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_probe_amd_gfx_arch()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_amd_gpu_present_via_pci()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_infer_amd_gfx_arch_from_gpu_name()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_infer_linux_amd_gfx_arch()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_amd_arch_index_family_for_gfx()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_trim_index_path_slashes()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^get_torch_index_url()/,/^}/p' "$INSTALL_SH"
} | sed "s|/usr/bin/nvidia-smi|$_FAKE_SMI_DIR/nvidia-smi-absent|g" \
} | sed -e "s|/usr/bin/nvidia-smi|$_FAKE_SMI_DIR/nvidia-smi-absent|g" \
-e "s|/opt/rocm|$_FAKE_ROCM_DIR|g" \
> "$_FUNC_FILE"
# Save system PATH so we always have basic tools (uname, grep, head, etc.)
@ -438,6 +463,7 @@ assert_eq "url override preserves fragment slash" "https://mirror.example.com/wh
rm -f "$_FUNC_FILE"
rm -rf "$_FAKE_SMI_DIR"
rm -rf "$_FAKE_ROCM_DIR"
rm -rf "$_TOOLS_DIR"
echo ""