diff --git a/.github/actions/setup-gpu-test-env/action.yml b/.github/actions/setup-gpu-test-env/action.yml index 1becddc7a..818704890 100644 --- a/.github/actions/setup-gpu-test-env/action.yml +++ b/.github/actions/setup-gpu-test-env/action.yml @@ -81,15 +81,31 @@ runs: import tomllib with open(sys.argv[1], "rb") as handle: - indexes = tomllib.load(handle)["tool"]["uv"]["index"] + uv_config = tomllib.load(handle)["tool"]["uv"] + indexes = uv_config["index"] cuda_extra = sys.argv[2] + + # Some indexes are reached only through [tool.uv.sources] and carry no + # CUDA variant in their name or URL (flashinfer-cubin serves a single + # https://flashinfer.ai/whl/ for every extra). We pass --no-sources + # below, so collect those by their declared extra or they are dropped. + source_indexes = { + entry["index"] + for value in uv_config.get("sources", {}).values() + for entry in (value if isinstance(value, list) else [value]) + if isinstance(entry, dict) + and entry.get("extra") == cuda_extra + and "index" in entry + } + print( "\n".join( index["url"] for index in indexes if index["name"].endswith(f"-{cuda_extra}") or f"/{cuda_extra}" in index["url"] + or index["name"] in source_indexes ) ) PY