Fixed install issues with windows builds

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

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}
 

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()
+        from . import spec as spec_mod
-        _check("torch", torch is not None, torch or "not installed")
+
+        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(

manager/env.py

@@ -173,17 +173,64 @@ def installed_torch():
         return None
 
 
-def torch_matches(spec):
+def torch_stack():
-    current = installed_torch()
+    """{package: version or 'ERROR: ...'} for torch, torchvision, torchaudio.
-    if current is None:
+
+    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)
+    _pip_install(["-r", spec.requirements_path()] + pins, dry_run=dry_run)
-    find_links = []
+    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)

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)
 

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