From 6b0449c326459ac188a97aeb1b83d66a3eb088be Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Mon, 27 Jul 2026 16:06:48 -0600 Subject: [PATCH] Fixed install issues with windows builds --- manager/README.md | 14 +++++- manager/detect.py | 8 ++- manager/doctor.py | 33 ++++++++++++- manager/env.py | 111 +++++++++++++++++++++++++++++++++++++++--- manager/spec.py | 37 ++++++++++++++ requirements_base.txt | 2 +- 6 files changed, 191 insertions(+), 14 deletions(-) diff --git a/manager/README.md b/manager/README.md index 75112e5..a103a7b 100644 --- a/manager/README.md +++ b/manager/README.md @@ -19,7 +19,7 @@ python3 -m manager doctor # diagnose problems `install.ps1`) just shell out to this CLI and stay dumb. Machine-readable output via `check --json` / `detect --json`. - **Hardware → spec mapping** is in [spec.py](spec.py). One universal torch - pin (2.12.0 / torchvision 0.27.0 / torchaudio 2.11.0) on every platform: + pin (2.13.0 / torchvision 0.28.0 / torchaudio 2.11.0) on every platform: cu130 wheels when the driver supports CUDA 13 (cu126 fallback for older drivers, refused outright on Blackwell GPUs which need cu130), same stack + Python 3.11 + `dgx_requirements.txt` on DGX/Grace, PyPI wheels on Mac, @@ -34,6 +34,18 @@ python3 -m manager doctor # diagnose problems (whl.natten.org) on Linux both arches; triton bundled with torch on Linux and `triton-windows` 3.7.x on Windows. No flash-attn/NATTEN/triton on Mac, no NATTEN on Windows (no wheels exist). +- **The torch stack is pinned against the resolver.** torch is an unpinned + transitive dep of timm/peft/accelerate/torchvision, so anything in + `requirements*.txt` that conflicts with what the pinned torch needs makes the + resolver silently backtrack to an older torch rather than fail — and prebuilt + accelerator wheels then reinstall a plain PyPI torch over the GPU one, + leaving torchvision/torchaudio's C++ extensions linked against a libtorch + that is gone. Every install pass after torch therefore carries a generated + constraints file (`torch==X+cu130`, ...) plus per-package `--find-links` to + the pytorch index, and `_verify_torch` re-checks (and repairs) the trio + before the optional accelerators are import-tested and again at the end. + Requirements files must never pin a torch dependency below what torch needs + (torch 2.13 wants `setuptools>=77.0.3`). - **Nothing global is ever installed.** FFmpeg (shared builds — the libs torchcodec dlopens) goes to `.ffmpeg/` ([ffmpeg.py](ffmpeg.py)), Node (when the system lacks >= 20) to `.node/` ([nodejs.py](nodejs.py)), the uv diff --git a/manager/detect.py b/manager/detect.py index 60e7980..4878157 100644 --- a/manager/detect.py +++ b/manager/detect.py @@ -50,9 +50,13 @@ def detect_nvidia(): driver = parts[2] if not gpus: return None - # Max CUDA version the driver supports only appears in the banner output + # Max CUDA version the driver supports only appears in the banner output. + # Windows/WDDM labels it "CUDA UMD Version:" on recent drivers, Linux just + # "CUDA Version:" — without the optional word we silently fall through to + # the "driver present but version unknown, assume current" path and can + # hand an old driver cu130 wheels it cannot run. banner = _run_quiet([smi]) or "" - m = re.search(r"CUDA Version:\s*([0-9]+\.[0-9]+)", banner) + m = re.search(r"CUDA(?:\s+[A-Z]+)?\s+Version:\s*([0-9]+\.[0-9]+)", banner) cuda_version = m.group(1) if m else None return {"gpus": gpus, "driver": driver, "cuda_version": cuda_version} diff --git a/manager/doctor.py b/manager/doctor.py index abb80a3..2b3043b 100644 --- a/manager/doctor.py +++ b/manager/doctor.py @@ -53,8 +53,37 @@ def run_doctor(): has_venv = env.venv_exists() _check("venv", has_venv, venv_dir() if has_venv else "not created yet") if has_venv: - torch = env.installed_torch() - _check("torch", torch is not None, torch or "not installed") + from . import spec as spec_mod + + stack = env.torch_stack() + torch = stack.get("torch") + if torch and torch.startswith("ERROR"): + torch = None + _check("torch", torch is not None, stack.get("torch") or "not installed") + # torchvision/torchaudio ship C++ extensions linked against libtorch — + # a version skew here fails at import, not at install + for name in ("torchvision", "torchaudio"): + found = stack.get(name) + _check( + name, + bool(found) and not found.startswith("ERROR"), + found or "not installed", + ) + try: + want = spec_mod.build_spec(d, allow_cpu=True) + _check( + "torch stack pins", + env.torch_matches(want), + "expected %s (%s)" + % ( + ", ".join( + "%s %s" % (k, v) for k, v in sorted(want.torch_packages.items()) + ), + want.backend, + ), + ) + except RuntimeError as e: + _check("torch stack pins", False, str(e)) if torch and d["backend"] == "cuda": try: out = subprocess.run( diff --git a/manager/env.py b/manager/env.py index 5f2e921..f376e49 100644 --- a/manager/env.py +++ b/manager/env.py @@ -173,17 +173,64 @@ def installed_torch(): return None -def torch_matches(spec): - current = installed_torch() - if current is None: +def torch_stack(): + """{package: version or 'ERROR: ...'} for torch, torchvision, torchaudio. + + Each one is imported rather than read from metadata: the failure mode this + guards against is a C++ extension (libtorchaudio, torchvision's ops) linked + against a libtorch that is no longer the installed one, and that only shows + up at import time — the recorded version looks perfectly fine. + """ + if not venv_exists(): + return {} + code = ( + "import importlib, json\n" + "r = {}\n" + "for m in ('torch', 'torchvision', 'torchaudio'):\n" + " try:\n" + " r[m] = importlib.import_module(m).__version__\n" + " except Exception as e:\n" + " r[m] = 'ERROR: %s: %s' % (type(e).__name__, e)\n" + "print(json.dumps(r))\n" + ) + try: + out = subprocess.run( + [venv_python(), "-c", code], + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + timeout=300, + env=clean_env(), + ) + return json.loads(out.stdout.decode().strip() or "{}") + except (OSError, subprocess.TimeoutExpired, ValueError): + return {} + + +def _version_matches(current, want, backend): + if not current or current.startswith("ERROR"): return False - want = spec.torch_packages["torch"] # local version tag carries the backend: "2.9.1+cu128" if "+" in current: version, local = current.split("+", 1) - return version == want and local == spec.backend + return version == want and local == backend # PyPI wheels (mac) have no local tag - return current == want and spec.backend in ("mps", "cpu") + return current == want and backend in ("mps", "cpu") + + +def torch_matches(spec): + """True only if the whole trio is at the pinned version AND imports. + + Checking torch alone is not enough — a resolver backtrack downgrades + torchvision while leaving torch untouched, and torchaudio keeps its version + number even when its extension can no longer load. + """ + stack = torch_stack() + if not stack: + return False + return all( + _version_matches(stack.get(name), version, spec.backend) + for name, version in spec.torch_packages.items() + ) def ensure_torch(spec, dry_run=False): @@ -208,6 +255,48 @@ def ensure_torch(spec, dry_run=False): return True +CONSTRAINTS_FILE = "aitk_torch_constraints.txt" + + +def _torch_pin_args(spec, dry_run=False): + """`--constraint`/`--find-links` args nailing torch down for later passes.""" + path = os.path.join(venv_dir(), CONSTRAINTS_FILE) + if not dry_run: + with open(path, "w") as f: + f.write( + "# Generated by the AI Toolkit manager (manager/env.py).\n" + "# Keeps requirements.txt and prebuilt accelerator wheels from\n" + "# replacing the GPU torch build. Do not edit.\n" + ) + f.write("\n".join(spec.torch_constraints()) + "\n") + args = ["--constraint", path] + for url in spec.torch_find_links(): + args += ["--find-links", url] + return args + + +def _verify_torch(spec, dry_run=False): + """Last line of defence: no install pass may leave torch swapped out. + + The constraints normally prevent this outright; this catches the cases they + can't (a package vendoring its own torch, a pip fallback that ignores the + constraint) before the venv is declared good. + """ + if dry_run or torch_matches(spec): + return False + stack = torch_stack() + found = ", ".join( + "%s %s" % (name, stack.get(name) or "missing") + for name in sorted(spec.torch_packages) + ) + warn( + "The PyTorch stack was disturbed during dependency install (%s) — " + "restoring the pinned %s build." % (found, spec.backend) + ) + _pip_install(spec.torch_args(), dry_run=dry_run) + return True + + # ---------------------------------------------------------------- requirements @@ -328,15 +417,20 @@ def ensure_requirements(spec, dry_run=False, force=False): if stale: info("Git pin changed — reinstalling: %s" % ", ".join(stale)) _pip_uninstall(stale, dry_run=dry_run) + # every pass below carries the torch pins so nothing can swap the GPU build + pins = _torch_pin_args(spec, dry_run=dry_run) info("Installing requirements from %s..." % spec.requirements_file) - _pip_install(["-r", spec.requirements_path()], dry_run=dry_run) - find_links = [] + _pip_install(["-r", spec.requirements_path()] + pins, dry_run=dry_run) + find_links = list(pins) for url in spec.find_links: find_links += ["--find-links", url] extras = _filter_extras(spec.extra_packages) if extras: info("Installing platform extras...") _pip_install(extras + find_links, dry_run=dry_run, upgrade=True) + # the optional import checks below only mean anything against the torch we + # actually intend to ship, so repair it first if something got through + _verify_torch(spec, dry_run=dry_run) # accelerators (flash-attn, NATTEN, ...): install one-by-one, warn on # failure — training works without them, so never fail the whole install for pkg in _filter_extras(spec.optional_packages): @@ -360,6 +454,7 @@ def ensure_requirements(spec, dry_run=False, force=False): "removing it (training falls back to native attention)." % name ) _pip_uninstall([name]) + _verify_torch(spec, dry_run=dry_run) if not dry_run: state = load_state() state["req_hash"] = requirements_hash(spec) diff --git a/manager/spec.py b/manager/spec.py index 3dd02e2..ce3d459 100644 --- a/manager/spec.py +++ b/manager/spec.py @@ -22,6 +22,13 @@ kernels smoke-tested on an RTX 5090 / sm120 with torch 2.13.0+cu130): - triton: bundled with torch on Linux (incl. aarch64; torch 2.13 bundles triton 3.7.1); triton-windows 3.7.x matches on Windows; nothing for MPS. +Requirements files must never pin anything torch itself depends on below what +the pinned torch needs (torch 2.13 wants setuptools>=77.0.3). The resolver does +not report that as a conflict — torch is an unpinned transitive dep of timm / +peft / accelerate / torchvision, so it just silently backtracks to an older +torch and drags torchvision down with it. `torch_constraints()` below turns +that class of mistake into a hard resolution error instead of a broken venv. + extra_packages are installed AFTER requirements.txt with --upgrade so they can override requirement pins (e.g. torchcodec). optional_packages are installed one-by-one and only warn on failure (accelerators the training code can live @@ -91,6 +98,36 @@ class EnvSpec(object): args += ["--index-url", self.torch_index] return args + def torch_constraints(self): + """Exact pins for the torch trio, carrying the backend local tag. + + Written to a constraints file and passed to every install pass after + torch itself, so nothing in requirements.txt (or a prebuilt accelerator + wheel that merely declares `torch`) can quietly swap the GPU build for + a PyPI one. Without this, a single conflicting pin anywhere in the tree + makes the resolver silently backtrack to an older torch instead of + failing, and the accelerator wheels then reinstall a plain PyPI torch + on top — which leaves torchaudio/torchvision linked against a libtorch + that is no longer there. + """ + tag = "+%s" % self.backend if self.torch_index else "" + return [ + "%s==%s%s" % (k, v, tag) for k, v in sorted(self.torch_packages.items()) + ] + + def torch_find_links(self): + """Per-package wheel pages making the constrained pins resolvable. + + Deliberately per-package `--find-links` rather than `--extra-index-url`: + the pytorch index also mirrors numpy/pillow/setuptools/... and with uv's + default first-index strategy an extra index would win for those too, + pinning them to whatever stale copy pytorch happens to host. + """ + if not self.torch_index: + return [] + base = self.torch_index.rstrip("/") + return ["%s/%s/" % (base, name) for name in sorted(self.torch_packages)] + def requirements_path(self): return os.path.join(REPO_ROOT, self.requirements_file) diff --git a/requirements_base.txt b/requirements_base.txt index 656b2c5..30f9451 100644 --- a/requirements_base.txt +++ b/requirements_base.txt @@ -35,7 +35,7 @@ python-slugify opencv-python pytorch-wavelets==1.3.0 matplotlib==3.10.1 -setuptools==69.5.1 +setuptools>=77.0.3 av==16.0.1 torchcodec==0.9.1 librosa==0.11.0