Fixed install issues with windows builds

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)