Skip to content

Commit 183529a

Browse files
wilke0818wilke0818
authored
Merge pull request #516 from sensein/20260512-204619-fix-canary-cuda-conflict
subprocess_venv: route torch/torchaudio via CUDA-aware PyTorch index
2 parents 1443cdf + 2c0d351 commit 183529a

17 files changed

Lines changed: 2342 additions & 14 deletions

File tree

‎README.md‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,20 @@ For more detailed information, check out our [**Documentation**](https://sensein
7979
3. CUDA libraries matching the CUDA version expected by the PyTorch wheels (e.g., the latest pytorch 2.8 expects cuda-12.8). To install those with conda, please do:
8080
- ```conda config --add channels nvidia```
8181
- ```conda install -y nvidia/label/cuda-12.8.1::cuda-libraries-dev```
82+
83+
**Hosts with newer system CUDA** (e.g., CUDA 12.9): the subprocess-venv backends (`nemo-canary-qwen`, `nemo`, `qwen-asr`) auto-detect the host's CUDA version via `nvidia-smi` and route their `torch`/`torchaudio` installs through the matching PyTorch wheel index (`cu128` / `cu126` / `cu124` / `cu121` / `cpu`). No manual configuration needed.
84+
85+
**Internal mirrors / unsupported CUDA / CPU fallback**: set the `SENSELAB_TORCH_INDEX_URL` environment variable to override the chosen index. Common values:
86+
87+
```bash
88+
# Force CPU wheels (e.g. testing CPU path on a GPU host, or unsupported CUDA)
89+
export SENSELAB_TORCH_INDEX_URL=https://download.pytorch.org/whl/cpu
90+
91+
# Internal PyPI mirror that proxies PyTorch wheels
92+
export SENSELAB_TORCH_INDEX_URL=https://pypi.internal.example.com/pytorch/cu128
93+
```
94+
95+
When no compatible `torch`+`torchaudio` binary pair exists for your host (rare; happens in the days after a CUDA major release), installation fails with a named `SenselabCudaCompatibilityError` that lists the detected host CUDA, the attempted index URL, and the recommended action — no opaque stack traces from inside `torchaudio`.
8296
4. Docker is required and must be running for some video models (e.g., MediaPipe-based estimators).
8397
Please follow the official installation instructions for your platform: [Install Docker](https://docs.docker.com/get-started/get-docker/).
8498
5. Some functionalities rely on HuggingFace models, and increasingly, models require authentication and signed license agreements. Instructions on how to generate a Hugging Face access token can be found here: https://huggingface.co/docs/hub/security-tokens

‎specs/20260512-204619-fix-canary-cuda-conflict/checklists/requirements.md‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
# Specification Quality Checklist: Resolve NeMo Canary torch/torchaudio CUDA mismatch on newer-CUDA hosts
2+
3+
**Purpose**: Validate specification completeness and quality before proceeding to planning
4+
**Created**: 2026-05-12
5+
**Feature**: [spec.md](../spec.md)
6+
7+
## Content Quality
8+
9+
- [X] No implementation details (languages, frameworks, APIs)
10+
- [X] Focused on user value and business needs
11+
- [X] Written for non-technical stakeholders
12+
- [X] All mandatory sections completed
13+
14+
## Requirement Completeness
15+
16+
- [X] No [NEEDS CLARIFICATION] markers remain
17+
- [X] Requirements are testable and unambiguous
18+
- [X] Success criteria are measurable
19+
- [X] Success criteria are technology-agnostic (no implementation details)
20+
- [X] All acceptance scenarios are defined
21+
- [X] Edge cases are identified
22+
- [X] Scope is clearly bounded
23+
- [X] Dependencies and assumptions identified
24+
25+
## Feature Readiness
26+
27+
- [X] All functional requirements have clear acceptance criteria
28+
- [X] User scenarios cover primary flows
29+
- [X] Feature meets measurable outcomes defined in Success Criteria
30+
- [X] No implementation details leak into specification
31+
32+
## Notes
33+
34+
- Items marked incomplete require spec updates before `/speckit.clarify` or `/speckit.plan`.
35+
- Validation pass 1 (2026-05-12): all items pass. The spec avoids implementation specifics (no mention of `uv pip`, index URLs, specific version pins) and frames the problem in terms of "binary pair compatibility" and "isolated backend environment" so the planning phase has freedom to pick the right resolution mechanism.
36+
- Validation pass 2 (2026-05-12, post-`/speckit.analyze`): all items still pass. Tightened SC-001 ("under a slow link" → "≥ 100 Mbps baseline") and added a concrete example for "newer-CUDA host" (CUDA > 12.8). Both edits are quality polish, not structural changes — the spec remains implementation-agnostic.
37+
- One judgment call worth flagging for `/speckit.plan`: SC-005 requires at least one other isolated-environment backend to benefit from the same mechanism in the same release. If during planning it becomes clear that the Qwen ASR backend already works fine on CUDA 12.9 (e.g., it doesn't ship torchaudio), the success criterion can be downgraded to "documented as out-of-scope-this-release" per the parenthetical already in SC-005.

‎specs/20260512-204619-fix-canary-cuda-conflict/contracts/ensure_venv.md‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
# Contract: `ensure_venv` behavior change
2+
3+
## Signature (unchanged)
4+
5+
```python
6+
def ensure_venv(
7+
name: str,
8+
requirements: list[str],
9+
python_version: Optional[str] = None,
10+
) -> Path: ...
11+
```
12+
13+
## Behavior change
14+
15+
### Before this fix
16+
17+
Inside the lock:
18+
1. Compare requested `requirements` against the marker's stored list.
19+
2. If matches → return cached venv.
20+
3. Otherwise: `uv venv ...` → `uv pip install <requirements> safetensors numpy torchaudio`.
21+
4. Write marker, return venv dir.
22+
23+
### After this fix
24+
25+
Inside the lock:
26+
1. Resolve the torch index first:
27+
- `env_override = os.getenv("SENSELAB_TORCH_INDEX_URL") or None` (empty string treated as unset).
28+
- If `env_override` is set → skip `detect_host_cuda()` (it would be wasted work; the picker short-circuits on override).
29+
- Else → `host_cuda = cuda_probe.detect_host_cuda()`.
30+
- `torch_index = cuda_probe.pick_torch_index(host_cuda, env_override=env_override)`.
31+
2. Compare requested `requirements` AND `stored["torch_index"]["url"] == torch_index.url` against the marker's stored data.
32+
3. If matches → return cached venv.
33+
4. Otherwise:
34+
a. `uv venv ...` (recreate fresh).
35+
b. `uv pip install --index-url <torch_index.url> --extra-index-url https://pypi.org/simple --python <venv-python> <requirements> safetensors numpy torchaudio`.
36+
c. On `subprocess.CalledProcessError`, parse stderr for "no matching distribution" / "could not find a version" patterns. If matched, raise `SenselabCudaCompatibilityError(host_cuda=..., attempted_index=..., failing_packages=[...])`. Otherwise, re-raise the original.
37+
d. Write the marker including `torch_index`.
38+
5. Return venv dir.
39+
40+
## New failure mode
41+
42+
`SenselabCudaCompatibilityError` raised when:
43+
- The chosen torch index has no wheel matching the venv's python version + platform.
44+
- The override URL provided via `SENSELAB_TORCH_INDEX_URL` doesn't serve a compatible wheel.
45+
46+
Error message format (single line, actionable):
47+
48+
```text
49+
SenselabCudaCompatibilityError: No torch+torchaudio wheel available for this host.
50+
Host CUDA: 12.9 (source: nvidia-smi)
51+
Attempted index: https://download.pytorch.org/whl/cu128
52+
Failing packages: ['torch==2.8.1+cu128', 'torchaudio==2.8.1+cu128']
53+
Action: downgrade CUDA, set SENSELAB_TORCH_INDEX_URL=https://download.pytorch.org/whl/cpu for CPU fallback, or wait for upstream wheels.
54+
```
55+
56+
## Cleanup on failure
57+
58+
If `uv pip install` fails, the venv directory is removed before raising. This guarantees no partially-built venv survives, which is what FR-005 and SC-004 require.
59+
60+
## Lock semantics (unchanged)
61+
62+
Per-name file lock at `~/.cache/senselab/venvs/<name>.lock` with 600-second timeout. Multiple processes calling `ensure_venv` with the same `name` serialize through this lock. Behavior under concurrency is unchanged by this fix.
63+
64+
## Marker format
65+
66+
See `data-model.md` "Marker schema" section. Summary: adds a `torch_index` field. Markers written by the pre-fix code (no `torch_index` key) trigger a one-time rebuild on first run after upgrade — handled by the same mismatch-clause that handles every other change to the marker payload.
67+
68+
## Test surface
69+
70+
Unit tests required in `src/tests/utils/subprocess_venv_test.py`:
71+
72+
- Pre-fix-shape marker (no `torch_index` key) → mismatch → rebuild path.
73+
- Marker with matching `torch_index.url` → cache hit.
74+
- Marker with non-matching `torch_index.url` → mismatch → rebuild.
75+
- `uv pip install` failure with "no matching distribution" stderr → `SenselabCudaCompatibilityError` raised, venv dir removed.
76+
- `uv pip install` failure with unrelated error → original `CalledProcessError` re-raised, venv dir removed.
77+
- Env override set → marker records `torch_index.source == "env-override"`, install argv uses override URL.

‎specs/20260512-204619-fix-canary-cuda-conflict/contracts/env-vars.md‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
# Contract: `SENSELAB_TORCH_INDEX_URL`
2+
3+
## Purpose
4+
5+
Operator escape hatch for hosts where the static CUDA → PyTorch wheel index map is wrong (rare CUDA versions, internal PyPI mirrors, air-gapped builds, ARM CUDA hosts, ROCm).
6+
7+
## Semantics
8+
9+
| State | Meaning |
10+
|---|---|
11+
| Unset (or empty string) | `cuda_probe.pick_torch_index()` consults the static map, picks the highest `cuXX ≤ host CUDA`, falls back to `cpu`. |
12+
| Set to a non-empty URL | That URL is used verbatim as the `--index-url` argument to `uv pip install`. The static map is bypassed for the index choice; the host CUDA probe still runs and is surfaced in any `SenselabCudaCompatibilityError` so the user can see whether the override was a sensible match for the actual host. |
13+
14+
The override URL is passed to `uv pip install` exactly as-is. Senselab does not validate the URL is reachable or that it serves PyTorch wheels — if the operator misconfigures it, the install will fail and surface a `SenselabCudaCompatibilityError` with the override URL named in the diagnostic.
15+
16+
## Examples
17+
18+
```bash
19+
# Force CPU wheels (e.g. a CUDA host where you want to verify CPU behavior)
20+
SENSELAB_TORCH_INDEX_URL=https://download.pytorch.org/whl/cpu uv run python scripts/analyze_audio.py ...
21+
22+
# Internal PyPI mirror that proxies PyTorch wheels
23+
SENSELAB_TORCH_INDEX_URL=https://pypi.internal.example.com/pytorch/cu128 uv run python ...
24+
25+
# ARM CUDA host (NVIDIA's PyPI mirror)
26+
SENSELAB_TORCH_INDEX_URL=https://pypi.nvidia.com uv run python ...
27+
28+
# Force a specific PyTorch CUDA version even when host CUDA is higher
29+
SENSELAB_TORCH_INDEX_URL=https://download.pytorch.org/whl/cu124 uv run python ...
30+
```
31+
32+
## Interaction with marker file
33+
34+
The marker file records the resolved index URL (whether from static map or override). If the user changes `SENSELAB_TORCH_INDEX_URL` between runs, the marker mismatch triggers a venv rebuild — the new URL is honored on the next install.
35+
36+
## Interaction with the rest of `uv pip install`
37+
38+
When the override is in effect, the install argv is:
39+
40+
```text
41+
uv pip install --index-url $SENSELAB_TORCH_INDEX_URL --extra-index-url https://pypi.org/simple --python <venv-python> <requirements> safetensors numpy torchaudio
42+
```
43+
44+
The `--extra-index-url` for PyPI is still present so non-torch packages (nemo_toolkit, qwen-asr, soundfile, etc.) continue to resolve from PyPI.
45+
46+
## Persistence
47+
48+
Senselab does not write the override anywhere outside the venv marker. Restart of the shell session resets the variable to unset. Users wanting a persistent override should add it to their shell profile or to the env file their workflow loads.
49+
50+
## Documentation pointers
51+
52+
- Listed in repo `CLAUDE.md` env-var index.
53+
- Listed in repo `README.md` "Troubleshooting / CUDA hosts" section (to be added in the docs task).
54+
- Mentioned in the `SenselabCudaCompatibilityError` message as a recommended user action.

‎specs/20260512-204619-fix-canary-cuda-conflict/data-model.md‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
# Phase 1 — Data Model
2+
3+
## Value objects
4+
5+
### `HostCuda`
6+
7+
Result of probing the host's runtime CUDA capability. Produced by `cuda_probe.detect_host_cuda()`.
8+
9+
| Field | Type | Description |
10+
|---|---|---|
11+
| `version` | `tuple[int, int] \| None` | Major + minor, e.g. `(12, 9)`. `None` when no CUDA detected. |
12+
| `source` | `Literal["nvidia-smi", "nvcc", "none"]` | Which probe succeeded. `"none"` only when both failed. |
13+
| `raw` | `str` | Raw stdout the probe parsed (or `""` for `"none"`). Kept for the diagnostic in FR-004. |
14+
15+
**Validation rules**:
16+
- `version is None` ⟺ `source == "none"`.
17+
- `version` has both elements `≥ 0` when present.
18+
- Immutable; produced fresh on every probe call (no class-level cache; caching is the caller's responsibility — `ensure_venv` records it in the marker).
19+
20+
### `TorchIndex`
21+
22+
The chosen index URL for the subsequent `uv pip install`. Produced by `cuda_probe.pick_torch_index(host_cuda, env_override=None)`.
23+
24+
| Field | Type | Description |
25+
|---|---|---|
26+
| `url` | `str` | Full HTTPS URL of the chosen index. |
27+
| `tag` | `str` | Short identifier — `"cu128"`, `"cu126"`, `"cu124"`, `"cu121"`, `"cpu"`, or `"override"` when the env var supplied an arbitrary URL. |
28+
| `cuda_version` | `tuple[int, int] \| None` | The index's CUDA version, not the host's. `None` when tag is `"cpu"` or `"override"`. |
29+
| `source` | `Literal["static-map", "env-override"]` | How this index was selected. |
30+
31+
**Validation rules**:
32+
- `tag == "override"` ⟺ `source == "env-override"`.
33+
- `tag == "cpu"` ⟺ `cuda_version is None and source == "static-map"`.
34+
- `url` starts with `https://`.
35+
36+
## Marker schema (`.senselab-installed`)
37+
38+
Stored at `~/.cache/senselab/venvs/<name>/.senselab-installed` after a successful venv build.
39+
40+
**Current (pre-fix)**:
41+
```json
42+
{
43+
"requirements": ["nemo_toolkit[asr,tts] @ git+...", "torch>=2.8,<2.9", ...],
44+
"python_version": "3.12"
45+
}
46+
```
47+
48+
**After this fix**:
49+
```json
50+
{
51+
"requirements": ["nemo_toolkit[asr,tts] @ git+...", "torch>=2.8,<2.9", ...],
52+
"python_version": "3.12",
53+
"torch_index": {
54+
"tag": "cu128",
55+
"url": "https://download.pytorch.org/whl/cu128",
56+
"source": "static-map"
57+
}
58+
}
59+
```
60+
61+
**Mismatch rules** (any one triggers a rebuild):
62+
- `stored.get("requirements") != sorted(requirements)` — existing behavior, unchanged.
63+
- `stored.get("torch_index", {}).get("url")` differs from the currently-resolved `TorchIndex.url` — new clause. Absence of `torch_index` in the existing marker (old shape) counts as a mismatch, which is the intended one-time rebuild for users upgrading through this fix.
64+
65+
No `schema_version` field — the marker is internal state, not a published contract. Adding fields by field-presence is sufficient; explicit versioning would be ceremony with no payoff.
66+
67+
## State transitions
68+
69+
`ensure_venv` state diagram:
70+
71+
```text
72+
┌─────────────────┐
73+
call ─────► │ Acquire lock │
74+
└────────┬────────┘
75+
▼
76+
┌─────────────────────────┐
77+
│ Read marker (if exists) │
78+
└────────┬────────────────┘
79+
▼
80+
┌───────── match? ──────────┐
81+
│ │
82+
match │ │ mismatch / missing
83+
▼ ▼
84+
Return venv_dir Probe host CUDA
85+
Pick torch index
86+
Apply env override
87+
│
88+
▼
89+
uv venv (recreate)
90+
uv pip install
91+
with --index-url
92+
│
93+
┌───────────┴───────────┐
94+
│ │
95+
success failure
96+
│ │
97+
▼ ▼
98+
Write marker w/ torch_index Wrap into
99+
Return venv_dir SenselabCudaCompatibilityError
100+
re-raise
101+
```
102+
103+
## Public API surface added
104+
105+
```python
106+
# src/senselab/utils/cuda_probe.py
107+
108+
@dataclass(frozen=True)
109+
class HostCuda:
110+
version: tuple[int, int] | None
111+
source: Literal["nvidia-smi", "nvcc", "none"]
112+
raw: str
113+
114+
@dataclass(frozen=True)
115+
class TorchIndex:
116+
url: str
117+
tag: str
118+
cuda_version: tuple[int, int] | None
119+
source: Literal["static-map", "env-override"]
120+
121+
class SenselabCudaCompatibilityError(RuntimeError):
122+
"""Raised when no torch/torchaudio binary pair is installable on this host."""
123+
124+
def detect_host_cuda() -> HostCuda: ...
125+
126+
def pick_torch_index(
127+
host_cuda: HostCuda,
128+
env_override: str | None = None,
129+
) -> TorchIndex: ...
130+
```
131+
132+
The two existing public functions in `subprocess_venv.py` (`ensure_venv`, `venv_python`) keep their signatures. The behavior of `ensure_venv` changes: it now probes CUDA and routes the install through the chosen index; on unsupported-CUDA hosts it raises `SenselabCudaCompatibilityError` instead of letting `uv pip install` fail with an opaque error.
133+
134+
## Backwards compatibility
135+
136+
- Existing marker files without a `torch_index` key: trigger one rebuild on first run after upgrade (intentional, see Mismatch rules).
137+
- Existing callers of `ensure_venv`: signature unchanged; no migration needed.
138+
- Existing subprocess backends (`canary_qwen.py`, `nemo.py`, `qwen.py`): no code change required. They benefit automatically.
139+
- Users without GPUs: behavior unchanged today — install used `cpu` wheels via PyPI's default torch resolution, post-fix they use `cpu` wheels via the explicit CPU index. Same wheel, different URL.

0 commit comments

Comments
 (0)