github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Fixed install issues with windows builds

Browse Source
This commit is contained in:
Jaret Burkett
2026-07-27 16:06:48 -06:00
parent 1e58c9a0f0
commit 6b0449c326
6 changed files with 191 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
manager/README.md
View File

@@ -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

8
manager/detect.py
View File

@@ -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}

33
manager/doctor.py
View File

@@ -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(

111
manager/env.py
View File

@@ -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)

37
manager/spec.py
View File

@@ -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)

2
requirements_base.txt
View File

@@ -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