subprocess_venv: route torch/torchaudio via CUDA-aware PyTorch index#516
Conversation
New module ``senselab.utils.cuda_probe`` provides: - ``detect_host_cuda()`` — parses ``CUDA Version: X.Y`` from ``nvidia-smi``'s default header, falls back to ``release X.Y`` from ``nvcc --version``, then ``HostCuda(version=None)`` when neither is available. Each probe has a 5s subprocess timeout. No driver→CUDA lookup table to maintain — both probes report the CUDA version directly. - ``pick_torch_index(host_cuda, env_override=None)`` — picks the highest static-map entry whose CUDA is ``<=`` host's, from ``[cu128, cu126, cu124, cu121]``; falls back to ``cpu`` index. When ``env_override`` is set + non-empty, returns that URL verbatim with ``tag="override"``. - ``SenselabCudaCompatibilityError`` — wraps wheel-not-found install failures with host CUDA + attempted index + failing packages and a one-line recommended action (downgrade CUDA / set ``SENSELAB_TORCH_INDEX_URL=https://download.pytorch.org/whl/cpu`` / wait for upstream wheels). Two frozen dataclasses (``HostCuda``, ``TorchIndex``) carry the probe and picker results so they're individually testable and the resolved index URL is auditable in the subprocess-venv marker file. Tests: ``src/tests/utils/cuda_probe_test.py`` — 15 unit tests, all ``subprocess.run`` calls mocked, no shell-out on CI. Coverage: both probe paths, both probe failures, all 5 index map entries, env override short-circuit + empty-string-treated-as-unset, error message contents. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
``ensure_venv`` now resolves a PyTorch wheel index for the host before
running ``uv pip install`` and routes the install through that index. On
hosts with system CUDA newer than the PyTorch default-wheels CUDA (e.g.
CUDA 12.9 against PyTorch's ``cu128`` default), this prevents ``torch``
and ``torchaudio`` from resolving to mismatched CUDA toolchains and
breaking their ABI contract at import time.
Resolution order inside the lock:
1. ``env_override = os.getenv("SENSELAB_TORCH_INDEX_URL") or None`` (empty
string treated as unset).
2. If ``env_override`` is set: skip the host CUDA probe — the override
wins regardless of host CUDA.
3. Otherwise: ``host_cuda = detect_host_cuda()`` followed by
``pick_torch_index(host_cuda, env_override=env_override)``.
The marker file gains a ``torch_index`` field; absence of that field in
an existing marker counts as a mismatch and triggers a rebuild — so
users upgrading through this fix get one automatic rebuild on first run.
No explicit schema-version field; the marker is internal cache state,
not a published contract.
Install argv now prefixes ``--index-url <torch_index.url>
--extra-index-url https://pypi.org/simple ...`` — the PyTorch index for
torch/torchaudio resolution, PyPI as the extra index so non-torch
packages (nemo_toolkit, qwen-asr, etc.) still resolve.
Failure handling: ``uv pip install`` failures pass through
``_classify_uv_failure`` to distinguish "no matching distribution" /
"could not find a version" errors (wrap into
``SenselabCudaCompatibilityError`` with named host/index/packages/action)
from unrelated failures (re-raise the original ``CalledProcessError``).
Either way the half-built venv directory is wiped before the raise so
the next run starts clean — addresses the recovery-from-broken-state
case.
Added in-source warnings to ``canary_qwen.py``, ``nemo.py``, and
``qwen.py`` near each ``*_REQUIREMENTS`` list: developers editing the
torch pins should not bypass ``ensure_venv`` with a per-backend install
shellout, since the CUDA routing is what guarantees the toolchain match.
Tests: ``src/tests/utils/subprocess_venv_test.py`` — 14 integration
tests, ``subprocess.run`` mocked. Coverage: marker-mismatch rebuild
paths (pre-fix-shape marker, matching index, different index), install
argv routing, env override short-circuit + empty-string-unset, wheel
not-found wrapping, unrelated-error pass-through, ``_classify_uv_failure``
unit cases, parameterized regression test walking all three real backend
requirement lists.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
README "System Requirements" section now mentions: - Subprocess-venv backends auto-route their ``torch``/``torchaudio`` install through the matching PyTorch wheel index based on host CUDA; no manual configuration needed. - ``SENSELAB_TORCH_INDEX_URL`` operator override for internal mirrors, unsupported CUDA versions, or forced CPU fallback. - Unsupported-CUDA failures surface as ``SenselabCudaCompatibilityError`` with named host CUDA, attempted index, and recommended action. Spec docs land under ``specs/20260512-204619-fix-canary-cuda-conflict/``: spec.md (3 prioritized user stories), plan.md (Constitution check + 6 research decisions), research.md, data-model.md, contracts/, quickstart.md (validation recipes for happy path, CPU fallback, negative path, regression, recovery), tasks.md (28 tasks across 6 phases), and checklists/requirements.md (passed twice). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a mechanism to resolve CUDA toolchain mismatches between torch and torchaudio in isolated subprocess virtual environments, specifically targeting hosts with newer system CUDA versions. A new cuda_probe utility detects the host's CUDA version and selects the most compatible official PyTorch wheel index, which is then used during the uv pip install process in ensure_venv. The implementation includes an environment variable override (SENSELAB_TORCH_INDEX_URL), automated venv rebuilds upon index changes, and structured error handling via SenselabCudaCompatibilityError. Feedback suggests improving the robustness of the uv failure classification by logging the full stderr to ensure diagnostic information is preserved if the tool's output format changes.
| matches = _FAILING_REQ_RE.findall(stderr) | ||
| if matches: | ||
| # De-duplicate preserving order. | ||
| seen: set[str] = set() | ||
| out: list[str] = [] | ||
| for m in matches: | ||
| if m not in seen: | ||
| seen.add(m) | ||
| out.append(m) | ||
| return out |
There was a problem hiding this comment.
The regex-based classification of uv failures relies on the specific error message format (backticked package names). While this works for current versions of uv, it is a best-effort diagnostic. If uv changes its output format, this helper might return ["<unknown>"]. Consider logging the full stderr in the SenselabCudaCompatibilityError to ensure the raw information is always available to the user.
Local-reviewer findings on the cuda_probe + subprocess_venv commits:
- MEDIUM (subprocess_venv): on the env-override path, host_cuda was
hardcoded to ``HostCuda(version=None)`` so the wrapped
``SenselabCudaCompatibilityError`` reported "Host CUDA: none" even on
real CUDA hosts using internal mirrors. Run ``detect_host_cuda()``
unconditionally now; the override still bypasses the static index map
but the probe result is preserved for diagnostic surface.
- MEDIUM (subprocess_venv): the failing-package extraction regex was
pulling EVERY backticked token from uv's stderr, so unrelated hints
like ``\`uv cache clean\`\`\` could leak into the
``SenselabCudaCompatibilityError`` message as fake "failing packages".
Anchored the capture to the "no matching distribution" / "could not
find a version" phrases so only the actual failing requirement is
surfaced.
- LOW (subprocess_venv): ``uv venv`` failure (e.g. requested Python
version unavailable) re-raised without wiping a partially-populated
venv directory. Mirror the install-failure cleanup so both failure
paths leave a clean cache state.
- LOW (cuda_probe): ``_PYTORCH_INDEX_MAP`` ordering ("highest CUDA
first") was asserted by comment only. Wrapped the literal in
``sorted(..., reverse=True)`` so a future contributor appending an
entry in the wrong position can't silently break routing.
- LOW (subprocess_venv): on the wrap-into-compat-error path, the
pre-classification ``logger.error`` was duplicating the wrapped
exception's already-formatted message and dumping unbounded uv
stderr. Demoted to ``logger.debug``; the pass-through path still logs
at error level since that's the only signal the user gets for
non-wheel failures.
Test updates: ``test_env_override_short_circuits_probe_*`` renamed and
reworked to assert the probe DOES run (not skips) on the override path
— the contract is the URL routing, not the probe-skip side effect.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The pre-fix-time local pre-commit pass was reading a stale mypy cache. Once invalidated (CI / fresh clone), five errors surfaced: - src/tests/utils/cuda_probe_test.py:44,46,67,69 — `_runner(*args: object)` was too generic for the indexed access `args[0]` (mypy "Value of type 'object' is not indexable"). Tightened to the actual ``subprocess.run`` positional signature `_runner(args: list[str], **kwargs: object)` — matches reality and unblocks indexing. - src/tests/utils/subprocess_venv_test.py:80 — `_SubprocessRecorder.hook` was typed `Optional[object]`, which mypy refuses to call. Tightened to `Optional[Callable[[list[list[str]]], subprocess.CompletedProcess]]` — the type the hook actually implements. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Two occurrences in research.md flagged by codespell on CI. Same fix shape — the hyphenated forms are non-standard per the project's codespell allowlist. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
Claude comment on comparison to #517: One thing worth flagging in review (if you want to): the static map currently tops out at cu128 and the picker falls back to "highest ≤ host" — that means a cu13.x Otherwise this seems strictly better than #517 minus one small change that I will PR separately |
|
The bug isn't just probe-fall-through — it's that even with the right cu128 index explicitly named, uv puts --extra-index-url pypi in higher precedence and PyPI's wheels win PR #516 —
|
|
send a PR (the workaround seems reasonable). |
…plit PR #516 routes the install through ``--index-url <cuda> --extra-index-url https://pypi.org/simple`` so non-torch packages can still resolve from PyPI. But uv treats ``--extra-index-url`` as having higher priority than ``--index-url`` (opposite of pip), so PyPI wins for every package — torch and torchaudio included. On hosts where PyPI ships those with mismatched ``+cu`` local-version tags (currently ``torch==X+cu129`` vs ``torchaudio==X`` with no tag, internally cu128), the resulting venv reproduces the same ABI mismatch this PR is meant to fix: RuntimeError: Detected that PyTorch and TorchAudio were compiled with different CUDA versions. PyTorch has CUDA version 12.9 whereas TorchAudio has CUDA version 12.8. Reproduced on MIT ORCD with the env override set (``SENSELAB_TORCH_INDEX_URL=https://download.pytorch.org/whl/cu128``) — override URL was honored in the marker but uv still resolved both wheels from PyPI. So the override path was non-functional too. Fix: split the install into two stages. Stage 1 — ``uv pip install --index-url <chosen> torch torchaudio`` with NO ``--extra-index-url``. The chosen CUDA index is unambiguously primary; both wheels and their ``nvidia-cuda-runtime-cu12`` transitives come from the same toolchain. Pins from the caller's ``requirements`` list flow through verbatim; multiple constraints for the same package (``["torch>=2.8", "torch<2.9"]``) are all forwarded so uv can combine them at resolve time. Stage 2 — ``uv pip install <rest of requirements> safetensors numpy`` with no index flags. Backend-pinned torch / torchaudio specs are filtered out so uv can't be tempted to re-resolve them from PyPI against the matched ``+cu128`` wheels installed in Stage 1. uv sees both already installed and satisfying constraints from their transitives' point of view, and leaves them alone. The PyTorch index now governs only the two packages it's designed for; stale wheels on the CUDA index for utilities like setuptools / pyarrow stay out of the picture. Add ``requires_torch=True`` parameter to ``ensure_venv`` for opt-out: subprocess venvs that genuinely don't touch torch (e.g. consuming a pure-Python GitHub repo) can pass ``requires_torch=False`` to skip the probe, skip Stage 1, and skip the IPC ``torchaudio`` append entirely — single install pass against default PyPI. Marker omits ``torch_index`` in that case; a later ``requires_torch=True`` call against the same name correctly invalidates and rebuilds via the existing URL-mismatch clause. Behavioral guarantees preserved from PR #516: marker schema (extended, not changed), cache-hit fast path, env override (``SENSELAB_TORCH_INDEX_URL``), ``SenselabCudaCompatibilityError`` wrapping on wheel-not-found errors (now scoped to Stage 1 where it semantically belongs), pass-through of unrelated ``CalledProcessError``, half-built-venv cleanup on failure, host-CUDA probe diagnostic, cross-backend routing. Tests: existing call-count assertions updated for the extra ``uv pip install``; the two install-argv tests split into seven new ones covering (1) Stage 1 names only ``--index-url`` with torch + torchaudio pinned per requirements, (2) Stage 1 falls back to bare names when requirements omit them (qwen backend), (3) Stage 2 carries no index flags + IPC deps, (4) Stage 2 filters torch / torchaudio specs from caller's requirements, (5) multiple constraints for the same package all forward through Stage 1 verbatim, (6) ``requires_torch=False`` skips the probe + Stage 1 + IPC torchaudio entirely, and (7) switching from ``requires_torch=False`` to ``True`` on the same name invalidates the cache and rebuilds. The cross-backend regression test now asserts both that Stage 1 routes through the chosen index AND that Stage 2 contains no torch / torchaudio specs at all — guarding against any future re-introduction of ``--extra-index-url`` on the torch install or a torch spec leaking into Stage 2. Review feedback addressed: - @satra: ``requires_torch`` opt-out for venvs that don't need torch - gemini-code-assist #1: filter torch / torchaudio specs from Stage 2 - gemini-code-assist #2: list-based ``_torch_install_specs`` keeps every matching constraint instead of clobbering through a dict Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…equirements PR #516 routes the install through ``--index-url <cuda> --extra-index-url https://pypi.org/simple`` so non-torch packages can still resolve from PyPI. But uv treats ``--extra-index-url`` as having higher priority than ``--index-url`` (opposite of pip), so PyPI wins for every package — torch and torchaudio included. On hosts where PyPI ships those with mismatched ``+cu`` local-version tags (currently ``torch==X+cu129`` vs ``torchaudio==X`` with no tag, internally cu128), the resulting venv reproduces the same ABI mismatch this PR is meant to fix: RuntimeError: Detected that PyTorch and TorchAudio were compiled with different CUDA versions. PyTorch has CUDA version 12.9 whereas TorchAudio has CUDA version 12.8. Reproduced on MIT ORCD with the env override set (``SENSELAB_TORCH_INDEX_URL=https://download.pytorch.org/whl/cu128``) — override URL was honored in the marker but uv still resolved both wheels from PyPI. So the override path was non-functional too. Fix: split the install into two stages and decide whether Stage 1 fires from the caller's ``requirements`` itself. Stage 1 — ``uv pip install --index-url <chosen>`` listing every ``torch`` / ``torchaudio`` spec found in ``requirements``, with NO ``--extra-index-url``. The chosen CUDA index is unambiguously primary; both wheels and their ``nvidia-cuda-runtime-cu12`` transitives come from the same toolchain. Multiple constraints for the same package (``["torch>=2.8", "torch<2.9"]``) are all forwarded so uv combines them at resolve time. Stage 2 — ``uv pip install <rest of requirements> safetensors numpy`` with no index flags. Backend-pinned torch / torchaudio specs are filtered out so uv can't be tempted to re-resolve them from PyPI against the matched ``+cu128`` wheels installed in Stage 1. Drop the unconditional ``torchaudio`` IPC append, in two ways: 1. ``_torch_install_specs`` no longer pads its return with bare ``torch`` / ``torchaudio`` names. If neither is in ``requirements``, it returns ``[]``. 2. ``ensure_venv`` treats ``_torch_install_specs(requirements) == []`` as the trigger to skip the probe entirely and run a single torch-free install pass against default PyPI. No ``nvidia-smi`` shellout, no ``torchaudio`` force-appended. This addresses @satra's review concern that the previous default forced ``torchaudio`` into every subprocess venv, even backends that don't use it. With the change, ``yamnet`` and ``continuous-ser`` — both of which read audio via ``soundfile`` in their worker scripts and never import torchaudio — get ~200 MB leaner venvs. Backends that genuinely need ``torch`` / ``torchaudio`` (including via a transitive dep) MUST pin them explicitly so Stage 1 routes them through the matched CUDA index; this is also why ``qwen.py``'s ``_REQUIREMENTS`` now names them directly even though ``qwen-asr==0.0.6`` would otherwise pull them transitively — relying on the transitive would let Stage 2's PyPI resolution split the local-version tags. The 7 backends that already pin ``torchaudio`` explicitly (canary, nemo, ppgs, sparc, coqui, s3prl) are unaffected. The earlier ``requires_torch=False`` parameter from a previous round of this review (an explicit opt-out flag) is removed in favor of the auto-detection — it's structurally equivalent and the explicit flag becomes redundant noise. Behavioral guarantees preserved from PR #516: marker schema (extended, not changed), cache-hit fast path, env override (``SENSELAB_TORCH_INDEX_URL``), ``SenselabCudaCompatibilityError`` wrapping on wheel-not-found errors (now scoped to Stage 1 where it semantically belongs), pass-through of unrelated ``CalledProcessError``, half-built-venv cleanup on failure, host-CUDA probe diagnostic surface, cross-backend routing. Tests: 34 in total covering (1) Stage 1 names only ``--index-url`` with torch + torchaudio pinned per requirements; (2) Stage 2 carries no index flags + IPC deps; (3) Stage 2 filters torch / torchaudio specs from caller's requirements; (4) multiple constraints for the same package all forward through Stage 1 verbatim; (5) requirements without torch / torchaudio skip the probe and Stage 1 entirely; (6) yamnet- style requirements get NO torchaudio in their install argv; (7) switching a venv name from torch-free to torch-using correctly invalidates the cache and rebuilds. The cross-backend regression test asserts both that Stage 1 routes through the chosen index AND that Stage 2 contains no torch / torchaudio specs at all. Review feedback addressed: - @satra (PR #518 review): drop the unconditional ``torchaudio`` IPC append; auto-detect torch routing from the caller's ``requirements`` instead of forcing every venv to pay for it. Adds explicit ``torch`` + ``torchaudio`` to qwen's ``_REQUIREMENTS`` so its transitive-dep resolution still flows through the CUDA index. - gemini-code-assist #1: filter torch / torchaudio specs from Stage 2. - gemini-code-assist #2: list-based ``_torch_install_specs`` keeps every matching constraint instead of clobbering through a dict. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
subprocess_venv: two-stage install to fix torch/torchaudio CUDA-tag s…
2 participants
Summary
Subprocess-venv backends (
nemo-canary-qwen,nemo,qwen-asr) were resolvingtorchandtorchaudiothrough PyPI's default index. On hosts with system CUDA newer than the PyTorch default-wheels CUDA (e.g. CUDA 12.9 againstcu128), the two libraries could end up compiled for different CUDA toolchains and break their ABI at import time.The shared
ensure_venvhelper insrc/senselab/utils/subprocess_venv.pynow:nvidia-smi(default headerCUDA Version: X.Y) →nvcc --version(release X.Y) → none.cu128/cu126/cu124/cu121/cpu), preferring the highest entry<=host CUDA.uv pip installthrough--index-url <chosen> --extra-index-url https://pypi.org/simpleso torch/torchaudio go through the PyTorch index while non-torch packages still resolve from PyPI..senselab-installedgains atorch_indexfield). Markers from the pre-fix code (notorch_indexkey) auto-rebuild on first run after upgrade.SenselabCudaCompatibilityErrorcarrying host CUDA + attempted index + failing packages + recommended action.Operator override: set
SENSELAB_TORCH_INDEX_URLto use a custom URL verbatim (internal mirrors, forced CPU fallback, unusual hosts). Empty string is treated as unset.Reusable across the three current subprocess-venv backends and any future one with no per-backend code change. In-source warnings landed near each backend's
*_REQUIREMENTSlist to prevent future bypass.Test plan
uv run pytest src/tests/utils/cuda_probe_test.py— 15 unit tests covering both probe paths, all index-map entries, env override behavior, error message contents.uv run pytest src/tests/utils/subprocess_venv_test.py— 14 integration tests covering marker-mismatch rebuild paths, install argv routing, env override short-circuit,SenselabCudaCompatibilityErrorwrapping, unrelated-error pass-through,_classify_uv_failureunit cases, parameterized regression test across all three real backend requirement lists.uv run pre-commit run --all-files— ruff + mypy + codespell + TOML/YAML/JSON formatters all clean.specs/20260512-204619-fix-canary-cuda-conflict/quickstart.md. Project CI has no CUDA 12.9 runner — this validation is a release-blocker checklist item, not a per-PR gate.Spec docs
Full spec, plan, research, contracts, data-model, and quickstart under
specs/20260512-204619-fix-canary-cuda-conflict/.🤖 Generated with Claude Code
Version
Published prerelease version:
1.3.1-alpha.31Changelog
🐛 Bug Fix
Authors: 3