RomanNum3ral
/
ai-toolkit
mirror of https://github.com/ostris/ai-toolkit.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #986 from ostris/dev 

Add AI Toolkit Manager script that auto installs/runs/and updates AI Toolkit.
This commit is contained in:
Jaret Burkett (Ostris) 2026-07-30 11:55:13 -06:00 committed by GitHub
parent 7e7053fc9a c8636478f9
commit 2cab330392
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
45 changed files with 4469 additions and 356 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
.gitignore vendored
View File

@ -124,6 +124,9 @@ celerybeat.pid
.venv
.python
.node
.ffmpeg
.mingit
.uv
env/
venv/
ENV/

69
README.md
View File

@ -67,6 +67,57 @@ AI Toolkit is an easy to use all in one training suite for diffusion models. I t
## Installation
### Install with the AI Toolkit Manager (experimental)
The recommended way to install and run AI Toolkit is with the **AI Toolkit
Manager**, built into this repo. The manager detects your hardware and sets up
the right PyTorch build, creates the python environment, and grabs local copies
of Node.js and FFmpeg — everything stays inside the ai-toolkit folder, nothing
is installed system-wide. On every launch the manager checks for updates and
applies them (your local changes are never overwritten — if you have modified
files, the update is skipped with a warning), then starts the UI at
`http://localhost:8675`.
The manager is still **experimental** — please let me know if you have any
issues with it. The manual instructions below still work if you prefer them
or run into problems.
The only requirement is **git** (on Windows the manager can even fetch a
portable git for updates, but you need one installed to clone the repo first).
```bash
git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
```
Then start the manager with the script for your platform:
Linux:
```bash
chmod +x run_linux.sh
./run_linux.sh
```
MacOS (Apple Silicon, experimental):
```bash
chmod +x run_mac.zsh
./run_mac.zsh
```
Windows: double-click `run_windows.bat` (or run it from a terminal).
You can also use the manager directly from a terminal (handy on headless
servers):
```bash
python3 -m manager install # first-time setup
python3 -m manager update # pull updates + sync dependencies
python3 -m manager launch # start the UI
python3 -m manager doctor # diagnose problems
```
### Manual installation
Requirements:
- python >=3.10 (3.12 recommended)
- Nvidia GPU with enough ram to do what you need
@ -81,7 +132,7 @@ cd ai-toolkit
python3 -m venv venv
source venv/bin/activate
# install torch first
pip3 install --no-cache-dir torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/cu128
pip3 install --no-cache-dir torch==2.13.0 torchvision==0.28.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu130
pip3 install -r requirements.txt
```
@ -97,24 +148,10 @@ git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
python -m venv venv
.\venv\Scripts\activate
pip install --no-cache-dir torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/cu128
pip install --no-cache-dir torch==2.13.0 torchvision==0.28.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu130
pip install -r requirements.txt
```
MacOS:
Experimental support for Silicon Macs is available. I do not have a Mac with enough RAM to fully test this
so please let me know if there are issues. There is a convience script to install and run on MacOS
locates at `./run_mac.zsh` that will install the dependencies locally and run the UI. To run this,
do the following:
```bash
git clone https://github.com/ostris/ai-toolkit.git
cd ai-toolkit
chmod +x run_mac.zsh
./run_mac.zsh
```
# AI Toolkit UI

2
dgx_instructions.md
View File

@ -37,7 +37,7 @@ conda activate ai-toolkit
**2) Install PyTorch**
```
pip3 install torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/cu130
pip3 install torch==2.13.0 torchvision==0.28.0 torchaudio==2.11.0 --index-url https://download.pytorch.org/whl/cu130
```

10
extensions_built_in/captioner/AceStepCaptioner.py
View File

@ -1,6 +1,9 @@
from typing import Optional
import librosa
try:
import librosa
except ImportError:
librosa = None
import numpy as np
import torch
import torchaudio
@ -41,6 +44,11 @@ KEY_NAMES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]
def analyze_audio(audio_path):
"""Extract BPM, key, and time signature from audio using librosa."""
if librosa is None:
raise ImportError(
"librosa is required for the AceStep captioner but is not "
"installed (no numba/llvmlite wheels for this platform yet)."
)
y, sr = librosa.load(audio_path, sr=22050, mono=True)
duration = librosa.get_duration(y=y, sr=sr)

33
extensions_built_in/captioner/BaseCaptioner.py
View File

@ -274,26 +274,21 @@ class BaseCaptioner(BaseExtensionProcess):
while True:
try:
if self.should_stop():
# Mark and update status (non-blocking; uses existing infra)
self.is_stopping = True
self._run_async_operation(
self._update_status("stopped", "Job stopped (remote)")
)
# Best-effort flush pending async ops
try:
asyncio.run(self.wait_for_all_async())
except RuntimeError:
pass
# Try to stop DB thread pool quickly
try:
self.thread_pool.shutdown(wait=False, cancel_futures=True)
except TypeError:
self.thread_pool.shutdown(wait=False)
if self.is_stopping:
# maybe_stop() already started the graceful shutdown;
# a second interrupt would only break its cleanup.
return
print("")
print("****************************************************")
print(" Stop signal received; terminating process. ")
print("****************************************************")
os.kill(os.getpid(), signal.SIGINT)
# Deliver a real KeyboardInterrupt to the main thread so
# on_error runs the normal shutdown (final DB write, last
# log). os.kill(pid, SIGINT) must not be used here: on
# Windows it is TerminateProcess and kills us instantly.
# Leave the thread pool alone -- on_error still needs it.
signal.raise_signal(signal.SIGINT)
return
time.sleep(interval_sec)
except Exception:
time.sleep(interval_sec)
@ -455,7 +450,11 @@ class BaseCaptioner(BaseExtensionProcess):
super(BaseCaptioner, self).on_error(e)
if self.is_ui_captioner:
try:
if not self.is_stopping:
if isinstance(e, KeyboardInterrupt):
# SIGINT (UI stop button or ctrl+c) is a stop, not an error
self.is_stopping = True
self.update_status("stopped", "Job stopped")
elif not self.is_stopping:
self.update_status("error", str(e))
asyncio.run(self.wait_for_all_async())
except Exception as db_err:

34
extensions_built_in/sd_trainer/DiffusionTrainer.py
View File

@ -39,7 +39,12 @@ class DiffusionTrainer(SDTrainer):
# Initialize the status
self._run_async_operation(self._update_status("running", "Starting"))
self._stop_watcher_started = False
# self.start_stop_watcher(interval_sec=2.0)
if os.name == "nt":
# On Windows the stop route cannot send us SIGINT from outside
# (no console to deliver a Ctrl+C to), so watch the stop flag
# and raise the interrupt from inside. On Linux the route
# sends a real SIGINT to the pid and this is unnecessary.
self.start_stop_watcher(interval_sec=2.0)
def start_stop_watcher(self, interval_sec: float = 5.0):
"""
@ -60,26 +65,21 @@ class DiffusionTrainer(SDTrainer):
while True:
try:
if self.should_stop():
# Mark and update status (non-blocking; uses existing infra)
self.is_stopping = True
self._run_async_operation(
self._update_status("stopped", "Job stopped (remote)")
)
# Best-effort flush pending async ops
try:
asyncio.run(self.wait_for_all_async())
except RuntimeError:
pass
# Try to stop DB thread pool quickly
try:
self.thread_pool.shutdown(wait=False, cancel_futures=True)
except TypeError:
self.thread_pool.shutdown(wait=False)
if self.is_stopping:
# maybe_stop() already started the graceful shutdown;
# a second interrupt would only break its cleanup.
return
print("")
print("****************************************************")
print(" Stop signal received; terminating process. ")
print("****************************************************")
os.kill(os.getpid(), signal.SIGINT)
# Deliver a real KeyboardInterrupt to the main thread so
# on_error runs the normal shutdown (final DB write, last
# log). os.kill(pid, SIGINT) must not be used here: on
# Windows it is TerminateProcess and kills us instantly.
# Leave the thread pool alone -- on_error still needs it.
signal.raise_signal(signal.SIGINT)
return
time.sleep(interval_sec)
except Exception:
time.sleep(interval_sec)

41
extensions_built_in/sd_trainer/UITrainer.py
View File

@ -33,7 +33,12 @@ class UITrainer(SDTrainer):
# Initialize the status
self._run_async_operation(self._update_status("running", "Starting"))
self._stop_watcher_started = False
# self.start_stop_watcher(interval_sec=2.0)
if os.name == "nt":
# On Windows the stop route cannot send us SIGINT from outside
# (no console to deliver a Ctrl+C to), so watch the stop flag
# and raise the interrupt from inside. On Linux the route
# sends a real SIGINT to the pid and this is unnecessary.
self.start_stop_watcher(interval_sec=2.0)
def start_stop_watcher(self, interval_sec: float = 5.0):
"""
@ -52,26 +57,21 @@ class UITrainer(SDTrainer):
while True:
try:
if self.should_stop():
# Mark and update status (non-blocking; uses existing infra)
self.is_stopping = True
self._run_async_operation(
self._update_status("stopped", "Job stopped (remote)")
)
# Best-effort flush pending async ops
try:
asyncio.run(self.wait_for_all_async())
except RuntimeError:
pass
# Try to stop DB thread pool quickly
try:
self.thread_pool.shutdown(wait=False, cancel_futures=True)
except TypeError:
self.thread_pool.shutdown(wait=False)
if self.is_stopping:
# maybe_stop() already started the graceful shutdown;
# a second interrupt would only break its cleanup.
return
print("")
print("****************************************************")
print(" Stop signal received; terminating process. ")
print("****************************************************")
os.kill(os.getpid(), signal.SIGINT)
# Deliver a real KeyboardInterrupt to the main thread so
# on_error runs the normal shutdown (final DB write, last
# log). os.kill(pid, SIGINT) must not be used here: on
# Windows it is TerminateProcess and kills us instantly.
# Leave the thread pool alone -- on_error still needs it.
signal.raise_signal(signal.SIGINT)
return
time.sleep(interval_sec)
except Exception:
time.sleep(interval_sec)
@ -217,7 +217,12 @@ class UITrainer(SDTrainer):
def on_error(self, e: Exception):
super(UITrainer, self).on_error(e)
if self.accelerator.is_main_process and not self.is_stopping:
if isinstance(e, KeyboardInterrupt):
# SIGINT (UI stop button or ctrl+c) is a stop, not an error
self.is_stopping = True
if self.accelerator.is_main_process:
self.update_status("stopped", "Job stopped")
elif self.accelerator.is_main_process and not self.is_stopping:
self.update_status("error", str(e))
self.update_db_key("step", self.last_save_step)
asyncio.run(self.wait_for_all_async())

95
manager/README.md Normal file
View File

@ -0,0 +1,95 @@
# AI Toolkit Manager
Self-contained install/update manager for this checkout of AI Toolkit. Runs
with any Python >= 3.8 and **no dependencies**, so it works before the
training environment exists.
```bash
python3 -m manager install # first-time setup: venv + torch + requirements
python3 -m manager check # is an update available / are deps out of sync?
python3 -m manager update # git pull, then sync deps + run migrations
python3 -m manager launch # start the web UI (http://localhost:8675)
python3 -m manager doctor # diagnose problems
```
## Design
- **The install logic lives in the repo it installs.** Every commit knows how
to install itself; external frontends (the desktop launcher, `install.sh`,
`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.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,
rocm7.1 (experimental) for AMD, `--cpu` to force a CPU install. **Torch
pins there must be updated together with the README install instructions,
run_mac.zsh, and dgx_instructions.md.**
- **Accelerators everywhere wheels exist**, via per-spec `extra_packages`
(installed after requirements with `--upgrade` so they override pins) and
`optional_packages` (installed one-by-one, warn-only on failure):
`torchcodec==0.15.0` on all platforms; flash-attn 2.8.3 prebuilt wheels
(mjun0812) on Linux x86_64/aarch64 + Windows; NATTEN 0.21.7 wheels
(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`).
- **`ui/package-lock.json` is never modified by installing.** A plain
`npm install` re-derives the lockfile for whichever machine runs it — on
Windows it strips the `libc` fields off the Linux-only optional binaries
(`@next/swc-linux-*`, rollup, lightningcss), on Linux it adds them back — so
with the install baked into `npm run build_and_start` every user got a dirty
tree on every launch, which then blocks `manager update` (a dirty tree aborts
the pull). `nodejs.ensure_ui_deps` owns the install instead: `npm install
--no-save`, gated on a hash of `ui/package.json` + `ui/package-lock.json`
(stored in the venv state), with the lockfile bytes snapshotted and restored
either way. `manager launch` therefore runs `npm run db_build_start`, not
`build_and_start`; the latter still exists for the manual
`cd ui && npm run build_and_start` flow in the README and calls the same
non-writing install via `npm run install_deps`.
- **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
binary (when absent) to `.uv/` ([uvbin.py](uvbin.py)) with uv-managed
Pythons kept in `.uv/python/` via `UV_PYTHON_INSTALL_DIR`, and on Windows
without git, portable MinGit to `.mingit/` ([gitwin.py](gitwin.py)) — all
inside the repo and gitignored. The first clone on a git-less Windows
box is handled by the bootstrap layer (install.ps1 / desktop launcher),
which downloads MinGit itself and moves it into the checkout afterwards. `manager launch` puts them on PATH (and
LD_LIBRARY_PATH on Linux) for the whole UI/training process tree, and a
generated `sitecustomize.py` in the venv exposes ffmpeg to any direct use
of the venv python (plus `os.add_dll_directory` on Windows).
- **Hostile-environment hardening** (learned from the community Windows
installer): every python/pip subprocess runs with PYTHONPATH/PYTHONHOME/
CONDA/PYENV/PIP_* scrubbed from the env; git runs with
`GIT_LFS_SKIP_SMUDGE=1`; git-pinned requirements (diffusers) are
force-reinstalled when requirements change since pip skips unchanged
version numbers; `launch` polls the UI port and opens the browser when
ready (`--no-browser` to disable, auto-skipped on headless boxes).
- **uv is used when present** (fast installs, auto-downloads the right
Python); plain `venv` + `pip` otherwise. The venv is created at `.venv/`
(an existing `venv/` is also respected, matching `ui/cron/pythonPath.ts`).
- **State** (requirements hash, applied migrations) lives inside the venv
(`aitk_manager_state.json`) — deleting the venv resets everything.
- **Update flow**: `update` pulls fast-forward only, then **re-execs**
`python -m manager sync` so the freshly pulled manager code — not the stale
in-memory copy — performs its own dependency sync and migrations.
**Local work is never overwritten**: a dirty tree aborts the update by
default (untracked files don't count), `--auto` (used by the run_* scripts)
warns and skips the pull instead so launching still works, and there is no
reset/clean anywhere — even `--force` relies on git itself refusing to
clobber modified files.
- **Migrations** ([migrations.py](migrations.py)): one-time post-update steps,
each applied at most once per environment.

6
manager/__init__.py Normal file
View File

@ -0,0 +1,6 @@
"""AI Toolkit install / update manager. Stdlib-only; see __main__.py."""
# Version of the CLI contract consumed by external frontends (desktop
# launcher, install scripts). Bump only on breaking changes to command
# names/flags or --json output shapes.
CLI_CONTRACT_VERSION = 1

249
manager/__main__.py Normal file
View File

@ -0,0 +1,249 @@
"""AI Toolkit manager CLI.
Runs with any Python >= 3.8 and no dependencies, so it works before the
training environment exists. This is the single entry point every installer
frontend (shell scripts, the desktop launcher, the web UI) shells out to.
python3 -m manager install first-time environment setup
python3 -m manager check [--json] is an update / dep sync needed?
python3 -m manager update git pull + dependency sync + migrations
python3 -m manager sync dependency sync only (no git pull)
python3 -m manager launch start the web UI
python3 -m manager detect [--json] show detected hardware
python3 -m manager doctor full environment diagnostics
"""
import argparse
import os
import subprocess
import sys
# allow `python manager/__main__.py` as well as `python -m manager`
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from manager import detect as detect_mod
from manager import env, gitops, launch, spec as spec_mod, util
from manager.util import die, info, ok, print_json, warn
def _resolve_spec(args):
detection = detect_mod.detect()
try:
return detection, spec_mod.build_spec(
detection, allow_cpu=getattr(args, "cpu", False)
)
except RuntimeError as e:
die(str(e))
def cmd_detect(args):
detection = detect_mod.detect()
try:
s = spec_mod.build_spec(detection, allow_cpu=True)
detection["spec"] = s.as_dict()
except RuntimeError as e:
detection["spec_error"] = str(e)
if args.json:
print_json(detection)
else:
backend = detection.get("spec", {}).get("backend", "unknown")
info("os=%s arch=%s backend=%s" % (detection["os"], detection["arch"], backend))
if detection["nvidia"]:
for gpu in detection["nvidia"]["gpus"]:
info("gpu: %s (%s)" % (gpu["name"], gpu["memory"]))
def cmd_install(args):
detection, s = _resolve_spec(args)
env.sync(s, detection, dry_run=args.dry_run, force=args.force)
if not args.dry_run:
ok("Install complete. Start the UI with: python3 -m manager launch")
def cmd_sync(args):
detection, s = _resolve_spec(args)
env.sync(s, detection, dry_run=args.dry_run, force=args.force)
def cmd_check(args):
_, s = _resolve_spec(args)
fetched = gitops.fetch()
behind = gitops.behind_count()
data = {
"version": _toolkit_version(),
"branch": gitops.current_branch(),
"commit": gitops.current_commit(),
"remote_commit": gitops.remote_commit(),
"dirty": gitops.is_dirty(),
"fetch_ok": fetched,
"behind": behind,
"incoming": gitops.incoming_log(),
"venv": env.venv_exists(),
"deps_in_sync": env.venv_exists()
and env.torch_matches(s)
and env.requirements_in_sync(s),
"backend": s.backend,
}
data["update_available"] = bool(behind) or not data["deps_in_sync"]
if args.json:
print_json(data)
return
info("AI Toolkit %s (%s @ %s)" % (data["version"], data["branch"], data["commit"]))
if not fetched:
warn("Could not reach the remote (offline?) — update status may be stale.")
if behind:
info("Update available: %d new commit(s)." % behind)
for line in data["incoming"]:
print(" " + line)
elif behind == 0:
ok("Code is up to date.")
if not data["deps_in_sync"]:
warn("Dependencies are out of sync. Run: python3 -m manager sync")
elif behind == 0:
ok("Dependencies are in sync.")
def cmd_update(args):
"""git pull (never destructive) + dependency sync.
Local work is sacred: a dirty tree either aborts (default), or with
--auto is skipped with a warning so run scripts can continue to launch.
We never reset/clean; even a forced pull is --ff-only, which git itself
aborts rather than overwriting local changes.
"""
auto = getattr(args, "auto", False)
skip_pull = False
if gitops.is_dirty() and not args.force:
if auto:
warn(
"Local changes detected — skipping the code update to protect "
"your work. Commit or stash your changes to receive updates."
)
skip_pull = True
else:
die(
"You have local changes to tracked files. Commit or stash them, "
"or re-run with --force to attempt the update anyway (git will "
"still refuse rather than overwrite your changes)."
)
if not skip_pull and not gitops.fetch():
if auto:
warn("Could not reach the git remote — skipping the update check.")
skip_pull = True
else:
die("Could not reach the git remote. Check your network and try again.")
if not skip_pull:
behind = gitops.behind_count()
if behind is None:
warn("Current branch has no upstream; skipping git pull.")
elif behind == 0:
ok("Code already up to date.")
else:
info("Pulling %d new commit(s)..." % behind)
gitops.pull_ff()
ok("Code updated to %s." % gitops.current_commit())
# Re-exec so the freshly pulled manager code runs its own dependency
# sync and migrations (the in-memory copy of this module is stale now).
cmd = [sys.executable, "-m", "manager", "sync"]
if args.dry_run:
cmd.append("--dry-run")
sys.exit(subprocess.call(cmd, cwd=util.REPO_ROOT))
# nothing was pulled — safe to sync with the code already loaded
detection, s = _resolve_spec(args)
env.sync(s, detection, dry_run=args.dry_run)
def cmd_launch(args):
sys.exit(launch.launch_ui(open_browser=not args.no_browser))
def cmd_doctor(args):
from manager import doctor
doctor.run_doctor()
def cmd_version(args):
print(_toolkit_version())
def _toolkit_version():
version = {}
try:
with open(os.path.join(util.REPO_ROOT, "version.py")) as f:
exec(f.read(), version)
return version.get("VERSION", "unknown")
except OSError:
return "unknown"
def main(argv=None):
parser = argparse.ArgumentParser(
prog="manager", description="AI Toolkit install / update manager"
)
sub = parser.add_subparsers(dest="command")
def add(name, fn, **kwargs):
p = sub.add_parser(name, **kwargs)
p.set_defaults(fn=fn)
return p
p = add("detect", cmd_detect, help="show detected hardware and env spec")
p.add_argument("--json", action="store_true")
for name, fn, help_text in (
("install", cmd_install, "first-time environment setup"),
("sync", cmd_sync, "sync dependencies for the current checkout"),
):
p = add(name, fn, help=help_text)
p.add_argument("--cpu", action="store_true", help="allow CPU-only install")
p.add_argument("--dry-run", action="store_true")
p.add_argument(
"--force",
action="store_true",
help="reinstall requirements even if in sync",
)
p = add("check", cmd_check, help="check for updates (use --json for machines)")
p.add_argument("--json", action="store_true")
p.add_argument("--cpu", action="store_true", help=argparse.SUPPRESS)
p = add("update", cmd_update, help="git pull + dependency sync + migrations")
p.add_argument("--cpu", action="store_true", help="allow CPU-only install")
p.add_argument("--dry-run", action="store_true")
p.add_argument(
"--force", action="store_true", help="update even with local changes"
)
p.add_argument(
"--auto",
action="store_true",
help="unattended mode (run scripts): on local changes or an unreachable "
"remote, warn and skip the code update instead of failing; deps still sync",
)
p = add("launch", cmd_launch, help="start the web UI")
p.add_argument(
"--no-browser",
action="store_true",
help="do not open a browser when the UI is ready",
)
add("doctor", cmd_doctor, help="diagnose the environment")
add("version", cmd_version, help="print the toolkit version")
args = parser.parse_args(argv)
if not getattr(args, "command", None):
parser.print_help()
return 1
util.set_json_mode(bool(getattr(args, "json", False)))
try:
args.fn(args)
except KeyboardInterrupt:
return 130
return 0
if __name__ == "__main__":
sys.exit(main())

114
manager/detect.py Normal file
View File

@ -0,0 +1,114 @@
"""Hardware / platform detection. Stdlib only, safe to run anywhere."""
import os
import platform
import re
import subprocess
from .util import which, IS_WINDOWS, IS_MAC, IS_LINUX
def _run_quiet(cmd):
try:
out = subprocess.run(
cmd, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, timeout=15
)
if out.returncode != 0:
return None
return out.stdout.decode("utf-8", errors="replace")
except (FileNotFoundError, subprocess.TimeoutExpired, OSError):
return None
def detect_nvidia():
"""Returns dict with gpus + driver info, or None if no working nvidia-smi."""
smi = which("nvidia-smi")
if not smi:
return None
fields = "name,memory.total,driver_version,compute_cap"
csv = _run_quiet([smi, "--query-gpu=" + fields, "--format=csv,noheader"])
if not csv:
# older drivers don't know the compute_cap field
csv = _run_quiet(
[
smi,
"--query-gpu=name,memory.total,driver_version",
"--format=csv,noheader",
]
)
if not csv:
return None
gpus = []
driver = None
for line in csv.strip().splitlines():
parts = [p.strip() for p in line.split(",")]
if len(parts) >= 3:
gpu = {"name": parts[0], "memory": parts[1]}
if len(parts) >= 4:
gpu["compute_cap"] = parts[3]
gpus.append(gpu)
driver = parts[2]
if not gpus:
return None
# 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(?:\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}
def detect_rocm():
"""Returns dict if an AMD ROCm stack is present, else None."""
if not IS_LINUX:
return None
smi = which("rocm-smi")
if not smi and not os.path.isdir("/opt/rocm"):
return None
gpus = []
out = _run_quiet([smi, "--showproductname"]) if smi else None
if out:
for m in re.finditer(r"Card [Ss]eries:\s*(.+)", out):
gpus.append({"name": m.group(1).strip()})
return {"gpus": gpus}
def detect():
"""Full platform detection. Returns a plain dict (json-serializable)."""
system = platform.system() # Linux / Darwin / Windows
arch = platform.machine().lower() # x86_64 / amd64 / arm64 / aarch64
if arch == "amd64":
arch = "x86_64"
if arch == "arm64" and not IS_MAC:
arch = "aarch64"
result = {
"os": {"Linux": "linux", "Darwin": "mac", "Windows": "windows"}.get(
system, system.lower()
),
"arch": arch,
"python": platform.python_version(),
"nvidia": None,
"rocm": None,
"backend": "cpu",
}
nvidia = detect_nvidia()
if nvidia:
result["nvidia"] = nvidia
result["backend"] = "cuda"
else:
rocm = detect_rocm()
if rocm:
result["rocm"] = rocm
result["backend"] = "rocm"
if IS_MAC:
result["backend"] = "mps" if arch == "arm64" else "cpu"
# DGX OS / Grace (GB10, DGX Spark): NVIDIA GPU on aarch64 Linux
result["is_dgx"] = bool(result["os"] == "linux" and arch == "aarch64" and nvidia)
return result

168
manager/doctor.py Normal file
View File

@ -0,0 +1,168 @@
"""Environment diagnostics: `python -m manager doctor`."""
import os
import shutil
import subprocess
import sys
from . import detect as detect_mod
from . import env, ffmpeg, gitops, nodejs
from .util import REPO_ROOT, clean_env, find_uv, venv_dir, venv_python
def _check(label, passed, detail=""):
if sys.stdout.isatty():
mark = "\033[32mOK\033[0m " if passed else "\033[31mFAIL\033[0m"
else:
mark = "OK " if passed else "FAIL"
print(" [%s] %-18s %s" % (mark, label, detail))
return passed
def run_doctor():
print("AI Toolkit doctor\n")
d = detect_mod.detect()
from . import gitwin
arch_detail = "%s %s" % (d["os"], d["arch"])
if d["os"] == "windows" and d["arch"] == "aarch64":
from . import spec as spec_mod_arch
try:
_s = spec_mod_arch.build_spec(d, allow_cpu=True)
if _s.backend == "cu134":
arch_detail += " (RTX Spark: native win_arm64 CUDA stack)"
else:
arch_detail += " (Windows-on-ARM: x64 stack via emulation)"
except RuntimeError:
pass
_check("os / arch", True, arch_detail)
git = gitwin.find_git()
_check(
"git",
git is not None,
git or "not found (manager sync installs a local copy on Windows)",
)
uv = find_uv()
_check("uv", True, uv or "not found (optional, recommended)")
if d["nvidia"]:
names = ", ".join(g["name"] for g in d["nvidia"]["gpus"])
_check(
"gpu",
True,
"%s (driver %s, CUDA %s)"
% (names, d["nvidia"]["driver"], d["nvidia"]["cuda_version"]),
)
elif d["rocm"]:
_check("gpu", True, "AMD ROCm (experimental)")
elif d["backend"] == "mps":
_check("gpu", True, "Apple Silicon (MPS)")
else:
_check("gpu", False, "no supported GPU detected")
has_venv = env.venv_exists()
_check("venv", has_venv, venv_dir() if has_venv else "not created yet")
if has_venv:
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(
[
venv_python(),
"-c",
"import torch; print(torch.cuda.is_available())",
],
stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL,
timeout=120,
env=clean_env(),
)
avail = out.stdout.decode().strip() == "True"
_check(
"torch sees gpu",
avail,
"" if avail else "torch.cuda.is_available() is False",
)
except (OSError, subprocess.TimeoutExpired):
_check("torch sees gpu", False, "could not query")
node_exe, node_major = nodejs.have_usable_node()
if node_exe:
_check("node", True, "%s (v%s)" % (node_exe, node_major))
else:
_check(
"node",
False,
"none >= %d found (manager sync installs a local copy)"
% nodejs.MIN_NODE_MAJOR,
)
if os.path.isfile(ffmpeg.ffmpeg_exe()):
# run it with the launch env so missing shared libs are caught
from . import launch
try:
out = subprocess.run(
[ffmpeg.ffmpeg_exe(), "-version"],
stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL,
timeout=30,
env=launch.build_env(),
)
works = out.returncode == 0
detail = (
out.stdout.decode().splitlines()[0]
if works
else "installed but fails to run"
)
except (OSError, subprocess.TimeoutExpired):
works, detail = False, "installed but fails to run"
_check("ffmpeg (local)", works, detail)
else:
_check("ffmpeg (local)", False, "not installed (manager sync installs it)")
try:
free_gb = shutil.disk_usage(REPO_ROOT).free / (1024**3)
_check("disk space", free_gb > 30, "%.0f GB free" % free_gb)
except OSError:
pass
branch = gitops.current_branch()
_check(
"git checkout",
True,
"%s @ %s%s"
% (branch, gitops.current_commit(), " (dirty)" if gitops.is_dirty() else ""),
)

707
manager/env.py Normal file
View File

@ -0,0 +1,707 @@
"""Python environment provisioning and dependency sync.
Strategy:
- If a venv already exists (.venv or venv), use it.
- Otherwise create one: prefer uv (downloads the exact Python version needed),
fall back to the running Python's venv module if it is new enough.
- Installs go through `uv pip` when uv is available (much faster), else pip.
State (torch backend, requirements hash, applied migrations) is stored inside
the venv so a deleted venv means a clean slate — which is correct.
"""
import json
import os
import subprocess
import sys
from .util import (
REPO_ROOT,
IS_WINDOWS,
clean_env,
die,
file_hash,
find_uv,
info,
ok,
run,
venv_dir,
venv_python,
warn,
)
STATE_FILE = "aitk_manager_state.json"
MIN_SYSTEM_PYTHON = (3, 10)
# ---------------------------------------------------------------- state
def state_path(venv=None):
return os.path.join(venv or venv_dir(), STATE_FILE)
def load_state():
try:
with open(state_path(), "r") as f:
return json.load(f)
except (OSError, ValueError):
return {}
def save_state(state):
with open(state_path(), "w") as f:
json.dump(state, f, indent=2)
# ---------------------------------------------------------------- venv
def venv_exists():
return os.path.isfile(venv_python())
def _venv_platform():
"""sysconfig platform of the existing venv ('win-amd64', 'win-arm64', ...)."""
if not venv_exists():
return None
try:
out = subprocess.run(
[venv_python(), "-c", "import sysconfig; print(sysconfig.get_platform())"],
stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL,
timeout=30,
env=clean_env(),
)
if out.returncode != 0:
return None
return out.stdout.decode().strip() or None
except (OSError, subprocess.TimeoutExpired):
return None
def _uv_python_platform(uv_python):
"""'win-arm64' / 'win-amd64' expected for a pinned uv interpreter request."""
if not uv_python:
return None
if "windows-aarch64" in uv_python:
return "win-arm64"
if "windows-x86_64" in uv_python:
return "win-amd64"
return None
def ensure_venv(spec, dry_run=False):
"""Create the venv if missing. Returns path to the venv python."""
if venv_exists():
# Switching stacks (e.g. Spark emulated x64 <-> native arm64) needs a
# different interpreter arch; the venv is disposable by design, so
# recreate it rather than install unresolvable wheels into it.
want = _uv_python_platform(spec.uv_python)
have = _venv_platform()
if want and have and want != have:
if dry_run:
info(
"[dry-run] venv is %s but this spec needs %s — would "
"recreate the venv." % (have, want)
)
return venv_python()
warn(
"Existing venv is %s but this spec needs %s — recreating the "
"venv (all packages will be reinstalled)." % (have, want)
)
import shutil
shutil.rmtree(venv_dir(), ignore_errors=True)
else:
return venv_python()
if venv_exists():
return venv_python()
target = venv_dir()
uv = find_uv()
# spec.uv_python pins the full interpreter build (arch included) where the
# default choice would be wrong — e.g. Windows-on-ARM must stay x86_64
python_request = spec.uv_python or spec.python_version
if dry_run:
info(
"[dry-run] would create venv at %s (python %s, via %s)"
% (target, python_request, "uv" if uv else "venv module")
)
return venv_python(target)
if uv:
info("Creating venv with uv (python %s) at %s" % (python_request, target))
run(
[uv, "venv", target, "--python", python_request, "--seed"],
env=clean_env(),
)
else:
if sys.version_info < MIN_SYSTEM_PYTHON:
die(
"Python %d.%d is too old (need >= %d.%d) and uv is not installed.\n"
"Install uv (https://docs.astral.sh/uv/) or a newer Python, then re-run."
% (sys.version_info[:2] + MIN_SYSTEM_PYTHON)
)
if spec.uv_python:
warn(
"uv not found — the venv needs the %s interpreter and the "
"system Python may be a different build. Install uv if the "
"torch install below fails to resolve." % spec.uv_python
)
pyver = "%d.%d" % (sys.version_info[:2])
if pyver != spec.python_version:
warn(
"Recommended Python is %s but using system Python %s "
"(install uv to get the exact version automatically)."
% (spec.python_version, pyver)
)
info("Creating venv at %s" % target)
run([sys.executable, "-m", "venv", target])
ok("Virtual environment ready.")
return venv_python(target)
def _pip_install(args, dry_run=False, upgrade=False, check=True):
"""Install into the venv, via uv pip if available. Returns exit code."""
uv = find_uv()
if uv:
cmd = [uv, "pip", "install", "--python", venv_python()]
else:
cmd = [venv_python(), "-m", "pip", "install"]
if upgrade:
cmd.append("--upgrade")
cmd += args
if dry_run:
info("[dry-run] would run: %s" % " ".join(cmd))
return 0
code, _ = run(cmd, stream=True, env=clean_env(), check=check)
return code
def _pip_install_no_deps(pkg, dry_run=False):
"""Install a single package with --no-deps. Returns exit code."""
uv = find_uv()
if uv:
cmd = [uv, "pip", "install", "--python", venv_python(), "--no-deps", pkg]
else:
cmd = [venv_python(), "-m", "pip", "install", "--no-deps", pkg]
if dry_run:
info("[dry-run] would run: %s" % " ".join(cmd))
return 0
code, _ = run(cmd, stream=True, env=clean_env(), check=False)
return code
def _pip_uninstall(packages, dry_run=False):
uv = find_uv()
if uv:
cmd = [uv, "pip", "uninstall", "--python", venv_python()] + packages
else:
cmd = [venv_python(), "-m", "pip", "uninstall", "-y"] + packages
if dry_run:
info("[dry-run] would run: %s" % " ".join(cmd))
return
# non-fatal: the package may simply not be installed yet
run(cmd, check=False, env=clean_env())
def venv_python_version():
"""'3.12' etc. from the venv interpreter, or None."""
if not venv_exists():
return None
try:
out = subprocess.run(
[venv_python(), "-c", "import sys; print('%d.%d' % sys.version_info[:2])"],
stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL,
timeout=30,
env=clean_env(),
)
if out.returncode != 0:
return None
return out.stdout.decode().strip() or None
except (OSError, subprocess.TimeoutExpired):
return None
# ---------------------------------------------------------------- torch
def installed_torch():
"""Returns torch.__version__ from the venv (e.g. '2.9.1+cu128'), or None."""
if not venv_exists():
return None
try:
out = subprocess.run(
[venv_python(), "-c", "import torch; print(torch.__version__)"],
stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL,
timeout=120,
env=clean_env(),
)
if out.returncode != 0:
return None
return out.stdout.decode().strip() or None
except (OSError, subprocess.TimeoutExpired):
return 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
# local version tag carries the backend: "2.9.1+cu128"
if "+" in current:
version, local = current.split("+", 1)
return version == want and local == backend
# PyPI wheels (mac) have no local tag
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):
if torch_matches(spec):
ok(
"PyTorch %s (%s) already installed."
% (spec.torch_packages["torch"], spec.backend)
)
return False
current = installed_torch()
if current:
info(
"PyTorch %s installed, need %s (%s) — reinstalling."
% (current, spec.torch_packages["torch"], spec.backend)
)
else:
info(
"Installing PyTorch %s (%s)..."
% (spec.torch_packages["torch"], spec.backend)
)
_pip_install(spec.torch_args(), dry_run=dry_run)
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
def requirements_hash(spec):
"""Hash of every requirements file plus the spec itself."""
req_files = [
os.path.join(REPO_ROOT, f)
for f in os.listdir(REPO_ROOT)
if f.startswith("requirements") and f.endswith(".txt")
]
req_files.append(os.path.join(REPO_ROOT, "dgx_requirements.txt"))
base = file_hash(req_files)
import hashlib
h = hashlib.sha256()
h.update(base.encode())
h.update(json.dumps(spec.as_dict(), sort_keys=True).encode())
return h.hexdigest()
def requirements_in_sync(spec):
if not venv_exists():
return False
return load_state().get("req_hash") == requirements_hash(spec)
def _git_pinned_packages(spec):
"""{package_name: full git+ requirement line} from the requirements files.
pip skips reinstalling a git pin whose version number didn't change even
when the commit hash did, so pins whose URL changed since the last sync
get uninstalled first to force the new commit (the trick the community
Windows installer uses for diffusers).
"""
pins = {}
seen_files = set()
def scan(path):
if path in seen_files or not os.path.isfile(path):
return
seen_files.add(path)
with open(path) as f:
for line in f:
line = line.strip()
if line.startswith("-r "):
scan(os.path.join(os.path.dirname(path), line[3:].strip()))
elif "git+" in line and not line.startswith("#"):
# e.g. git+https://github.com/huggingface/diffusers.git@<sha>
tail = line.split("/")[-1]
name = tail.split(".git")[0].split("@")[0]
if name:
pins[name] = line
scan(spec.requirements_path())
return pins
def _stale_git_pins(spec):
"""Git-pinned packages whose pin (commit) changed since the last sync."""
current = _git_pinned_packages(spec)
state = load_state()
stored = state.get("git_pins")
if stored is None:
# no record of what's installed (pre-tracking env): if deps were ever
# installed here, play it safe and force-reinstall all git pins once
return list(current) if state.get("req_hash") else []
return [name for name, line in current.items() if stored.get(name) != line]
# optional packages whose import name differs from the distribution name
_IMPORT_ALIASES = {"flash_linear_attention": "fla"}
# companion dists to remove on rollback (fla-core provides the `fla` module
# itself; leaving it behind would keep the broken import resolvable)
_ROLLBACK_EXTRAS = {"flash_linear_attention": ["fla-core"]}
def _optional_names(pkg):
"""(distribution, import) names for an optional package spec or wheel URL."""
if "://" in pkg:
name = os.path.basename(pkg).split("-")[0]
else:
name = pkg
for sep in ("==", ">=", "<=", "<", ">", "["):
name = name.split(sep)[0]
name = name.strip().replace("-", "_")
return name, _IMPORT_ALIASES.get(name, name)
def _venv_import_ok(module_name):
try:
out = subprocess.run(
[venv_python(), "-c", "import %s" % module_name],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
timeout=180,
env=clean_env(),
)
return out.returncode == 0
except (OSError, subprocess.TimeoutExpired):
return False
def _filter_extras(extras):
"""Drop wheel URLs whose cpXY tag doesn't match the venv python."""
pyver = venv_python_version()
cp_tag = "cp" + pyver.replace(".", "") if pyver else None
kept = []
for pkg in extras:
if "cp3" in pkg and cp_tag and cp_tag not in pkg:
warn(
"Skipping %s (built for a different python than venv %s)."
% (os.path.basename(pkg), pyver)
)
continue
kept.append(pkg)
return kept
def ensure_requirements(spec, dry_run=False, force=False):
if not force and requirements_in_sync(spec):
ok("Requirements already in sync.")
return False
# force git-pinned deps (diffusers) onto a newly pinned commit — pip won't
# reinstall them on its own because the version number stays the same
stale = _stale_git_pins(spec)
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)
find_links = list(pins)
for url in spec.find_links:
find_links += ["--find-links", url]
info("Installing requirements from %s..." % spec.requirements_file)
# find_links included: on Spark the requirements themselves resolve
# self-built wheels (opencv, soxr, ...) from the spark wheel set
_pip_install(["-r", spec.requirements_path()] + find_links, dry_run=dry_run)
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.
# No --upgrade here: every optional spec is an exact pin or wheel URL, and
# uv's -U eagerly upgrades the whole dependency closure, blowing past
# requirements.txt pins (numpy/transformers) that only the requirements
# pass enforces.
for pkg in _filter_extras(spec.optional_packages):
label = os.path.basename(pkg) if "://" in pkg else pkg
info("Installing optional accelerator: %s" % label)
code = _pip_install([pkg] + find_links, dry_run=dry_run, check=False)
if code != 0:
warn("Optional package failed to install (continuing): %s" % label)
continue
if dry_run:
continue
# prebuilt accelerator wheels are sometimes built against a torch
# nightly and fail to load against the release ABI — verify the
# import and roll back rather than leaving a broken wheel installed
dist_name, import_name = _optional_names(pkg)
if not _venv_import_ok(import_name):
warn(
"%s installed but fails to import against this torch build — "
"removing it (training falls back to native attention)."
% import_name
)
_pip_uninstall([dist_name] + _ROLLBACK_EXTRAS.get(dist_name, []))
# packages whose dependency metadata is unsatisfiable on this platform but
# which work fine without it (e.g. tensorboard's grpcio on win_arm64)
for pkg in spec.no_deps_packages:
info("Installing (no-deps): %s" % pkg)
code = _pip_install_no_deps(pkg, dry_run=dry_run)
if code != 0:
warn("No-deps package failed to install (continuing): %s" % pkg)
_verify_torch(spec, dry_run=dry_run)
if not dry_run:
state = load_state()
state["req_hash"] = requirements_hash(spec)
state["backend"] = spec.backend
state["git_pins"] = _git_pinned_packages(spec)
save_state(state)
return True
# ---------------------------------------------------------------- sitecustomize
def _msvc_runtime_env():
"""{env: value} + [bin dirs] from vcvarsarm64, for triton's runtime JIT.
Triton compiles its kernel launcher stubs with cl.exe at runtime (cached
afterwards in ~/.triton), which needs INCLUDE/LIB and cl on PATH. Capture
the values once at sync time and bake them into sitecustomize.
"""
vcvars = (
r"C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools"
r"\VC\Auxiliary\Build\vcvarsarm64.bat"
)
if not os.path.isfile(vcvars):
return {}, []
try:
# string form: list2cmdline would mangle the nested quoting
out = subprocess.run(
'cmd /s /c "call "%s" >nul && set"' % vcvars,
stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL,
timeout=120,
)
if out.returncode != 0:
return {}, []
env = {}
for line in out.stdout.decode(errors="replace").splitlines():
if "=" in line:
k, _, v = line.partition("=")
env[k.upper()] = v
bin_dirs = [
d for d in env.get("PATH", "").split(os.pathsep)
if os.path.isfile(os.path.join(d, "cl.exe"))
][:1]
keep = {k: env[k] for k in ("INCLUDE", "LIB") if env.get(k)}
if keep and bin_dirs:
keep["CC"] = "cl"
return keep, bin_dirs
except (OSError, subprocess.TimeoutExpired):
pass
return {}, []
def write_sitecustomize(dry_run=False, spec=None):
"""Drop a sitecustomize.py into the venv exposing runtime DLL dirs.
sitecustomize is imported automatically at interpreter startup, so ANY use
of the venv python (UI-spawned training jobs, run.py from a terminal) gets
.ffmpeg/bin on PATH — and on Windows, os.add_dll_directory so torchcodec
finds the FFmpeg DLLs. On Spark the spec also carries the CUDA/cuDNN/APL
bin dirs, because the native torch wheel does not bundle its DLLs.
"""
from . import ffmpeg
if not venv_exists():
return
try:
out = subprocess.run(
[
venv_python(),
"-c",
"import sysconfig; print(sysconfig.get_paths()['purelib'])",
],
stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL,
timeout=30,
env=clean_env(),
)
site_packages = out.stdout.decode().strip()
except (OSError, subprocess.TimeoutExpired):
site_packages = ""
if not site_packages or not os.path.isdir(site_packages):
warn("Could not locate venv site-packages — skipping sitecustomize.")
return
target = os.path.join(site_packages, "sitecustomize.py")
dll_dirs = [ffmpeg.bin_dir()] + list(getattr(spec, "runtime_dll_dirs", []) or [])
runtime_env = dict(getattr(spec, "runtime_env", {}) or {})
if getattr(spec, "backend", None) == "cu134":
# triton's runtime launcher JIT needs the MSVC environment
msvc_env, msvc_bins = _msvc_runtime_env()
runtime_env.update(msvc_env)
dll_dirs += msvc_bins
content = (
"# Generated by the AI Toolkit manager (manager/env.py). Do not edit;\n"
"# regenerated on every `manager sync`.\n"
"import os\n"
"for _k, _v in %r.items():\n"
" os.environ.setdefault(_k, _v)\n"
"_DLL_DIRS = %r\n"
"_FFMPEG_LIB = %r\n"
"for _d in _DLL_DIRS:\n"
" if os.path.isdir(_d):\n"
" os.environ['PATH'] = _d + os.pathsep + os.environ.get('PATH', '')\n"
" if hasattr(os, 'add_dll_directory'):\n"
" try:\n"
" os.add_dll_directory(_d)\n"
" except OSError:\n"
" pass\n"
"if os.path.isdir(_FFMPEG_LIB):\n"
" # inherited by child processes (the ffmpeg/ffprobe executables\n"
" # need it to find their own shared libs)\n"
" _prev = os.environ.get('LD_LIBRARY_PATH', '')\n"
" if _FFMPEG_LIB not in _prev.split(os.pathsep):\n"
" os.environ['LD_LIBRARY_PATH'] = (\n"
" _FFMPEG_LIB + ((os.pathsep + _prev) if _prev else '')\n"
" )\n"
) % (runtime_env, dll_dirs, ffmpeg.lib_dir())
if dry_run:
info("[dry-run] would write %s" % target)
return
with open(target, "w") as f:
f.write(content)
# ---------------------------------------------------------------- sync
def sync(spec, detection, dry_run=False, force=False):
"""Bring the environment fully up to date for this checkout."""
from . import ffmpeg, gitwin, migrations, nodejs, uvbin
for note in spec.notes:
warn(note)
uvbin.ensure_uv(dry_run=dry_run)
gitwin.ensure_git(dry_run=dry_run)
if spec.backend == "cu134":
# native Spark stack: provision CUDA/cuDNN/APL runtime DLLs, VC
# redist, and (best-effort) MSVC for triton's kernel launcher JIT
from . import sparkdeps
sparkdeps.ensure_spark_runtime(dry_run=dry_run)
ensure_venv(spec, dry_run=dry_run)
# sitecustomize must exist BEFORE any torch import check below: on Spark
# the torch wheel is unbundled and only imports once the CUDA/cuDNN/BLAS
# DLL dirs from the spec are exposed to the interpreter
write_sitecustomize(dry_run=dry_run, spec=spec)
changed_torch = ensure_torch(spec, dry_run=dry_run)
# a torch reinstall can clobber pinned deps; force req pass afterwards
ensure_requirements(spec, dry_run=dry_run, force=force or changed_torch)
ffmpeg.ensure_ffmpeg(detection, dry_run=dry_run, spec=spec)
nodejs.ensure_node(detection, dry_run=dry_run)
nodejs.ensure_ui_deps(dry_run=dry_run)
write_sitecustomize(dry_run=dry_run, spec=spec)
migrations.run_pending(dry_run=dry_run)
ok("Environment is up to date.")

179
manager/ffmpeg.py Normal file
View File

@ -0,0 +1,179 @@
"""Local (never global) FFmpeg provisioning into <repo>/.ffmpeg/.
Why: the toolkit shells out to ffmpeg/ffprobe for video work, and torchcodec
dlopens the FFmpeg *shared libraries* at runtime. Installing FFmpeg system-wide
(apt/winget/brew) is exactly what we want to avoid, so we download a portable
build next to the repo:
- Linux / Windows: BtbN shared builds (bin/ + lib/ with .so/.dll) — the shared
libs are what torchcodec needs. The FFmpeg major version must be one the
pinned torchcodec supports.
- macOS: Martin Riedl static ffmpeg/ffprobe executables (no shared libs
published; torchcodec on mac keeps whatever it uses today).
Exposure to the rest of the system:
- `manager launch` prepends .ffmpeg/bin to PATH (and .ffmpeg/lib to
LD_LIBRARY_PATH on Linux) so the UI and every training job it spawns see it.
- env.py writes a sitecustomize.py into the venv that prepends .ffmpeg/bin to
PATH and (on Windows) calls os.add_dll_directory — so torchcodec finds the
DLLs in ANY use of the venv python, not just via `manager launch`.
"""
import os
import shutil
import stat
import tempfile
from .util import (
IS_MAC,
IS_WINDOWS,
REPO_ROOT,
download,
extract_archive,
info,
ok,
warn,
)
FFMPEG_DIR = os.path.join(REPO_ROOT, ".ffmpeg")
# FFmpeg 8 on all BtbN platforms — torchcodec 0.15 (pinned in spec.py)
# supports ffmpeg up to 8. Bump these together with the torchcodec pin.
_BTBN = "https://github.com/BtbN/FFmpeg-Builds/releases/download/latest/"
_RIEDL = (
"https://ffmpeg.martin-riedl.de/redirect/latest/macos/{arch}/release/{tool}.zip"
)
_SOURCES = {
("linux", "x86_64"): _BTBN + "ffmpeg-n8.1-latest-linux64-gpl-shared-8.1.tar.xz",
("linux", "aarch64"): _BTBN + "ffmpeg-n8.1-latest-linuxarm64-gpl-shared-8.1.tar.xz",
("windows", "x86_64"): _BTBN + "ffmpeg-n8.1-latest-win64-gpl-shared-8.1.zip",
# Windows-on-ARM in the emulated-x64 stack gets the x64 build, NOT BtbN's
# winarm64 one: an x64 torchcodec can only dlopen x64 FFmpeg DLLs, and the
# exes run fine under emulation.
("windows", "aarch64"): _BTBN + "ffmpeg-n8.1-latest-win64-gpl-shared-8.1.zip",
}
# Native Spark stack: the self-built win_arm64 torchcodec is linked against
# (and dlopens) arm64 FFmpeg 8 — LGPL to keep the distributed torchcodec wheel
# clean, matching C:\Dev spark build scripts / the wheel-set build recipe.
_SPARK_NATIVE_SOURCE = _BTBN + "ffmpeg-n8.1-latest-winarm64-lgpl-shared-8.1.zip"
def bin_dir():
return os.path.join(FFMPEG_DIR, "bin")
def lib_dir():
return os.path.join(FFMPEG_DIR, "lib")
def ffmpeg_exe():
return os.path.join(bin_dir(), "ffmpeg.exe" if IS_WINDOWS else "ffmpeg")
def is_installed(source_url):
marker = os.path.join(FFMPEG_DIR, ".source")
if not os.path.isfile(ffmpeg_exe()) or not os.path.isfile(marker):
return False
with open(marker) as f:
return f.read().strip() == source_url
def _mark_installed(source_url):
with open(os.path.join(FFMPEG_DIR, ".source"), "w") as f:
f.write(source_url)
def _install_btbn(url):
tmp = tempfile.mkdtemp(prefix="aitk_ffmpeg_")
try:
archive = os.path.join(tmp, os.path.basename(url))
download(url, archive, label="ffmpeg")
extract_archive(archive, tmp)
# archives contain a single top-level dir with bin/ lib/ include/
inner = None
for name in os.listdir(tmp):
path = os.path.join(tmp, name)
if os.path.isdir(path) and os.path.isdir(os.path.join(path, "bin")):
inner = path
break
if inner is None:
warn("Unexpected ffmpeg archive layout — skipping ffmpeg install.")
return False
if os.path.isdir(FFMPEG_DIR):
shutil.rmtree(FFMPEG_DIR)
shutil.move(inner, FFMPEG_DIR)
return True
finally:
shutil.rmtree(tmp, ignore_errors=True)
def _install_mac(detection):
arch = "arm64" if detection["arch"] == "arm64" else "amd64"
tmp = tempfile.mkdtemp(prefix="aitk_ffmpeg_")
try:
os.makedirs(bin_dir(), exist_ok=True)
for tool in ("ffmpeg", "ffprobe"):
url = _RIEDL.format(arch=arch, tool=tool)
archive = os.path.join(tmp, tool + ".zip")
download(url, archive, label=tool)
extract_archive(archive, tmp)
src = os.path.join(tmp, tool)
if not os.path.isfile(src):
warn("Unexpected %s archive layout — skipping." % tool)
return False
dest = os.path.join(bin_dir(), tool)
shutil.move(src, dest)
os.chmod(
dest, os.stat(dest).st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
)
return True
finally:
shutil.rmtree(tmp, ignore_errors=True)
def source_url(detection, spec=None):
if detection["os"] == "mac":
arch = "arm64" if detection["arch"] == "arm64" else "amd64"
return _RIEDL.format(arch=arch, tool="ffmpeg")
if (
getattr(spec, "backend", None) == "cu134"
and (detection["os"], detection["arch"]) == ("windows", "aarch64")
):
return _SPARK_NATIVE_SOURCE
return _SOURCES.get((detection["os"], detection["arch"]))
def ensure_ffmpeg(detection, dry_run=False, spec=None):
url = source_url(detection, spec=spec)
if url is None:
warn(
"No portable FFmpeg source for %s/%s — skipping local ffmpeg."
% (detection["os"], detection["arch"])
)
return False
if is_installed(url):
ok("Local FFmpeg already installed (.ffmpeg/).")
return False
if dry_run:
info("[dry-run] would install local FFmpeg from %s into %s" % (url, FFMPEG_DIR))
return False
installed = (
_install_mac(detection) if detection["os"] == "mac" else _install_btbn(url)
)
if installed:
_mark_installed(url)
ok("Local FFmpeg installed at %s" % FFMPEG_DIR)
return installed
def env_additions():
"""(path_dirs, ld_library_dirs) to prepend when launching anything."""
paths = []
lib_paths = []
if os.path.isdir(bin_dir()):
paths.append(bin_dir())
if not IS_WINDOWS and not IS_MAC and os.path.isdir(lib_dir()):
lib_paths.append(lib_dir())
return paths, lib_paths

92
manager/gitops.py Normal file
View File

@ -0,0 +1,92 @@
"""Git operations for update checking and pulling. Stdlib only."""
import os
from . import gitwin
from .util import run, REPO_ROOT, die
def _git(args, capture=True, check=True):
git = gitwin.find_git()
if not git:
die(
"git was not found. On Windows run `python -m manager sync` to "
"install a local copy; elsewhere install git with your package "
"manager (or xcode-select --install on macOS)."
)
# skip LFS payloads — the repo may carry LFS files not needed at runtime
env = os.environ.copy()
env["GIT_LFS_SKIP_SMUDGE"] = "1"
return run([git] + args, cwd=REPO_ROOT, capture=capture, check=check, env=env)
def current_branch():
_, out = _git(["rev-parse", "--abbrev-ref", "HEAD"])
return out
def current_commit():
_, out = _git(["rev-parse", "--short", "HEAD"])
return out
def is_dirty():
_, out = _git(["status", "--porcelain"])
# untracked files don't block updates; modified/staged tracked files do
for line in (out or "").splitlines():
if not line.startswith("??"):
return True
return False
def fetch():
code, _ = _git(["fetch", "--quiet"], capture=False, check=False)
return code == 0
def upstream():
code, out = _git(
["rev-parse", "--abbrev-ref", "--symbolic-full-name", "@{u}"], check=False
)
if code != 0:
return None
return out
def behind_count():
"""Commits the local branch is behind its upstream. None if no upstream."""
up = upstream()
if not up:
return None
code, out = _git(["rev-list", "--count", "HEAD..@{u}"], check=False)
if code != 0 or out is None:
return None
try:
return int(out)
except ValueError:
return None
def remote_commit():
code, out = _git(["rev-parse", "--short", "@{u}"], check=False)
return out if code == 0 else None
def incoming_log(limit=15):
code, out = _git(
["log", "--oneline", "HEAD..@{u}", "--max-count=%d" % limit], check=False
)
if code != 0 or not out:
return []
return out.splitlines()
def pull_ff():
"""Fast-forward pull. Dies with guidance on failure."""
code, _ = _git(["pull", "--ff-only"], capture=False, check=False)
if code != 0:
die(
"git pull --ff-only failed. Your branch has local commits or "
"conflicts with the remote. Resolve manually (git stash / git rebase) "
"and re-run the update."
)

76
manager/gitwin.py Normal file
View File

@ -0,0 +1,76 @@
"""Local (never global) Git provisioning for Windows into <repo>/.mingit/.
MinGit is the official minimal, portable Git for Windows — a plain zip meant
for embedding in tools (no installer, no registry, no PATH changes). On
Windows, if the system has no git, `manager sync` drops one here so updates
keep working from any terminal. The first clone (before this repo exists) is
handled the same way by the bootstrap layer (install.ps1 / desktop launcher),
which then moves its MinGit into the fresh checkout as .mingit/.
Linux and macOS have no sane portable git (glibc / Xcode CLT entanglement),
and git is effectively always available there — so this is Windows-only and
the manager just errors with install instructions elsewhere.
"""
import os
import platform
import shutil
import tempfile
from .util import IS_WINDOWS, REPO_ROOT, download, extract_archive, info, ok, warn
MINGIT_DIR = os.path.join(REPO_ROOT, ".mingit")
# Update this pin together with nothing else — it's independent of torch etc.
_MINGIT_TAG = "v2.55.0.windows.3"
_MINGIT_VERSION = "2.55.0.3"
def _mingit_url():
# git is a standalone subprocess (nothing dlopens it), so unlike
# node/ffmpeg it can be native arm64 on Windows-on-ARM
arm = platform.machine().lower() in ("arm64", "aarch64")
flavor = "arm64" if arm else "64-bit"
return "https://github.com/git-for-windows/git/releases/download/%s/MinGit-%s-%s.zip" % (
_MINGIT_TAG,
_MINGIT_VERSION,
flavor,
)
def local_git_exe():
return os.path.join(MINGIT_DIR, "cmd", "git.exe")
def find_git():
"""Path/command for git: repo-local MinGit first (Windows), then system."""
if IS_WINDOWS and os.path.isfile(local_git_exe()):
return local_git_exe()
return shutil.which("git")
def ensure_git(dry_run=False):
"""Windows-only: provision .mingit/ when the system has no git."""
if not IS_WINDOWS:
return False
if find_git():
return False
if dry_run:
info("[dry-run] would install MinGit into %s" % MINGIT_DIR)
return False
tmp = tempfile.mkdtemp(prefix="aitk_mingit_")
try:
archive = os.path.join(tmp, "mingit.zip")
download(_mingit_url(), archive, label="MinGit")
# MinGit zips have no top-level folder: cmd/, mingw64/, etc at the root
extracted = os.path.join(tmp, "mingit")
extract_archive(archive, extracted)
if not os.path.isfile(os.path.join(extracted, "cmd", "git.exe")):
warn("Unexpected MinGit archive layout — skipping local git install.")
return False
if os.path.isdir(MINGIT_DIR):
shutil.rmtree(MINGIT_DIR)
shutil.move(extracted, MINGIT_DIR)
ok("Local Git (MinGit) installed at %s" % MINGIT_DIR)
return True
finally:
shutil.rmtree(tmp, ignore_errors=True)

114
manager/launch.py Normal file
View File

@ -0,0 +1,114 @@
"""Launch the AI Toolkit web UI (ui/ -> npm run db_build_start)."""
import os
import subprocess
import threading
from . import ffmpeg, nodejs
from .util import (
IS_LINUX,
IS_WINDOWS,
REPO_ROOT,
clean_env,
die,
info,
venv_dir,
venv_python,
)
UI_PORT = 8675
UI_URL = "http://localhost:%d" % UI_PORT
BROWSER_POLL_SECONDS = 300
def build_env():
"""Scrubbed env with local node, ffmpeg, and the venv on PATH.
Everything the UI worker spawns (training jobs) inherits this, so the
local ffmpeg/node are visible to the whole process tree.
"""
env = clean_env()
path_dirs = []
if os.path.isdir(nodejs.node_bin_dir()):
path_dirs.append(nodejs.node_bin_dir())
ff_paths, ff_libs = ffmpeg.env_additions()
path_dirs += ff_paths
vbin = os.path.join(venv_dir(), "Scripts" if IS_WINDOWS else "bin")
if os.path.isdir(vbin):
path_dirs.append(vbin)
if path_dirs:
env["PATH"] = os.pathsep.join(path_dirs) + os.pathsep + env.get("PATH", "")
if ff_libs:
env["LD_LIBRARY_PATH"] = os.pathsep.join(
ff_libs + [env.get("LD_LIBRARY_PATH", "")]
).rstrip(os.pathsep)
return env
def _headless():
return IS_LINUX and not (
os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY")
)
def _open_browser_when_ready(stop_event):
"""Poll the UI port in the background; open the browser once it responds."""
import time
import urllib.request
import webbrowser
waited = 0
while not stop_event.is_set() and waited < BROWSER_POLL_SECONDS:
try:
urllib.request.urlopen(UI_URL, timeout=2).close()
webbrowser.open(UI_URL)
return
except OSError:
time.sleep(2)
waited += 2
def launch_ui(open_browser=True):
if not os.path.isfile(venv_python()):
die("No Python environment found. Run: python3 -m manager install")
env = build_env()
npm = nodejs.find_npm(env)
if not npm:
die(
"Node.js was not found. Run `python3 -m manager sync` to install a "
"local copy, or install Node.js >= %d from https://nodejs.org."
% nodejs.MIN_NODE_MAJOR
)
_, major = nodejs.have_usable_node()
if major is not None and major < nodejs.MIN_NODE_MAJOR:
die(
"Node.js v%d found, but >= %d is required. Run `python3 -m manager sync`."
% (major, nodejs.MIN_NODE_MAJOR)
)
# deps are installed here rather than by the npm script so the lockfile is
# never rewritten (see nodejs.ensure_ui_deps); a no-op once they're in sync
nodejs.ensure_ui_deps(env=env)
info("Starting AI Toolkit UI (%s) ..." % UI_URL)
stop_event = threading.Event()
if open_browser and not _headless():
threading.Thread(
target=_open_browser_when_ready, args=(stop_event,), daemon=True
).start()
proc = subprocess.Popen(
[npm, "run", "db_build_start"], cwd=os.path.join(REPO_ROOT, "ui"), env=env
)
try:
return proc.wait()
except KeyboardInterrupt:
proc.terminate()
try:
proc.wait(timeout=10)
except subprocess.TimeoutExpired:
proc.kill()
return 130
finally:
stop_event.set()

34
manager/migrations.py Normal file
View File

@ -0,0 +1,34 @@
"""One-time migration steps that run after an update.
Add a migration when an update needs more than a dependency sync (moving
files, converting configs, clearing caches, etc). Each runs at most once per
environment; applied ids are recorded in the venv state file.
def _example(dry_run):
...
MIGRATIONS = [
{"id": "2026-07-example-cache-move", "run": _example},
]
"""
from .util import info
from . import env
MIGRATIONS = []
def run_pending(dry_run=False):
if not MIGRATIONS:
return
state = env.load_state()
applied = set(state.get("migrations", []))
for migration in MIGRATIONS:
if migration["id"] in applied:
continue
info("Running migration: %s" % migration["id"])
if not dry_run:
migration["run"](dry_run)
applied.add(migration["id"])
state["migrations"] = sorted(applied)
env.save_state(state)

245
manager/nodejs.py Normal file
View File

@ -0,0 +1,245 @@
"""Local (never global) Node.js provisioning into <repo>/.node/.
A system Node >= 20 is used when present; otherwise an official portable
build is downloaded next to the repo (same approach run_mac.zsh already uses).
Nothing is ever installed system-wide.
"""
import os
import shutil
import subprocess
import tempfile
from .util import (
IS_MAC,
IS_WINDOWS,
REPO_ROOT,
clean_env,
download,
extract_archive,
file_hash,
info,
ok,
run,
warn,
which,
)
NODE_DIR = os.path.join(REPO_ROOT, ".node")
# Node 24 is the current LTS line and matches the dgx_instructions.md guidance.
NODE_VERSION = "24.11.1"
MIN_NODE_MAJOR = 20
UI_DIR = os.path.join(REPO_ROOT, "ui")
UI_STATE_KEY = "ui_deps_hash"
def node_bin_dir():
# windows zips have node.exe/npm.cmd at the archive root; unix under bin/
return NODE_DIR if IS_WINDOWS else os.path.join(NODE_DIR, "bin")
def local_node_exe():
return os.path.join(node_bin_dir(), "node.exe" if IS_WINDOWS else "node")
def _node_info(exe):
"""(major, arch) for a node executable, e.g. (24, 'x64'), or (None, None)."""
try:
out = subprocess.run(
[exe, "-p", "process.version + ' ' + process.arch"],
stdout=subprocess.PIPE,
stderr=subprocess.DEVNULL,
timeout=15,
)
if out.returncode != 0:
return None, None
version, arch = out.stdout.decode().strip().split() # v24.11.1 x64
return int(version.lstrip("v").split(".")[0]), arch
except (OSError, subprocess.TimeoutExpired, ValueError):
return None, None
def _usable_major(exe):
"""Version if this node can run the UI, else None.
On Windows only an x64 node qualifies — the UI's native modules resolve
for the node process arch and Prisma ships no windows-arm64 query engine,
so a native ARM64 node loads the UI but dies on the first DB call. The
x64 node runs under emulation on ARM hosts.
"""
major, arch = _node_info(exe)
if major is None or major < MIN_NODE_MAJOR:
return None
if IS_WINDOWS and arch != "x64":
return None
return major
def _dist_url(detection):
arch = detection["arch"]
if detection["os"] == "linux":
plat = "linux-arm64" if arch == "aarch64" else "linux-x64"
ext = "tar.xz"
elif detection["os"] == "mac":
plat = "darwin-arm64" if arch == "arm64" else "darwin-x64"
ext = "tar.gz"
elif detection["os"] == "windows":
# win-x64 even on ARM hosts: Prisma has no windows-arm64 engine, so an
# arm64 node can't run the UI (see _usable_major); x64 node emulates fine
plat = "win-x64"
ext = "zip"
else:
return None, None
name = "node-v%s-%s" % (NODE_VERSION, plat)
return "https://nodejs.org/dist/v%s/%s.%s" % (NODE_VERSION, name, ext), name
def have_usable_node():
"""(exe, major) for the best available node: local .node/ first, then system."""
local = local_node_exe()
if os.path.isfile(local):
major = _usable_major(local)
if major is not None:
return local, major
system = which("node")
if system:
major = _usable_major(system)
if major is not None:
return system, major
return None, None
def ensure_node(detection, dry_run=False):
exe, major = have_usable_node()
if exe:
ok("Node.js v%d found (%s)." % (major, exe))
return False
system = which("node")
if system:
_, arch = _node_info(system)
if IS_WINDOWS and arch and arch != "x64":
info(
"System Node.js is %s but the UI needs an x64 node on Windows "
"(Prisma has no windows-arm64 engine) — installing a local x64 "
"copy." % arch
)
url, inner_name = _dist_url(detection)
if url is None:
warn(
"No portable Node.js build for this platform — install Node >= %d manually."
% MIN_NODE_MAJOR
)
return False
if dry_run:
info(
"[dry-run] would install portable Node.js v%s into %s"
% (NODE_VERSION, NODE_DIR)
)
return False
tmp = tempfile.mkdtemp(prefix="aitk_node_")
try:
archive = os.path.join(tmp, os.path.basename(url))
download(url, archive, label="node v%s" % NODE_VERSION)
extract_archive(archive, tmp)
inner = os.path.join(tmp, inner_name)
if not os.path.isdir(inner):
warn("Unexpected Node.js archive layout — skipping node install.")
return False
if os.path.isdir(NODE_DIR):
shutil.rmtree(NODE_DIR)
shutil.move(inner, NODE_DIR)
ok("Portable Node.js v%s installed at %s" % (NODE_VERSION, NODE_DIR))
return True
finally:
shutil.rmtree(tmp, ignore_errors=True)
# ---------------------------------------------------------------- ui deps
def npm_env():
"""Scrubbed env with the local .node/ ahead of PATH."""
env = clean_env()
if os.path.isdir(node_bin_dir()):
env["PATH"] = node_bin_dir() + os.pathsep + env.get("PATH", "")
return env
def find_npm(env=None):
"""npm from the local .node/ copy when we installed one, else the system."""
path = (env or npm_env()).get("PATH")
return shutil.which("npm.cmd" if IS_WINDOWS else "npm", path=path)
def _ui_deps_hash():
return file_hash(
[
os.path.join(UI_DIR, "package-lock.json"),
os.path.join(UI_DIR, "package.json"),
]
)
def ensure_ui_deps(env=None, dry_run=False):
"""Install ui/node_modules without ever rewriting ui/package-lock.json.
A plain `npm install` treats the lockfile as writable and re-derives it for
the machine doing the install: on Windows it strips the `libc` fields off
the Linux-only optional binaries (@next/swc-linux-*, rollup, lightningcss),
on Linux it puts them back. Since the UI installs on every launch, everyone
ends up with a modified lockfile they never asked for — which then blocks
`manager update`, because a dirty tree aborts the pull.
`--no-save` keeps the resolution but stops the write-back, and stays an
incremental install (seconds) rather than the full wipe-and-refetch of
`npm ci`. The lockfile is snapshotted and restored regardless, so this
holds even if a future npm changes what `--no-save` covers.
Gated on a hash of the manifests so steady-state launches skip npm.
"""
from . import env as env_mod
lock = os.path.join(UI_DIR, "package-lock.json")
if not os.path.isfile(lock):
warn("ui/package-lock.json is missing — skipping UI dependency install.")
return False
want = _ui_deps_hash()
if env_mod.venv_exists() and env_mod.load_state().get(UI_STATE_KEY) == want:
if os.path.isdir(os.path.join(UI_DIR, "node_modules")):
ok("UI dependencies already installed.")
return False
if dry_run:
info("[dry-run] would run: npm install --no-save (in %s)" % UI_DIR)
return False
env = env or npm_env()
npm = find_npm(env)
if npm is None:
warn("npm not found — skipping UI dependency install.")
return False
info("Installing UI dependencies...")
with open(lock, "rb") as f:
before = f.read()
code, _ = run(
[npm, "install", "--no-save", "--no-audit", "--no-fund"],
cwd=UI_DIR,
env=env,
stream=True,
check=False,
)
with open(lock, "rb") as f:
after = f.read()
if after != before:
with open(lock, "wb") as f:
f.write(before)
info("Reverted npm's edit to ui/package-lock.json (managed by the repo).")
if code != 0:
warn("UI dependency install failed — the UI may not start.")
return False
if env_mod.venv_exists():
state = env_mod.load_state()
state[UI_STATE_KEY] = want
env_mod.save_state(state)
ok("UI dependencies ready.")
return True

224
manager/sparkdeps.py Normal file
View File

@ -0,0 +1,224 @@
"""RTX Spark native-stack runtime provisioning (Windows on ARM, cu134).
Goal: a fresh Spark machine runs run_windows.bat and gets as close to
zero-manual-setup as licensing allows. The native wheels (torch etc.) do not
bundle CUDA / cuDNN / BLAS DLLs, and triton's launcher JIT wants MSVC. Policy:
NVIDIA components are never redistributed by us.
- CUDA 13.4 toolkit (developer preview): MANUAL install — the preview EULA
requires NVIDIA's own click-through, so the manager only detects it and
prints instructions when missing. This is the single manual step.
- cuDNN (arm64): auto-downloaded from NVIDIA's own official installer URL and
installed silently — fetched directly from NVIDIA, not redistributed.
- Arm Performance Libraries: auto-install via winget (official Arm package).
- MSVC Build Tools (triton torch.compile JIT only): auto-install via winget;
failure downgrades gracefully (training works, no torch.compile).
- VC redistributable (arm64): auto-install via winget when msvcp140 missing.
Everything is best-effort with warnings; the training stack itself only hard-
requires the CUDA + cuDNN + APL DLL dirs.
"""
import glob
import os
import subprocess
from .util import download, info, ok, warn, which
CUDA_DOWNLOAD_PAGE = (
"https://developer.nvidia.com/cuda-13-4-0-download-archive"
"?target_os=Windows&target_arch=arm64"
)
# NVIDIA's official public installer for cuDNN on Windows arm64. Downloaded
# straight from NVIDIA at install time (we do not redistribute it). Update
# together with the wheel set when moving to a newer cuDNN.
CUDNN_INSTALLER_URL = (
"https://developer.download.nvidia.com/compute/cudnn/9.25.0/"
"local_installers/cudnn_9.25.0_windows_arm64.exe"
)
# System install roots, newest version preferred (globs, not pinned versions)
_CUDA_BIN_GLOB = r"C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v*\bin\arm64"
_CUDNN_BIN_GLOB = r"C:\Program Files\NVIDIA\CUDNN\v*\bin\*\arm64"
_ARMPL_BIN_GLOB = r"C:\Program Files\Arm Performance Libraries\armpl_*\bin"
_VS_BUILDTOOLS = (
r"C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools"
)
def _newest(pattern):
matches = sorted(glob.glob(pattern))
return matches[-1] if matches else None
def cuda_bin_dir():
return _newest(_CUDA_BIN_GLOB)
def cuda_root():
d = cuda_bin_dir()
# <root>\bin\arm64 -> <root>
return os.path.dirname(os.path.dirname(d)) if d else None
def cudnn_bin_dir():
return _newest(_CUDNN_BIN_GLOB)
def armpl_bin_dir():
return _newest(_ARMPL_BIN_GLOB)
def resolve_dll_dirs():
"""All runtime DLL dirs for the native stack (existing ones only)."""
return [d for d in (cuda_bin_dir(), cudnn_bin_dir(), armpl_bin_dir()) if d]
def runtime_complete():
return bool(cuda_bin_dir() and cudnn_bin_dir() and armpl_bin_dir())
def triton_tool_env():
"""TRITON_*_PATH env for ptxas etc. from the system CUDA install.
Our triton wheel deliberately does NOT bundle NVIDIA's compiler tools
(developer-preview licensing); resolve them from the user's toolkit.
"""
root = cuda_root()
if not root:
return {}
env = {}
for var, exe in (
("TRITON_PTXAS_PATH", "ptxas.exe"),
("TRITON_PTXAS_BLACKWELL_PATH", "ptxas.exe"),
("TRITON_CUOBJDUMP_PATH", "cuobjdump.exe"),
("TRITON_NVDISASM_PATH", "nvdisasm.exe"),
):
path = os.path.join(root, "bin", exe)
if os.path.isfile(path):
env[var] = path
return env
def check_cuda():
"""CUDA toolkit is the one manual install (preview EULA). Detect + guide."""
if cuda_bin_dir():
return True
warn(
"The CUDA 13.4 toolkit (arm64) is not installed. NVIDIA's developer "
"preview license requires installing it manually:\n"
" 1. Download from %s\n"
" 2. Install with default settings, then re-run this setup.\n"
"The RTX Spark developer driver (R616+) is required as well."
% CUDA_DOWNLOAD_PAGE
)
return False
def ensure_cudnn(dry_run=False):
"""Fetch + silently run NVIDIA's official cuDNN installer if missing."""
if cudnn_bin_dir():
return True
if dry_run:
info("[dry-run] would download and install cuDNN from NVIDIA")
return False
import tempfile
tmp = tempfile.mkdtemp(prefix="aitk_cudnn_")
try:
exe = os.path.join(tmp, os.path.basename(CUDNN_INSTALLER_URL))
download(CUDNN_INSTALLER_URL, exe, label="cuDNN (from NVIDIA)")
info("Installing cuDNN (silent)...")
code = subprocess.call([exe, "-s"])
if code != 0:
warn("cuDNN installer exited with %d." % code)
return cudnn_bin_dir() is not None
finally:
import shutil
shutil.rmtree(tmp, ignore_errors=True)
def have_msvc():
return bool(
glob.glob(os.path.join(_VS_BUILDTOOLS, "VC", "Tools", "MSVC", "*",
"bin", "Hostarm64", "arm64", "cl.exe"))
)
def _winget_install(args, label, dry_run=False):
winget = which("winget")
if not winget:
warn("winget not available — cannot auto-install %s." % label)
return False
if dry_run:
info("[dry-run] would winget install %s" % label)
return False
info("Installing %s (one-time, may take several minutes)..." % label)
code = subprocess.call(
[winget, "install", "--exact", "--source", "winget",
"--accept-source-agreements", "--accept-package-agreements"] + args,
stdout=subprocess.DEVNULL,
)
if code != 0:
warn("%s install failed (winget exit %d)." % (label, code))
return code == 0
def ensure_armpl(dry_run=False):
if armpl_bin_dir():
return True
return _winget_install(
["--id", "Arm.ArmPerformanceLibraries"],
"Arm Performance Libraries",
dry_run=dry_run,
)
def ensure_vcredist(dry_run=False):
"""VC runtime (msvcp140 etc.) — required by the native wheels."""
sysdir = os.path.join(os.environ.get("SystemRoot", r"C:\Windows"), "System32")
if os.path.isfile(os.path.join(sysdir, "msvcp140.dll")):
return True
return _winget_install(
["--id", "Microsoft.VCRedist.2015+.arm64"],
"Visual C++ Redistributable (arm64)",
dry_run=dry_run,
)
def ensure_msvc(dry_run=False):
"""MSVC Build Tools — only needed for triton's runtime kernel launchers.
Best effort: without it, training still works; torch.compile / triton
JIT is unavailable until the user installs Build Tools.
"""
if have_msvc():
return True
done = _winget_install(
["--id", "Microsoft.VisualStudio.2022.BuildTools", "--override",
"--quiet --wait --norestart "
"--add Microsoft.VisualStudio.Workload.VCTools "
"--add Microsoft.VisualStudio.Component.VC.Tools.ARM64 "
"--add Microsoft.VisualStudio.Component.Windows11SDK.26100"],
"MSVC Build Tools (for torch.compile/triton)",
dry_run=dry_run,
)
if not done and not dry_run:
warn(
"torch.compile/triton kernel JIT will be unavailable until MSVC "
"Build Tools are installed; training itself is unaffected."
)
return done
def ensure_spark_runtime(dry_run=False):
"""Full best-effort provisioning for the native Spark stack."""
check_cuda()
ensure_vcredist(dry_run=dry_run)
ensure_cudnn(dry_run=dry_run)
ensure_armpl(dry_run=dry_run)
ensure_msvc(dry_run=dry_run)
if runtime_complete():
ok("Spark native runtime present (CUDA + cuDNN + Arm PL).")

482
manager/spec.py Normal file
View File

@ -0,0 +1,482 @@
"""Maps detected hardware to an environment spec.
This is the single source of truth for "what does this machine need to run
this commit of AI Toolkit". One universal torch version across all platforms;
per-platform accelerator extras (flash-attn, NATTEN, triton) wherever prebuilt
wheels exist. **Update the pins below together with the README install
instructions and run_mac.zsh.**
Wheel coverage for the pinned set (verified 2026-07, flash-attn + NATTEN GPU
kernels smoke-tested on an RTX 5090 / sm120 with torch 2.13.0+cu130):
- torch 2.13.0: cu126/cu130 wheels for linux x86_64 + aarch64 + windows; PyPI
wheels for mac arm64. torchaudio is in maintenance mode — 2.11.0 is the
current release and is torch-version-agnostic (no torch dep in metadata).
- torchcodec 0.15.0: supports torch >= 2.11, wheels on all platforms.
- flash-attn 2.8.3: prebuilt by mjun0812/flash-attention-prebuild-wheels for
{cu126,cu130} x {cp310..cp314} x {linux x86_64, linux aarch64, windows}.
NOTE: the torch2.12 linux wheels there were built against a torch nightly
and fail to import on 2.12.0 final — the torch2.13 batches (v0.9.47+) are
verified good. Re-verify imports whenever bumping torch.
- NATTEN 0.21.7: prebuilt at whl.natten.org for {cu126,cu130,cu132} x
{cp310..cp314} x {linux x86_64, linux aarch64}. No Windows/mac wheels.
- 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.
- flash-linear-attention 0.5.2: pure-Python (py3-none-any) Triton kernels —
installs anywhere, needs triton>=3.3 at runtime. Installed bare (no backend
extra) so it never pulls its own torch/triton over our pinned stack; usable
on every platform with a working triton (CUDA linux/windows, Spark, ROCm),
not on MPS/CPU.
- Windows-on-ARM (verified 2026-07): there are NO win_arm64 wheels for any of
the CUDA stack — torch cu130, triton-windows, flash-attn and Prisma's node
engine are all x64-only. The supported configuration is therefore the x64
stack end to end (Python, torch, Node) running under Windows' x64 emulation,
with the GPU driven natively by the NVIDIA driver. build_spec() pins the
interpreter arch explicitly so uv can never flip the venv to native aarch64
(where torch would not resolve). Revisit if pytorch.org ever publishes
win_arm64 CUDA wheels.
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
without). Wheel URLs containing a cpXY tag are skipped automatically if the
venv python doesn't match.
"""
import os
from .util import REPO_ROOT
# ---- version pins (edit these to move the fleet forward) -------------------
TORCH = {"torch": "2.13.0", "torchvision": "0.28.0", "torchaudio": "2.11.0"}
TORCH_TAG = "2.13" # as it appears in flash-attn / natten wheel names
TORCHCODEC = "torchcodec==0.15.0"
TRITON_WINDOWS = "triton-windows>=3.7,<3.8"
NATTEN_VERSION = "0.21.7"
NATTEN_FIND_LINKS = "https://whl.natten.org"
# pure-Python triton kernels; bare install (no [cuda]/[rocm] extra) on purpose —
# the extras only add torch/triton pins we already manage per-platform
FLA = "flash-linear-attention==0.5.2"
FLASH_ATTN_VERSION = "2.8.3"
_FA_BASE = (
"https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/"
)
# (os, arch) -> (release tag, wheel platform tag) — tags are per torch
# version; these carry the torch2.13 builds
_FA_BUILDS = {
("linux", "x86_64"): ("v0.9.47", "manylinux_2_24_x86_64.manylinux_2_28_x86_64"),
("linux", "aarch64"): ("v0.9.48", "manylinux_2_34_aarch64"),
("windows", "x86_64"): ("v0.9.52", "win_amd64"),
}
# helper build tools some sdists need on Windows
_WIN_HELPERS = ["wheel", "setuptools", "poetry-core", "hf_xet"]
PYTORCH_INDEX = "https://download.pytorch.org/whl/"
# ---- NVIDIA RTX Spark: native Windows-on-ARM CUDA ---------------------------
# The CUDA 13.4 developer preview added native win_arm64 CUDA. No public wheels
# exist for the GPU stack, so we build them ourselves (torch from main +
# pytorch/pytorch#190448, plus torchvision/torchaudio/torchcodec and the deps
# with no win_arm64 wheels). The manager installs them from a find-links
# source: the local wheels/spark/ dir during development, or the hosted URL
# once published. Without that source (or with a pre-13.4 driver) Spark
# machines fall back to the emulated x64 stack below, which also works.
SPARK_BACKEND = "cu134"
SPARK_TORCH = {
"torch": "2.14.0.dev20260727",
"torchvision": "0.29.0.dev20260727",
"torchaudio": "2.11.0.dev20260727",
}
SPARK_WHEELS_DIR = os.path.join(REPO_ROOT, "wheels", "spark")
# GitHub's expanded_assets endpoint serves plain HTML anchors — a valid
# pip/uv find-links page pointing at the release assets.
SPARK_WHEELS_URL = (
"https://github.com/ostris/ai-toolkit-spark-wheels/releases/"
"expanded_assets/cu134-20260727"
)
SPARK_UV_PYTHON = "cpython-3.12-windows-aarch64-none"
# Runtime DLL homes for the unbundled native torch (TH_BINARY_BUILD=0) are
# resolved dynamically (system installs of any version, else the downloadable
# runtime bundle) — see sparkdeps.resolve_dll_dirs().
def _spark_wheels_source():
"""find-links source holding the self-built win_arm64 wheels, or None."""
if os.path.isdir(SPARK_WHEELS_DIR):
for name in os.listdir(SPARK_WHEELS_DIR):
if name.startswith("torch-") and "win_arm64" in name:
return SPARK_WHEELS_DIR
return SPARK_WHEELS_URL
def _spark_capable(detection):
"""Driver new enough for native win_arm64 CUDA (R616+ reports CUDA 13.4)."""
nvidia = detection.get("nvidia") or {}
try:
cuda = tuple(int(x) for x in (nvidia.get("cuda_version") or "").split("."))
except ValueError:
return False
return cuda >= (13, 4)
class EnvSpec(object):
def __init__(
self,
backend,
torch_packages,
torch_index=None,
python_version="3.12",
requirements_file="requirements.txt",
extra_packages=None,
optional_packages=None,
find_links=None,
notes=None,
uv_python=None,
torch_links=None,
no_deps_packages=None,
runtime_dll_dirs=None,
):
self.backend = backend # cu134 / cu130 / cu126 / rocm7.1 / mps / cpu
self.torch_packages = torch_packages # {name: version}
self.torch_index = torch_index # None = PyPI
self.python_version = python_version
self.requirements_file = requirements_file
self.extra_packages = extra_packages or []
self.optional_packages = optional_packages or []
self.find_links = find_links or []
self.notes = notes or []
# full uv interpreter request (e.g. "cpython-3.12-windows-x86_64-none")
# when the venv arch must not be left to uv's default; None = just
# python_version
self.uv_python = uv_python
# --find-links sources for the torch trio itself (self-built wheels);
# used when torch_index is None
self.torch_links = torch_links or []
# installed with --no-deps after everything else (e.g. tensorboard on
# Spark, whose grpcio dep has no win_arm64 wheels but is only needed
# for the server, not the log writer)
self.no_deps_packages = no_deps_packages or []
# extra DLL dirs the venv needs at runtime (unbundled CUDA/cuDNN/BLAS
# on Spark); baked into sitecustomize.py, missing dirs skipped
self.runtime_dll_dirs = runtime_dll_dirs or []
# env vars every venv python needs (sitecustomize setdefault)
self.runtime_env = {}
def torch_args(self):
args = ["%s==%s" % (k, v) for k, v in sorted(self.torch_packages.items())]
if self.torch_index:
args += ["--index-url", self.torch_index]
for links in self.torch_links:
args += ["--find-links", links]
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)
def as_dict(self):
return {
"backend": self.backend,
"torch_packages": self.torch_packages,
"torch_index": self.torch_index,
"python_version": self.python_version,
"requirements_file": self.requirements_file,
"extra_packages": self.extra_packages,
"optional_packages": self.optional_packages,
"find_links": self.find_links,
"notes": self.notes,
"uv_python": self.uv_python,
"torch_links": self.torch_links,
"no_deps_packages": self.no_deps_packages,
"runtime_dll_dirs": self.runtime_dll_dirs,
"runtime_env": self.runtime_env,
}
def _flash_attn_url(flavor, os_name, arch, python_version):
build = _FA_BUILDS.get((os_name, arch))
if build is None:
return None
tag, plat = build
cp = "cp" + python_version.replace(".", "")
return "%s%s/flash_attn-%s+%storch%s-%s-%s-%s.whl" % (
_FA_BASE,
tag,
FLASH_ATTN_VERSION,
flavor,
TORCH_TAG,
cp,
cp,
plat,
)
def _natten_pin(flavor):
# natten wheel local tags use the full torch version without dots: torch2120cu130
return "natten==%s+torch%s%s" % (
NATTEN_VERSION,
TORCH["torch"].replace(".", ""),
flavor.replace(".", ""),
)
def _cuda_flavor(detection):
"""Pick a cuda wheel flavor the installed driver can actually run."""
nvidia = detection.get("nvidia") or {}
cuda = None
if nvidia.get("cuda_version"):
try:
cuda = tuple(int(x) for x in nvidia["cuda_version"].split("."))
except ValueError:
cuda = None
if cuda is None:
# driver present but version unknown — assume current
return "cu130", []
if cuda >= (13, 0):
return "cu130", []
# non-GPU rows (the NPU on ARM hybrids) report compute_cap as "[N/A]"
caps = []
for g in nvidia.get("gpus", []):
try:
caps.append(float(g.get("compute_cap")))
except (TypeError, ValueError):
pass
has_blackwell = any(c >= 12.0 for c in caps)
if cuda >= (12, 6):
if has_blackwell:
raise RuntimeError(
"Blackwell GPU detected but the NVIDIA driver only supports "
"CUDA %s. Blackwell needs the cu130 build — update your "
"driver to 580+ and re-run." % nvidia["cuda_version"]
)
return "cu126", [
"NVIDIA driver only supports CUDA %s — installing cu126 wheels. "
"Updating your driver is recommended." % nvidia["cuda_version"]
]
raise RuntimeError(
"NVIDIA driver only supports CUDA %s, which is too old for the pinned "
"torch build. Update your NVIDIA driver, then re-run install."
% nvidia["cuda_version"]
)
def _cuda_spec(detection):
os_name = detection["os"]
arch = detection["arch"]
flavor, notes = _cuda_flavor(detection)
python_version = "3.12"
requirements = (
"dgx_requirements.txt" if detection.get("is_dgx") else "requirements.txt"
)
if detection.get("is_dgx"):
# the old "Python 3.11 on DGX OS" constraint was for conda/system
# installs; uv provisions 3.12 and all aarch64 cp312 wheels exist now
notes = notes + [
"DGX OS / Grace detected: using %s wheels and dgx_requirements.txt."
% flavor
]
extras = [TORCHCODEC]
optional = [FLA]
find_links = []
# Windows-on-ARM runs the x64 wheel stack (see module docstring), so wheel
# selection uses x86_64 there regardless of the host arch.
wheel_arch = "x86_64" if os_name == "windows" else arch
fa_url = _flash_attn_url(flavor, os_name, wheel_arch, python_version)
if fa_url:
optional.append(fa_url)
if os_name == "linux":
optional.append(_natten_pin(flavor))
find_links.append(NATTEN_FIND_LINKS)
elif os_name == "windows":
extras = _WIN_HELPERS + extras + [TRITON_WINDOWS]
notes = notes + ["NATTEN has no Windows wheels — skipping it."]
return EnvSpec(
flavor,
TORCH,
torch_index=PYTORCH_INDEX + flavor,
python_version=python_version,
requirements_file=requirements,
extra_packages=extras,
optional_packages=optional,
find_links=find_links,
notes=notes,
)
def _spark_spec(detection, wheels_source):
"""Native win_arm64 CUDA 13.4 stack from self-built wheels (RTX Spark)."""
from . import sparkdeps
dll_dirs = sparkdeps.resolve_dll_dirs()
spec = _make_spark_spec(detection, wheels_source, dll_dirs)
# OpenCV's runtime CPU detection is blind to FP16/DOTPROD on Windows
# ARM64 and aborts the process at import even though the N1X supports
# both; the self-built cv2 wheels rely on this skip.
spec.runtime_env["OPENCV_SKIP_CPU_BASELINE_CHECK"] = "1"
# our triton wheel does not bundle NVIDIA's compiler tools (preview
# licensing) — point its knobs at the user's CUDA toolkit
spec.runtime_env.update(sparkdeps.triton_tool_env())
return spec
def _make_spark_spec(detection, wheels_source, dll_dirs):
return EnvSpec(
SPARK_BACKEND,
SPARK_TORCH,
torch_index=None,
python_version="3.12",
requirements_file="spark_requirements.txt",
# all pins resolve from the spark wheel set (self-built win_arm64):
# torchcodec, plus our triton port (torch.compile / compiled flex
# attention) — see the wheel-set build notes
extra_packages=_WIN_HELPERS + [
"torchcodec==0.15.0",
"triton==3.8.0+git8743423b",
],
# import-verified with rollback, like accelerators on other platforms
optional_packages=[
"flash-attn==2.8.3+cu134torch2.14",
"natten==0.21.7",
FLA,
],
find_links=[wheels_source],
notes=[
"RTX Spark native mode: win_arm64 CUDA %s stack from the "
"ai-toolkit wheel set (CUDA 13.4 developer preview), including "
"self-built flash-attn, NATTEN and triton (torch.compile)."
% SPARK_BACKEND,
],
uv_python=SPARK_UV_PYTHON,
torch_links=[wheels_source],
no_deps_packages=["tensorboard"],
runtime_dll_dirs=dll_dirs,
)
def build_spec(detection, allow_cpu=False):
"""Returns EnvSpec, or raises RuntimeError with a user-facing message."""
if (
detection["os"] == "windows"
and detection["arch"] == "aarch64"
and detection.get("backend") == "cuda"
and os.environ.get("AITK_SPARK_NATIVE", "1") != "0"
and _spark_capable(detection)
):
from . import sparkdeps
wheels_source = _spark_wheels_source()
# the CUDA toolkit is the one manual install (preview EULA) — without
# it the native wheels cannot run, so fall back to the x64 stack
if wheels_source and sparkdeps.cuda_bin_dir():
return _spark_spec(detection, wheels_source)
spec = _build_spec(detection, allow_cpu=allow_cpu)
if detection["os"] == "windows" and detection["arch"] == "aarch64":
# Emulated-x64 fallback (no native wheel source, old driver, or
# AITK_SPARK_NATIVE=0). Pin the venv interpreter to x64 explicitly:
# uv currently defaults to an emulated x86_64 Python on arm64 hosts,
# but says it will flip to native aarch64 once it considers support
# mature — which would silently leave a venv where the cu130 torch
# wheels don't resolve.
spec.uv_python = "cpython-%s-windows-x86_64-none" % spec.python_version
spec.notes.append(
"Windows-on-ARM detected: using the x64 stack under Windows' "
"emulation (GPU work still runs natively via the NVIDIA driver). "
"Native mode needs a CUDA 13.4+ driver and the Spark wheel set."
)
return spec
def _build_spec(detection, allow_cpu=False):
os_name = detection["os"]
if os_name == "mac":
notes = [
"flash-attn / NATTEN / triton / flash-linear-attention are "
"unavailable on macOS."
]
if detection["backend"] != "mps":
notes.append("Intel Mac detected — training will be extremely slow.")
return EnvSpec(
"mps",
TORCH,
python_version="3.12",
extra_packages=[TORCHCODEC],
notes=notes,
)
if detection["backend"] == "cuda":
return _cuda_spec(detection)
if detection["backend"] == "rocm":
return EnvSpec(
"rocm7.1",
TORCH,
torch_index=PYTORCH_INDEX + "rocm7.1",
extra_packages=[TORCHCODEC],
# runs on ROCm via the triton bundled with rocm torch
optional_packages=[FLA],
notes=[
"AMD ROCm support is experimental and largely untested.",
"flash-attn / NATTEN prebuilt wheels are unavailable for ROCm.",
],
)
# CPU fallback
if not allow_cpu:
raise RuntimeError(
"No supported GPU detected (NVIDIA CUDA, AMD ROCm, or Apple Silicon). "
"Training on CPU is not practical. Pass --cpu to install anyway."
)
return EnvSpec(
"cpu",
TORCH,
torch_index=PYTORCH_INDEX + "cpu",
extra_packages=[TORCHCODEC],
notes=["CPU-only install: training will be impractically slow."],
)

275
manager/util.py Normal file
View File

@ -0,0 +1,275 @@
"""Shared helpers for the AI Toolkit manager. Stdlib only."""
import json
import os
import platform
import shutil
import subprocess
import sys
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
IS_WINDOWS = platform.system() == "Windows"
IS_MAC = platform.system() == "Darwin"
IS_LINUX = platform.system() == "Linux"
USER_AGENT = (
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
"(KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
)
# When --json is used, human output goes to stderr so stdout stays machine-readable
_json_mode = False
def set_json_mode(enabled):
global _json_mode
_json_mode = enabled
def _supports_color(stream):
if os.environ.get("NO_COLOR"):
return False
return hasattr(stream, "isatty") and stream.isatty()
def _emit(prefix, msg, color):
stream = sys.stderr if _json_mode else sys.stdout
if _supports_color(stream):
stream.write("\033[%sm%s\033[0m %s\n" % (color, prefix, msg))
else:
stream.write("%s %s\n" % (prefix, msg))
stream.flush()
def info(msg):
_emit("[*]", msg, "36")
def ok(msg):
_emit("[+]", msg, "32")
def warn(msg):
_emit("[!]", msg, "33")
def error(msg):
_emit("[x]", msg, "31")
def die(msg, code=1):
error(msg)
sys.exit(code)
def print_json(data):
sys.stdout.write(json.dumps(data, indent=2) + "\n")
sys.stdout.flush()
def run(cmd, cwd=None, capture=False, check=True, env=None, stream=False):
"""Run a command. capture=True returns stdout text (stripped).
stream=True inherits stdio so the user sees live output.
Returns (returncode, stdout_or_None).
"""
kwargs = {"cwd": cwd or REPO_ROOT}
if env is not None:
kwargs["env"] = env
if capture:
kwargs["stdout"] = subprocess.PIPE
kwargs["stderr"] = subprocess.PIPE
elif _json_mode and not stream:
# keep stdout clean in json mode
kwargs["stdout"] = sys.stderr
try:
proc = subprocess.run(cmd, **kwargs)
except FileNotFoundError:
if check:
die("Command not found: %s" % cmd[0])
return 127, None
out = None
if capture:
out = proc.stdout.decode("utf-8", errors="replace").strip()
if check and proc.returncode != 0:
detail = ""
if capture and proc.stderr:
detail = "\n" + proc.stderr.decode("utf-8", errors="replace").strip()
die("Command failed (%d): %s%s" % (proc.returncode, " ".join(cmd), detail))
return proc.returncode, out
def which(name):
return shutil.which(name)
def find_uv():
"""Find uv: repo-local .uv/ first, then PATH, then common install dirs."""
home = os.path.expanduser("~")
candidates = [
os.path.join(REPO_ROOT, ".uv", "uv.exe" if IS_WINDOWS else "uv"),
]
for c in candidates:
if os.path.isfile(c) and os.access(c, os.X_OK):
return c
uv = shutil.which("uv")
if uv:
return uv
candidates = [
os.path.join(home, ".local", "bin", "uv"),
os.path.join(home, ".cargo", "bin", "uv"),
]
if IS_WINDOWS:
local = os.environ.get("LOCALAPPDATA", "")
if local:
candidates.append(os.path.join(local, "uv", "uv.exe"))
candidates.append(os.path.join(home, ".local", "bin", "uv.exe"))
for c in candidates:
if os.path.isfile(c) and os.access(c, os.X_OK):
return c
return None
def venv_dir():
"""Existing venv dir (.venv preferred, matching ui/cron/pythonPath.ts), else default target."""
for name in (".venv", "venv"):
d = os.path.join(REPO_ROOT, name)
if os.path.isdir(d):
return d
return os.path.join(REPO_ROOT, ".venv")
def venv_python(venv=None):
venv = venv or venv_dir()
if IS_WINDOWS:
return os.path.join(venv, "Scripts", "python.exe")
return os.path.join(venv, "bin", "python3")
# Env vars that let a system/conda/pyenv Python leak into our subprocesses.
# Scrubbed from every python/pip/node invocation (mirrors what the community
# Windows installer learned the hard way).
_SCRUB_VARS = (
"PYTHONPATH",
"PYTHONHOME",
"PYTHON",
"PYTHONSTARTUP",
"PYTHONUSERBASE",
"PYTHONEXECUTABLE",
"PIP_CONFIG_FILE",
"PIP_REQUIRE_VIRTUALENV",
"VIRTUAL_ENV",
"CONDA_PREFIX",
"CONDA_DEFAULT_ENV",
"PYENV_ROOT",
"PYENV_VERSION",
)
def clean_env(extra=None):
"""os.environ copy with Python-hijacking vars removed.
Also points uv's managed-python store into the repo (.uv/python) so
interpreter downloads never land outside the checkout.
"""
env = os.environ.copy()
for var in _SCRUB_VARS:
env.pop(var, None)
env.setdefault("UV_PYTHON_INSTALL_DIR", os.path.join(REPO_ROOT, ".uv", "python"))
if extra:
env.update(extra)
return env
def download(url, dest, label=None):
"""Download url to dest (stdlib only), logging progress every ~10%.
Sends a browser User-Agent: some mirrors (ffmpeg.martin-riedl.de) return
403 for the default "Python-urllib/x.y" agent.
"""
import urllib.request
label = label or os.path.basename(dest)
info("Downloading %s ..." % label)
last = [-1]
def report(read, total):
if total <= 0:
return
pct = min(100, int(read * 100 / total))
if pct >= last[0] + 10:
last[0] = pct
info(" %s: %d%%" % (label, pct))
req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
tmp = dest + ".part"
try:
with urllib.request.urlopen(req) as resp, open(tmp, "wb") as out:
total = int(resp.headers.get("Content-Length") or 0)
read = 0
while True:
chunk = resp.read(1 << 16)
if not chunk:
break
out.write(chunk)
read += len(chunk)
report(read, total)
except Exception as e: # noqa: BLE001 - surface any network failure clearly
if os.path.exists(tmp):
os.remove(tmp)
# Windows: python's ssl verifies against the OS cert store but never
# triggers Windows' on-demand intermediate-CA fetching, so on a fresh
# machine github downloads can fail CERTIFICATE_VERIFY_FAILED even
# though the chain is fine. curl (bundled since Win10, schannel-based)
# does fetch intermediates — fall back to it before giving up.
if not _download_with_curl(url, tmp, label):
die("Download failed for %s: %s" % (url, e))
os.replace(tmp, dest)
def _download_with_curl(url, tmp, label):
curl = shutil.which("curl")
if not curl:
return False
info(" %s: retrying with curl..." % label)
try:
code = subprocess.call(
[curl, "-fSL", "--retry", "3", "-o", tmp, url],
stdout=subprocess.DEVNULL,
stderr=subprocess.DEVNULL,
)
except OSError:
return False
if code != 0 or not os.path.exists(tmp):
if os.path.exists(tmp):
os.remove(tmp)
return False
return True
def extract_archive(archive, dest_dir):
"""Extract .zip / .tar.* into dest_dir (created if needed)."""
import tarfile
import zipfile
os.makedirs(dest_dir, exist_ok=True)
if archive.endswith(".zip"):
with zipfile.ZipFile(archive) as z:
z.extractall(dest_dir)
else:
with tarfile.open(archive) as t:
t.extractall(dest_dir)
def file_hash(paths):
import hashlib
h = hashlib.sha256()
for p in sorted(paths):
if os.path.isfile(p):
with open(p, "rb") as f:
h.update(f.read())
return h.hexdigest()

78
manager/uvbin.py Normal file
View File

@ -0,0 +1,78 @@
"""Local (never system) uv provisioning into <repo>/.uv/.
uv is a single static binary (no deps, no Python needed). Keeping it inside
the repo means fast installs + exact-version Python provisioning with zero
footprint outside the checkout. The run scripts also install into this same
dir (via the astral installer with UV_INSTALL_DIR) for the no-python
bootstrap case; util.find_uv() checks here first either way.
"""
import os
import platform
import shutil
import stat
import tempfile
from .util import IS_MAC, IS_WINDOWS, REPO_ROOT, download, extract_archive, ok, warn
UV_DIR = os.path.join(REPO_ROOT, ".uv")
_RELEASE = "https://github.com/astral-sh/uv/releases/latest/download/"
def _asset():
arch = platform.machine().lower()
if arch == "amd64":
arch = "x86_64"
if IS_WINDOWS:
return "uv-x86_64-pc-windows-msvc.zip"
if IS_MAC:
return "uv-%s-apple-darwin.tar.gz" % (
"aarch64" if arch == "arm64" else "x86_64"
)
return "uv-%s-unknown-linux-gnu.tar.gz" % (
"aarch64" if arch in ("aarch64", "arm64") else "x86_64"
)
def local_uv_exe():
return os.path.join(UV_DIR, "uv.exe" if IS_WINDOWS else "uv")
def ensure_uv(dry_run=False):
"""Download the uv binary into .uv/ if it isn't available anywhere."""
from .util import find_uv
if find_uv():
return False
if dry_run:
from .util import info
info("[dry-run] would download uv into %s" % UV_DIR)
return False
tmp = tempfile.mkdtemp(prefix="aitk_uv_")
try:
asset = _asset()
archive = os.path.join(tmp, asset)
download(_RELEASE + asset, archive, label="uv")
extract_archive(archive, tmp)
# zip: uv.exe at root; tarballs: uv-<triple>/uv — search the tree
exe_name = "uv.exe" if IS_WINDOWS else "uv"
found = None
for root, _dirs, files in os.walk(tmp):
if exe_name in files:
found = os.path.join(root, exe_name)
break
if not found:
warn("Unexpected uv archive layout — continuing without uv.")
return False
os.makedirs(UV_DIR, exist_ok=True)
dest = local_uv_exe()
shutil.move(found, dest)
os.chmod(
dest, os.stat(dest).st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
)
ok("uv installed at %s" % dest)
return True
finally:
shutil.rmtree(tmp, ignore_errors=True)

4
requirements_base.txt
View File

@ -1,4 +1,4 @@
torchao==0.10.0
torchao==0.17.0
safetensors
git+https://github.com/huggingface/diffusers.git@c943837899b16cbae2f619b8dd4f7bb6f07dd81a
#pip install git+https://github.com/huggingface/diffusers.git@refs/pull/13432/head
@ -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

10
run.py
View File

@ -3,9 +3,10 @@ import sys
from dotenv import load_dotenv
# Load the .env file if it exists
load_dotenv()
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = os.getenv("HF_HUB_ENABLE_HF_TRANSFER", "1")
os.environ["HF_XET_HIGH_PERFORMANCE"] = os.getenv("HF_XET_HIGH_PERFORMANCE", "1")
os.environ["HF_HUB_DISABLE_XET"] = os.getenv("HF_HUB_DISABLE_XET", "0")
os.environ["NO_ALBUMENTATIONS_UPDATE"] = "1"
os.environ["OPENCV_FFMPEG_LOGLEVEL"] = "-8"
seed = None
if "SEED" in os.environ:
try:
@ -14,6 +15,13 @@ if "SEED" in os.environ:
print(f"Invalid SEED value: {os.environ['SEED']}. SEED must be an integer.")
sys.path.insert(0, os.getcwd())
# The UI launches jobs with no console; keep anything we shell out to (torch
# compiles, HF git downloads) from flashing a console window. Must come before
# any import that might spawn a subprocess.
from toolkit.win_console import suppress_child_consoles
suppress_child_consoles()
# must come before ANY torch or fastai imports
# import toolkit.cuda_malloc

58
run_linux.sh Executable file
View File

@ -0,0 +1,58 @@
#!/usr/bin/env bash
# Update-and-run script for Linux — thin bootstrap over the in-repo manager.
#
# Everything (venv via uv-managed Python, torch for your GPU, requirements,
# portable Node.js and FFmpeg, dependency updates) is handled by
# `python -m manager`; this script only makes sure uv + a Python interpreter
# exist, then delegates. Works on desktop and headless boxes alike.
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
# ── Banner ─────────────────────────────────────────────────────────
echo ""
printf '\033[36m'
cat << 'BANNER'
_ ___ _____ _ _ _ _
/ \ |_ _| |_ _| ___ ___ | || | __(_)| |_
/ _ \ | | | | / _ \ / _ \| || |/ /| || __|
/ ___ \ | | | | | (_) || (_) | || < | || |_
/_/ \_\|___| |_| \___/ \___/|_||_|\_\|_| \__|
BANNER
printf '\033[0m'
printf '\033[90m AI Toolkit Manager — Linux\033[0m\n'
echo ""
# ── 1. Ensure uv (prebuilt static binary, kept inside the repo) ─────
export PATH="$SCRIPT_DIR/.uv:$PATH"
export UV_PYTHON_INSTALL_DIR="$SCRIPT_DIR/.uv/python"
if ! command -v uv >/dev/null 2>&1; then
echo "Downloading uv (package/python manager) into .uv/ ..."
curl -LsSf https://astral.sh/uv/install.sh | \
UV_INSTALL_DIR="$SCRIPT_DIR/.uv" UV_NO_MODIFY_PATH=1 sh
fi
# ── 2. Find a Python to run the manager (stdlib-only, needs >= 3.9) ─
find_python() {
for cmd in python3 python; do
if command -v "$cmd" >/dev/null 2>&1; then
if "$cmd" -c 'import sys; sys.exit(0 if sys.version_info >= (3, 9) else 1)' 2>/dev/null; then
echo "$cmd"
return 0
fi
fi
done
return 1
}
PYTHON="$(find_python || true)"
if [[ -z "$PYTHON" ]]; then
echo "No system Python found — provisioning one with uv..."
uv python install 3.12
PYTHON="$(uv python find 3.12)"
fi
# ── 3. Sync the environment and start the UI ────────────────────────
cd "$SCRIPT_DIR"
"$PYTHON" -m manager update --auto
exec "$PYTHON" -m manager launch

175
run_mac.zsh
View File

@ -1,5 +1,9 @@
#!/usr/bin/env zsh
# Update-and-run script for macOS — portable Python 3.12 + PyTorch
# Update-and-run script for macOS — thin bootstrap over the in-repo manager.
#
# Everything (venv via uv-managed Python, torch, requirements, portable
# Node.js and FFmpeg, dependency updates) is handled by `python -m manager`;
# this script only makes sure uv + a Python interpreter exist, then delegates.
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
@ -15,152 +19,39 @@ cat << 'BANNER'
/_/ \_\|___| |_| \___/ \___/|_||_|\_\|_| \__|
BANNER
echo "\033[0m"
echo "\033[90m macOS Setup & Launcher\033[0m"
echo "\033[90m AI Toolkit Manager — macOS\033[0m"
echo ""
VENV_DIR="$SCRIPT_DIR/.venv"
PIP="$VENV_DIR/bin/pip"
PYTHON="$VENV_DIR/bin/python3"
PYTHON_VERSION="3.12.8"
RELEASE_TAG="20241219"
# --- Package versions (update these as needed) ---
NODE_VERSION="23.11.1"
TORCH_VERSION="2.11.0"
TORCHVISION_VERSION="0.26.0"
TORCHAUDIO_VERSION="2.11.0"
# Detect architecture
ARCH="$(uname -m)"
if [[ "$ARCH" == "arm64" ]]; then
PLATFORM="aarch64-apple-darwin"
elif [[ "$ARCH" == "x86_64" ]]; then
PLATFORM="x86_64-apple-darwin"
else
echo "Error: Unsupported architecture: $ARCH"
exit 1
# ── 1. Ensure uv (prebuilt static binary, kept inside the repo) ─────
export PATH="$SCRIPT_DIR/.uv:$PATH"
export UV_PYTHON_INSTALL_DIR="$SCRIPT_DIR/.uv/python"
if ! command -v uv >/dev/null 2>&1; then
echo "Downloading uv (package/python manager) into .uv/ ..."
curl -LsSf https://astral.sh/uv/install.sh | \
UV_INSTALL_DIR="$SCRIPT_DIR/.uv" UV_NO_MODIFY_PATH=1 sh
fi
# ── 1. Download standalone Python if needed ─────────────────────────
PYTHON_DIR="$SCRIPT_DIR/.python"
PYTHON_BIN="$PYTHON_DIR/bin/python3"
if [[ ! -x "$PYTHON_BIN" ]]; then
TARBALL="cpython-${PYTHON_VERSION}+${RELEASE_TAG}-${PLATFORM}-install_only.tar.gz"
URL="https://github.com/indygreg/python-build-standalone/releases/download/${RELEASE_TAG}/${TARBALL}"
TMPDIR_DL="$(mktemp -d)"
trap 'rm -rf "$TMPDIR_DL"' EXIT
echo "Downloading standalone Python ${PYTHON_VERSION} (${PLATFORM})..."
curl -fSL --progress-bar -o "$TMPDIR_DL/$TARBALL" "$URL"
echo "Extracting..."
tar -xzf "$TMPDIR_DL/$TARBALL" -C "$TMPDIR_DL"
# Move to permanent location (the archive extracts to a "python" folder)
rm -rf "$PYTHON_DIR"
mv "$TMPDIR_DL/python" "$PYTHON_DIR"
rm -rf "$TMPDIR_DL"
trap - EXIT
echo "Standalone Python installed to $PYTHON_DIR"
fi
# ── 2. Create venv if it doesn't exist ──────────────────────────────
if [[ ! -d "$VENV_DIR" ]]; then
echo "Creating virtual environment at $VENV_DIR..."
"$PYTHON_BIN" -m venv "$VENV_DIR"
echo "Virtual environment created."
fi
# ── 3. Download / update portable Node.js ──────────────────────────
NODE_DIR="$SCRIPT_DIR/.node"
NODE_BIN="$NODE_DIR/bin/node"
NEED_NODE=false
if [[ ! -x "$NODE_BIN" ]]; then
NEED_NODE=true
elif [[ "$("$NODE_BIN" --version 2>/dev/null)" != "v${NODE_VERSION}" ]]; then
echo "Node.js version mismatch (want v${NODE_VERSION}, have $("$NODE_BIN" --version))."
NEED_NODE=true
fi
if $NEED_NODE; then
if [[ "$ARCH" == "arm64" ]]; then
NODE_ARCH="arm64"
else
NODE_ARCH="x64"
fi
NODE_TARBALL="node-v${NODE_VERSION}-darwin-${NODE_ARCH}.tar.gz"
NODE_URL="https://nodejs.org/dist/v${NODE_VERSION}/${NODE_TARBALL}"
TMPDIR_DL="$(mktemp -d)"
trap 'rm -rf "$TMPDIR_DL"' EXIT
echo "Downloading Node.js v${NODE_VERSION} (darwin-${NODE_ARCH})..."
curl -fSL --progress-bar -o "$TMPDIR_DL/$NODE_TARBALL" "$NODE_URL"
echo "Extracting..."
tar -xzf "$TMPDIR_DL/$NODE_TARBALL" -C "$TMPDIR_DL"
rm -rf "$NODE_DIR"
mv "$TMPDIR_DL/node-v${NODE_VERSION}-darwin-${NODE_ARCH}" "$NODE_DIR"
rm -rf "$TMPDIR_DL"
trap - EXIT
echo "Node.js v${NODE_VERSION} installed to $NODE_DIR"
else
echo "Node.js v${NODE_VERSION} is up to date."
fi
# ── 4. Install / update PyTorch packages ────────────────────────────
# Helper: returns 0 if the package is installed at the exact version
pkg_ok() {
local pkg="$1" want="$2"
local got
got="$("$PIP" show "$pkg" 2>/dev/null | awk '/^Version:/{print $2}')" || true
[[ "$got" == "$want" ]]
# ── 2. Find a Python to run the manager (stdlib-only, needs >= 3.9) ─
find_python() {
for cmd in python3 python; do
if command -v "$cmd" >/dev/null 2>&1; then
if "$cmd" -c 'import sys; sys.exit(0 if sys.version_info >= (3, 9) else 1)' 2>/dev/null; then
echo "$cmd"
return 0
fi
fi
done
return 1
}
PKGS_TO_INSTALL=()
pkg_ok "torch" "$TORCH_VERSION" || PKGS_TO_INSTALL+=("torch==$TORCH_VERSION")
pkg_ok "torchvision" "$TORCHVISION_VERSION" || PKGS_TO_INSTALL+=("torchvision==$TORCHVISION_VERSION")
pkg_ok "torchaudio" "$TORCHAUDIO_VERSION" || PKGS_TO_INSTALL+=("torchaudio==$TORCHAUDIO_VERSION")
if (( ${#PKGS_TO_INSTALL[@]} )); then
echo "Installing / updating: ${PKGS_TO_INSTALL[*]}"
"$PIP" install "${PKGS_TO_INSTALL[@]}"
else
echo "PyTorch packages are up to date."
PYTHON="$(find_python || true)"
if [[ -z "$PYTHON" ]]; then
echo "No system Python found — provisioning one with uv..."
uv python install 3.12
PYTHON="$(uv python find 3.12)"
fi
# ── 5. Install / update requirements.txt ────────────────────────────
REQUIREMENTS="$SCRIPT_DIR/requirements.txt"
REQ_HASH_FILE="$VENV_DIR/.requirements_hash"
if [[ -f "$REQUIREMENTS" ]]; then
# Hash all requirements files (follows -r includes)
CURRENT_HASH="$(cat "$SCRIPT_DIR"/requirements*.txt 2>/dev/null | shasum -a 256 | awk '{print $1}')"
STORED_HASH=""
[[ -f "$REQ_HASH_FILE" ]] && STORED_HASH="$(cat "$REQ_HASH_FILE")"
if [[ "$CURRENT_HASH" != "$STORED_HASH" ]]; then
echo "Installing / updating requirements.txt..."
"$PIP" install -r "$REQUIREMENTS"
echo "$CURRENT_HASH" > "$REQ_HASH_FILE"
else
echo "Requirements are up to date."
fi
fi
# ── 6. Build and start the UI ───────────────────────────────────────
export PATH="$NODE_DIR/bin:$VENV_DIR/bin:$PATH"
echo ""
echo "Starting UI..."
cd "$SCRIPT_DIR/ui"
npm run build_and_start
# ── 3. Sync the environment and start the UI ────────────────────────
cd "$SCRIPT_DIR"
"$PYTHON" -m manager update --auto
exec "$PYTHON" -m manager launch

1
run_modal.py
View File

@ -55,7 +55,6 @@ image = (
"toml",
"pydantic",
"omegaconf",
"k-diffusion",
"open_clip_torch",
"timm",
"prodigyopt",

77
run_windows.bat Normal file
View File

@ -0,0 +1,77 @@
@echo off&&cd /d %~dp0
REM Update-and-run script for Windows - thin bootstrap over the in-repo manager.
REM
REM Everything (venv via uv-managed Python, torch for your GPU, requirements,
REM portable Node.js / FFmpeg / Git, dependency updates) is handled by
REM `python -m manager`; this script only makes sure uv + a Python interpreter
REM exist, then delegates.
setlocal EnableDelayedExpansion
Title AI Toolkit
echo.
echo _ ___ _____ _ _ _ _
echo / \ ^|_ _^| ^|_ _^| ___ ___ ^| ^|^| ^| __(_)^| ^|_
echo / _ \ ^| ^| ^| ^| / _ \ / _ \^| ^|^| ^|/ /^| ^|^| __^|
echo / ___ \ ^| ^| ^| ^| ^| (_) ^|^| (_) ^| ^|^| ^< ^| ^|^| ^|_
echo /_/ \_\^|___^| ^|_^| \___/ \___/^|_^|^|_^|\_\^|_^| \__^|
echo.
echo AI Toolkit Manager - Windows
echo.
REM Clear env vars that let a stray conda/pyenv/system Python hijack things
set PYTHONPATH=
set PYTHONHOME=
set PYTHONSTARTUP=
set PYTHONUSERBASE=
set PIP_CONFIG_FILE=
set VIRTUAL_ENV=
set CONDA_PREFIX=
set CONDA_DEFAULT_ENV=
set PYENV_ROOT=
set PYENV_VERSION=
REM ---- 1. Ensure uv (prebuilt static binary, kept inside the repo) ----
set "PATH=%~dp0.uv;%PATH%"
set "UV_PYTHON_INSTALL_DIR=%~dp0.uv\python"
where uv.exe >nul 2>&1
if errorlevel 1 (
echo Downloading uv ^(package/python manager^) into .uv\ ...
powershell -NoProfile -ExecutionPolicy ByPass -Command ^
"$env:UV_INSTALL_DIR = Join-Path '%~dp0' '.uv'; $env:UV_NO_MODIFY_PATH = '1'; irm https://astral.sh/uv/install.ps1 | iex"
where uv.exe >nul 2>&1
if errorlevel 1 (
echo ERROR: uv download failed. See https://docs.astral.sh/uv/
pause
exit /b 1
)
)
REM ---- 2. Find a Python to run the manager (stdlib-only, needs 3.9+) ----
set "PY="
for %%C in (python.exe py.exe) do (
if not defined PY (
%%C -c "import sys; sys.exit(0 if sys.version_info >= (3, 9) else 1)" >nul 2>&1
if not errorlevel 1 set "PY=%%C"
)
)
if not defined PY (
echo No system Python found - provisioning one with uv...
uv python install 3.12
for /f "delims=" %%P in ('uv python find 3.12') do set "PY=%%P"
)
if not defined PY (
echo ERROR: could not find or install a Python interpreter.
pause
exit /b 1
)
REM ---- 3. Sync the environment and start the UI ----
"%PY%" -m manager update --auto
if errorlevel 1 (
echo.
echo Setup failed - see output above.
pause
exit /b 1
)
"%PY%" -m manager launch
pause

242
scripts/test_quantizations.py
View File

@ -17,15 +17,20 @@ Measures, per qtype:
accumulated through the block stack
- weight reconstruction error and one-time quantize (conversion) time
Runs on CUDA or on Apple Silicon (MPS). See DEVICE NOTES below for what MPS
can and cannot measure, and which qtypes it cannot run at all.
Usage:
python scripts/test_quantizations.py --gpu 1
python scripts/test_quantizations.py --gpu 1 --qtypes bf16 convrot8
python scripts/test_quantizations.py --device mps
"""
import argparse
import math
import os
import sys
import threading
import time
# set cuda bus ordering to be pcie
@ -35,6 +40,124 @@ sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch # noqa: E402
# ---------------------------------------------------------------- DEVICE NOTES
#
# MPS differs from CUDA in three ways that matter to a benchmark, so the numbers
# it prints are not interchangeable with CUDA's:
#
# 1. No peak-memory API. torch.mps has current_allocated_memory() but no
# max_memory_allocated()/reset_peak_memory_stats(), so the vram peaks are
# SAMPLED by a side thread polling the allocator (see _MpsPeakSampler)
# instead of read exactly. Steady-state values (resident weights) are exact.
# 2. No fp8 dtype. torch.float8_e4m3fn is undefined on MPS, which rules out the
# qfloat8 qtype and convrot4 (its nvfp4 block scales are stored as e4m3).
# MPS_UNSUPPORTED lists them; they are dropped from a default run.
# 3. No int8/fp4 tensor cores, so the convrot backends run their W8A16 fallback
# path rather than the W8A8/W4A4 fast path. Latency here says what Apple
# hardware does, not what the format is worth on a GPU that can run it.
def pick_device(name: str, gpu: int) -> torch.device:
if name == "auto":
if torch.cuda.is_available():
name = "cuda"
elif torch.backends.mps.is_available():
name = "mps"
else:
raise SystemExit("no cuda or mps device available")
if name == "cuda":
if not torch.cuda.is_available():
raise SystemExit("--device cuda requested but cuda is not available")
device = torch.device(f"cuda:{gpu}")
torch.cuda.set_device(device)
return device
if name == "mps":
if not torch.backends.mps.is_available():
raise SystemExit("--device mps requested but mps is not available")
return torch.device("mps")
raise SystemExit(f"unsupported device {name!r}")
def describe_device(device: torch.device) -> str:
if device.type == "cuda":
p = torch.cuda.get_device_properties(device)
return (f"{device} ({p.name}, sm_{p.major}{p.minor}, "
f"{p.total_memory / 1e9:.0f} GB)")
import platform
return (f"{device} (Apple {platform.machine()}, "
f"{torch.mps.recommended_max_memory() / 1e9:.0f} GB recommended max)")
def sync(device: torch.device) -> None:
if device.type == "cuda":
torch.cuda.synchronize(device)
elif device.type == "mps":
torch.mps.synchronize()
def empty_cache(device: torch.device) -> None:
if device.type == "cuda":
torch.cuda.empty_cache()
elif device.type == "mps":
torch.mps.empty_cache()
def mem_allocated(device: torch.device) -> int:
if device.type == "cuda":
return torch.cuda.memory_allocated(device)
if device.type == "mps":
return torch.mps.current_allocated_memory()
return 0
class _MpsPeakSampler:
"""Approximate max_memory_allocated for MPS by polling the allocator.
Tensors are allocated on the calling thread as ops are enqueued, so a
fine-grained poll from a side thread does observe the transients; on the
probe cases it recovered exact expected sizes (a 512 MiB transient and a
24 MiB matmul output). It can still miss a transient shorter than the poll
interval, so treat MPS peaks as a lower bound, not a hard number.
"""
def __init__(self, device, interval=0.0002):
self.device, self.interval = device, interval
self.peak = 0
self._stop = threading.Event()
def __enter__(self):
self.peak = mem_allocated(self.device)
self._thread = threading.Thread(target=self._poll, daemon=True)
self._thread.start()
return self
def _poll(self):
while not self._stop.is_set():
self.peak = max(self.peak, mem_allocated(self.device))
time.sleep(self.interval)
def __exit__(self, *exc):
sync(self.device)
self.peak = max(self.peak, mem_allocated(self.device))
self._stop.set()
self._thread.join()
return False
def measure_peak(run, device: torch.device, base: int) -> int:
"""Peak bytes allocated during run(), over base. run() is called once first
so lazy init and compilation are not charged to the steady-state peak."""
run()
sync(device)
if device.type == "cuda":
torch.cuda.reset_peak_memory_stats(device)
run()
sync(device)
return torch.cuda.max_memory_allocated(device) - base
with _MpsPeakSampler(device) as sampler:
run()
return sampler.peak - base
# (tokens, in_features, out_features) — FLUX/Wan-class projections
SPEED_SHAPES = [
(4096, 3072, 3072),
@ -54,6 +177,12 @@ QTYPES = [
"convrotint2", "convrotbitnet", "convrotcomfyw4a4",
]
# qtypes that cannot run on MPS at all (see DEVICE NOTES: no fp8 dtype).
# quanto's float8 path is worse than a hard failure — it catches the dtype error,
# prints "Failed to quantize", and leaves the layer in bf16, so it would benchmark
# as bf16 under a quantized label. Dropped for the same reason.
MPS_UNSUPPORTED = {"qfloat8", "float8", "convrot4"}
STACK_KEY = f"{VRAM_BLOCKS}-block stack"
@ -93,11 +222,11 @@ def fp_weight(module: torch.nn.Linear) -> torch.Tensor:
def bench(fn, iters: int, device) -> float:
for _ in range(max(3, iters // 5)):
fn()
torch.cuda.synchronize(device)
sync(device)
t0 = time.perf_counter()
for _ in range(iters):
fn()
torch.cuda.synchronize(device)
sync(device)
return (time.perf_counter() - t0) / iters * 1000 # ms
@ -195,20 +324,15 @@ def run_speed(qtype: str, device, iters: int, results: dict):
)
except Exception as e:
print(f" [{qtype}] compiled train failed for {m}x{k}->{n}: {e}")
torch.cuda.empty_cache()
empty_cache(device)
def _stack_fwd_peak(blocks, x, device, base) -> int:
# warm up first so lazy-init allocations (and compilation) are not counted
# as steady-state peak
with torch.no_grad():
stack_forward(blocks, x)
torch.cuda.synchronize(device)
torch.cuda.reset_peak_memory_stats(device)
with torch.no_grad():
stack_forward(blocks, x)
torch.cuda.synchronize(device)
return torch.cuda.max_memory_allocated(device) - base
def fwd():
with torch.no_grad():
stack_forward(blocks, x)
return measure_peak(fwd, device, base)
def _stack_train_peak(blocks, x, device, base, checkpoint) -> int:
@ -217,24 +341,19 @@ def _stack_train_peak(blocks, x, device, base, checkpoint) -> int:
xi = x.detach().requires_grad_(True)
stack_forward(blocks, xi, checkpoint).float().pow(2).mean().backward()
train_step()
torch.cuda.synchronize(device)
torch.cuda.reset_peak_memory_stats(device)
train_step()
torch.cuda.synchronize(device)
return torch.cuda.max_memory_allocated(device) - base
return measure_peak(train_step, device, base)
def run_vram(qtype: str, device, results: dict):
torch.cuda.empty_cache()
base = torch.cuda.memory_allocated(device)
empty_cache(device)
base = mem_allocated(device)
blocks = make_stack(device)
for b in blocks:
for i in range(len(b)):
b[i] = convert(b[i], qtype)
torch.cuda.empty_cache()
results[(qtype, "vram_weights")] = torch.cuda.memory_allocated(device) - base
empty_cache(device)
results[(qtype, "vram_weights")] = mem_allocated(device) - base
x = torch.randn(VRAM_TOKENS, 3072, device=device, dtype=torch.bfloat16)
@ -256,7 +375,7 @@ def run_vram(qtype: str, device, results: dict):
print(f" [{qtype}] compiled vram measurement failed: {e}")
blocks = x = None # release before the allocator accounting of the next run
torch.cuda.empty_cache()
empty_cache(device)
def run_drift(qtype: str, device, results: dict):
@ -270,7 +389,7 @@ def run_drift(qtype: str, device, results: dict):
lin = convert(lin, qtype)
y_q = lin(x).float()
results[(qtype, "drift", (m, k, n))] = ((y_q - y_ref).norm() / y_ref.norm()).item()
torch.cuda.empty_cache()
empty_cache(device)
blocks = make_stack(device)
x = torch.randn(VRAM_TOKENS, 3072, device=device, dtype=torch.bfloat16)
@ -282,24 +401,24 @@ def run_drift(qtype: str, device, results: dict):
y_q = stack_forward(blocks, x).float()
results[(qtype, "drift", STACK_KEY)] = ((y_q - y_ref).norm() / y_ref.norm()).item()
blocks = x = None
torch.cuda.empty_cache()
empty_cache(device)
def run_quality_and_quantize_time(qtype: str, device, results: dict):
torch.manual_seed(0)
lin = make_layer(3072, 3072, device)
w0 = lin.weight.detach().float().clone()
torch.cuda.synchronize(device)
sync(device)
t0 = time.perf_counter()
lin = convert(lin, qtype)
torch.cuda.synchronize(device)
sync(device)
results[(qtype, "quantize_ms")] = (time.perf_counter() - t0) * 1000
if qtype == "bf16":
results[(qtype, "weight_err")] = 0.0
else:
wq = fp_weight(lin)
results[(qtype, "weight_err")] = ((wq - w0).norm() / w0.norm()).item()
torch.cuda.empty_cache()
empty_cache(device)
def print_speed_table(title: str, kind: str, qts, results):
@ -325,23 +444,37 @@ def print_speed_table(title: str, kind: str, qts, results):
def main():
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--device", default="auto", choices=["auto", "cuda", "mps"],
help="accelerator to run on (default: cuda if present, else mps)")
ap.add_argument("--gpu", type=int, default=0, help="cuda device id to run on")
ap.add_argument("--qtypes", nargs="+", default=QTYPES, help=f"subset of {QTYPES}")
ap.add_argument("--qtypes", nargs="+", default=None, help=f"subset of {QTYPES}")
ap.add_argument("--iters", type=int, default=50, help="timing iterations per case")
args = ap.parse_args()
device = torch.device(f"cuda:{args.gpu}")
torch.cuda.set_device(device)
props = torch.cuda.get_device_properties(device)
print(f"device: cuda:{args.gpu} ({props.name}, sm_{props.major}{props.minor}, "
f"{props.total_memory / 1e9:.0f} GB)")
print(f"torch {torch.__version__}\n")
device = pick_device(args.device, args.gpu)
print(f"device: {describe_device(device)}")
print(f"torch {torch.__version__}")
explicit_qtypes = args.qtypes is not None
qtypes = args.qtypes if explicit_qtypes else list(QTYPES)
if device.type == "mps":
blocked = [qt for qt in qtypes if qt in MPS_UNSUPPORTED]
if blocked and not explicit_qtypes:
qtypes = [qt for qt in qtypes if qt not in MPS_UNSUPPORTED]
print(f"note: skipping {', '.join(blocked)} — no fp8 dtype on MPS")
elif blocked:
# asked for by name: try anyway, but say what is expected to happen
print(f"note: {', '.join(blocked)} need an fp8 dtype MPS does not have; "
"expect them to fail or to silently stay in bf16")
print("note: MPS vram peaks are sampled, not exact; convrot backends run "
"their W8A16 fallback (no int8/fp4 tensor cores)")
print()
# warm the toolkit import chain (module imports + custom-op registration) so it
# isn't charged to the first qtype's quantize timing
if any(qt != "bf16" for qt in args.qtypes):
if any(qt != "bf16" for qt in qtypes):
from toolkit.util.ostris_quant import get_ostris_quantizer
for qt in args.qtypes:
for qt in qtypes:
if qt != "bf16":
get_ostris_quantizer(qt)
@ -351,15 +484,36 @@ def main():
torch._dynamo.config.cache_size_limit = 4096
results = {}
for qt in args.qtypes:
ran = []
for qt in qtypes:
print(f"benchmarking {qt} ...")
torch._dynamo.reset() # drop the previous qtype's compiled artifacts
run_quality_and_quantize_time(qt, device, results)
run_drift(qt, device, results)
run_speed(qt, device, args.iters, results)
run_vram(qt, device, results)
try:
run_quality_and_quantize_time(qt, device, results)
run_drift(qt, device, results)
run_speed(qt, device, args.iters, results)
run_vram(qt, device, results)
except Exception as e:
# one unsupported backend should not cost the whole run
print(f" [{qt}] FAILED, dropped from the tables: {type(e).__name__}: {e}")
for key in list(results):
if key[0] == qt:
del results[key]
empty_cache(device)
continue
if qt != "bf16" and results.get((qt, "weight_err")) == 0.0:
print(f" [{qt}] quantization was a no-op (weights unchanged) — "
"dropped so it is not reported as a quantized result")
for key in list(results):
if key[0] == qt:
del results[key]
continue
ran.append(qt)
qts = args.qtypes
if not ran:
raise SystemExit("no qtype completed successfully")
qts = ran
print_speed_table("layer latency, inference", "inf", qts, results)
print_speed_table("layer latency, inference (compiled)", "inf_comp", qts, results)
print_speed_table("layer latency, train fwd+bwd", "train", qts, results)

84
spark_requirements.txt Normal file
View File

@ -0,0 +1,84 @@
# NVIDIA RTX Spark — native Windows-on-ARM (win_arm64) requirement set.
#
# Mirrors requirements_base.txt with pins adjusted to versions that actually
# publish win_arm64 wheels (or that we build ourselves — see manager/spec.py
# spark handling). Keep in sync with requirements_base.txt when bumping pins.
#
# Deliberate differences from requirements_base.txt:
# - scipy>=1.17 first release with win_arm64 wheels (also unlocks numpy 2)
# - matplotlib>=3.11 first release with win_arm64 wheels (base: 3.10.1)
# - av>=17 first release with win_arm64 wheels (base: 16.0.1)
# - numba/llvmlite rc only prerelease win_arm64 wheels exist (librosa dep)
# - tensorboard REMOVED: its grpcio dep has no win_arm64 wheels; the
# manager installs tensorboard --no-deps (writer path
# only needs protobuf) as a spark extra
# - hf_transfer REMOVED: no win_arm64 wheels; hf-xet (arm64 OK) is the
# modern replacement and installs via manager extras
# - torchcodec pin REMOVED: the manager installs the locally built
# win_arm64 torchcodec wheel as an extra
# - opencv-python, soxr, pywavelets, brotli(gradio), kornia-rs have no
# published win_arm64 wheels at any version — supplied from the ai-toolkit
# spark wheel set (self-built), resolved via --find-links
numpy>=2,<3
scipy>=1.17
torchao==0.17.0
safetensors
git+https://github.com/huggingface/diffusers.git@c943837899b16cbae2f619b8dd4f7bb6f07dd81a
transformers==5.5.3
lycoris-lora==1.8.3
flatten_json
pyyaml
oyaml
kornia
invisible-watermark
einops
accelerate
toml
albumentations==1.4.15
albucore==0.0.16
pydantic
omegaconf
open_clip_torch
timm==1.0.22
prodigyopt
controlnet_aux==0.0.10
python-dotenv
bitsandbytes
lpips
pytorch_fid
optimum-quanto==0.2.4
sentencepiece
huggingface_hub==1.23.0
peft==0.18.1
gradio
python-slugify
# exact pins: newer opencv versions exist on PyPI only as sdists (which fail
# to build on MSVC arm64 — dnn __fp16); force the self-built 4.12 wheels
opencv-python==4.12.0.88
opencv-python-headless==4.12.0.88
pytorch-wavelets==1.3.0
matplotlib>=3.11
setuptools>=77.0.3
av>=17
# librosa chain: numba/llvmlite have no cp312 win_arm64 wheels upstream — the
# rc pins below resolve from the self-built spark wheel set (llvmlite built
# against our own LLVM 22 static build)
librosa==0.11.0
numba==0.67.0rc1
llvmlite==0.49.0rc1
soxr==1.1.0
mutagen==1.47.0
soundfile
# tensorboard is installed --no-deps by the manager (its grpcio dep has no
# win_arm64 wheels; grpc is only needed by the tensorboard server, not the
# SummaryWriter path the toolkit uses). Its remaining runtime deps:
absl-py
markdown
werkzeug
tensorboard-data-server
packaging
protobuf
six

10
toolkit/config_modules.py
View File

@ -711,11 +711,15 @@ class ModelConfig:
if self.layer_offloading and self.qtype_te == "qfloat8":
self.qtype_te = "float8"
# Mac mps only works with torachao uint
# MPS has no fp8 dtype, so qfloat8 has to become an 8 bit integer format.
# convrot8, not torchao int8: measured on an M3 against bf16, convrot8
# trains at 0.79x and holds 1.04 GB of resident weight where torchao int8
# trains at 0.52x and holds 1.21 GB, and convrot8 quantizes in 19ms
# against 2.8s. See scripts/test_quantizations.py --device mps.
if torch.backends.mps.is_available() and self.qtype == "qfloat8":
self.qtype = "int8"
self.qtype = "convrot8"
if torch.backends.mps.is_available() and self.qtype_te == "qfloat8":
self.qtype_te = "int8"
self.qtype_te = "convrot8"
# 0 is off and 1.0 is 100% of the layers
self.layer_offloading_transformer_percent = kwargs.get("layer_offloading_transformer_percent", 1.0)

44
toolkit/data_loader.py
View File

@ -533,6 +533,7 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
size_database=self.size_database,
dataset_root=dataset_folder,
encode_control_in_text_embeddings=self.sd.encode_control_in_text_embeddings if self.sd else False,
encode_first_frame_in_text_embeddings=getattr(self.sd, 'encode_first_frame_in_text_embeddings', False) if self.sd else False,
text_embedding_space_version=self.sd.text_embedding_space_version if self.sd else "sd1",
te_padding_side=self.sd.te_padding_side if self.sd else "right",
latent_space_version=latent_space_version,
@ -606,6 +607,14 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
self.setup_controls()
self.epoch_num += 1
def __getstate__(self):
# on Windows/macOS dataloader workers are spawned, which pickles the dataset.
# sd (the model) is not picklable (weakrefs, cuda tensors) and is only needed
# for caching, which runs in the main process before iteration starts.
state = self.__dict__.copy()
state['sd'] = None
return state
def __len__(self):
if self.dataset_config.buckets:
return len(self.batch_indices)
@ -652,6 +661,13 @@ class AiToolkitDataset(LatentCachingMixin, ControlCachingMixin, CLIPCachingMixin
return self._get_single_item(item)
def dto_collation(batch: List['FileItemDTO']):
# must be a module level function so spawned dataloader workers can pickle it
return DataLoaderBatchDTO(
file_items=batch
)
def get_dataloader_from_datasets(
dataset_options,
batch_size=1,
@ -692,22 +708,26 @@ def get_dataloader_from_datasets(
# todo build scheduler that can get buckets from all datasets that match
# todo and evenly distribute reg images
def dto_collation(batch: List['FileItemDTO']):
# create DTO batch
batch = DataLoaderBatchDTO(
file_items=batch
)
return batch
# check if is caching latents
dataloader_kwargs = {}
if is_native_windows() or is_macos():
dataloader_kwargs['num_workers'] = 0
else:
dataloader_kwargs['num_workers'] = dataset_config_list[0].num_workers
dataloader_kwargs['num_workers'] = dataset_config_list[0].num_workers
if dataloader_kwargs['num_workers'] > 0:
dataloader_kwargs['prefetch_factor'] = dataset_config_list[0].prefetch_factor
# keep workers alive across epochs. Without this, spawn platforms (Windows/macOS)
# boot new worker processes every epoch, which can take longer than the epoch
# itself on small datasets. The dataset is static after epoch 0 (setup_epoch only
# does work on the first call) and per-epoch shuffling happens in the main process
# sampler, so workers never hold stale state.
dataloader_kwargs['persistent_workers'] = True
# spawned workers re-import the full stack at boot and would repeat every
# import-time warning the parent already printed. Children inherit these env
# vars; the parent is unaffected since its imports already happened.
os.environ.setdefault('PYTHONWARNINGS', 'ignore::FutureWarning')
os.environ.setdefault('TORCH_LOGS', '-torch.utils._pytree')
os.environ.setdefault('DIFFUSERS_VERBOSITY', 'error')
os.environ.setdefault('NO_ALBUMENTATIONS_UPDATE', '1')
if has_buckets:
# make sure they all have buckets

3
toolkit/data_transfer_object/data_loader.py
View File

@ -65,6 +65,9 @@ class FileItemDTO(
self.encode_control_in_text_embeddings = kwargs.get(
"encode_control_in_text_embeddings", False
)
self.encode_first_frame_in_text_embeddings = kwargs.get(
"encode_first_frame_in_text_embeddings", False
)
self.te_padding_side = kwargs.get("te_padding_side", "right")
self.latent_space_version = kwargs.get("latent_space_version", "sd1")
self.text_embedding_space_version = kwargs.get("text_embedding_space_version", "sd1")

123
toolkit/dataloader_mixins.py
View File

@ -560,6 +560,35 @@ class ImageProcessingDTOMixin:
unique_frame_idxs = sorted(set(frames_to_extract))
processed_frames = {} # frame_idx -> processed frame (duplicates reuse it)
def process_frame(rgb_frame):
# Convert to PIL Image
img = Image.fromarray(rgb_frame)
# Apply the same processing as for single images
img = img.convert('RGB')
if self.flip_x:
img = img.transpose(Image.FLIP_LEFT_RIGHT)
if self.flip_y:
img = img.transpose(Image.FLIP_TOP_BOTTOM)
# Apply bucketing
img = img.resize((self.scale_to_width, self.scale_to_height), Image.BICUBIC)
img = img.crop((
self.crop_x,
self.crop_y,
self.crop_x + self.crop_width,
self.crop_y + self.crop_height
))
# Apply transform if provided
if transform:
img = transform(img)
return img
decode_with_pyav = False
# Set frame position
pos = unique_frame_idxs[0]
if pos > 0:
@ -580,10 +609,6 @@ class ImageProcessingDTOMixin:
if ret:
pos += 1
else:
# Try to provide more detailed error information
actual_frame = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
frame_pos_info = f"Requested frame: {frame_idx}, Actual frame position: {actual_frame}"
# Try to read the next available frame as a fallback
fallback_success = False
for fallback_offset in [1, -1, 5, -5, 10, -10]:
@ -600,38 +625,44 @@ class ImageProcessingDTOMixin:
pos = fallback_pos + 1
break
else:
# No fallback worked, raise a more detailed exception
video_info = f"Video: {self.path}, Total frames: {total_frames}, FPS: {video_fps}"
raise Exception(f"Failed to read frame {frame_idx} from video. {frame_pos_info}. {video_info}")
# No fallback worked. cv2's bundled ffmpeg cannot decode some codecs
# at all (e.g. AV1 has no software decoder there), so retry the
# remaining frames with PyAV below.
decode_with_pyav = True
break
# Convert BGR to RGB
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# Convert to PIL Image
img = Image.fromarray(frame)
# Apply the same processing as for single images
img = img.convert('RGB')
if self.flip_x:
img = img.transpose(Image.FLIP_LEFT_RIGHT)
if self.flip_y:
img = img.transpose(Image.FLIP_TOP_BOTTOM)
# Apply bucketing
img = img.resize((self.scale_to_width, self.scale_to_height), Image.BICUBIC)
img = img.crop((
self.crop_x,
self.crop_y,
self.crop_x + self.crop_width,
self.crop_y + self.crop_height
))
# Apply transform if provided
if transform:
img = transform(img)
processed_frames[frame_idx] = img
processed_frames[frame_idx] = process_frame(frame)
if decode_with_pyav:
# cv2 could not decode this video (e.g. AV1: OpenCV's bundled ffmpeg has no
# software AV1 decoder). Decode the still-missing frames in one sequential
# PyAV pass; PyAV ships libdav1d so it handles codecs cv2 cannot.
import av
needed = {i for i in unique_frame_idxs if i not in processed_frames}
last_av_frame = None
with av.open(self.path) as container:
decoded_idx = -1
for av_frame in container.decode(video=0):
decoded_idx += 1
last_av_frame = av_frame
if decoded_idx in needed:
processed_frames[decoded_idx] = process_frame(av_frame.to_ndarray(format='rgb24'))
needed.discard(decoded_idx)
if not needed:
break
if needed:
if last_av_frame is None:
video_info = f"Video: {self.path}, Total frames: {total_frames}, FPS: {video_fps}"
raise Exception(f"Failed to read frames {sorted(needed)} from video with both cv2 and PyAV. {video_info}")
# metadata frame count overshot the real stream; reuse the last decoded frame
tail_frame = process_frame(last_av_frame.to_ndarray(format='rgb24'))
for frame_idx in needed:
processed_frames[frame_idx] = tail_frame
# assemble in extraction order; stretched clips repeat decoded frames
frames = [processed_frames[frame_idx] for frame_idx in frames_to_extract]
@ -1899,6 +1930,13 @@ class TextEmbeddingFileItemDTOMixin:
# if we have a control image, cache the path
if self.encode_control_in_text_embeddings and self.control_path is not None:
item["control_path"] = self.control_path
# first-frame vision conditioning changes the embedding content -> new cache key
elif (
getattr(self, "encode_first_frame_in_text_embeddings", False)
and self.dataset_config.do_i2v
and (self.dataset_config.auto_frame_count or self.dataset_config.num_frames > 1)
):
item["first_frame_in_te"] = True
return item
def get_text_embedding_path(self: 'FileItemDTO', recalculate=False):
@ -1984,6 +2022,27 @@ class TextEmbeddingCachingMixin:
else:
ctrl_img = ctrl_img_list
prompt_embeds: PromptEmbeds = self.sd.encode_prompt(file_item.caption, control_images=ctrl_img)
elif (
getattr(self.sd, 'encode_first_frame_in_text_embeddings', False)
and self.dataset_config.do_i2v
and (self.dataset_config.auto_frame_count or self.dataset_config.num_frames > 1)
):
# video item: encode the clip's FIRST FRAME into the text embeddings
# as a vision reference, matching sampling (where the ctrl image goes
# into the embeds and is held as the clean first frames)
file_item.load_and_process_image(self.transform, only_load_latents=True)
frames = file_item.tensor # (T, C, H, W) or (C, H, W), in [-1, 1]
first = frames[0] if frames.dim() == 4 else frames
ctrl_img = (
((first + 1.0) / 2.0)
.clamp(0, 1)
.unsqueeze(0)
.to(self.sd.device_torch, dtype=self.sd.torch_dtype)
)
if self.sd.has_multiple_control_images:
ctrl_img = [ctrl_img]
prompt_embeds: PromptEmbeds = self.sd.encode_prompt(file_item.caption, control_images=ctrl_img)
file_item.tensor = None
else:
prompt_embeds: PromptEmbeds = self.sd.encode_prompt(file_item.caption)
# save it

252
toolkit/util/convrot_quant.py
View File

@ -95,9 +95,17 @@ def regular_hadamard(rot_size: int, device, dtype=torch.bfloat16) -> torch.Tenso
key = (rot_size, str(device), dtype)
def build():
# fp32, not fp64: entries stay exactly +-1 through the krons and rot_size
# is a power of 4, so dividing by its (power-of-two) root is exact — this
# build is bit-identical to an fp64 one at every supported rot_size, and
# verified so. fp64 is not merely unnecessary here but harmful: MPS has no
# float64, and on a cache miss inside a torch.compile trace inductor lifts
# the matrix in as a graph constant and dies moving it to the device
# ("Cannot convert a MPS Tensor to float64").
r4 = torch.tensor(
[[1.0, 1, 1, -1], [1, 1, -1, 1], [1, -1, 1, 1], [-1, 1, 1, 1]],
dtype=torch.float64,
dtype=torch.float32,
device="cpu",
)
h = r4.clone()
while h.shape[0] < rot_size:
@ -1040,19 +1048,223 @@ def _int8_gemm_supported(device) -> bool:
global _warned_no_int8
if not supported and not _warned_no_int8:
_warned_no_int8 = True
print_acc(
f"ConvRot: int8 matmul (torch._int_mm) is not usable on this device "
f"({device}). Inference falls back to dequantized bf16 matmuls: correct "
"output but NO speedup, and inference activations stay unquantized "
"(W8A16 numerics instead of W8A8). The training path is unaffected "
"(it always simulates W8A8 via fake-quant)."
)
if _mps_int8pack_device_ok(device):
print_acc(
f"ConvRot: int8 tensor-core matmul (torch._int_mm) is not implemented "
f"on this device ({device}). Inference uses torch._weight_int8pack_mm "
"instead — the int8 weight feeds the gemm directly, with no bf16 "
"dequant transient, which is the whole win on a bandwidth-bound Apple "
"GPU. Activations stay unquantized (W8A16 numerics instead of W8A8); "
"Apple GPUs have no int8 tensor cores, so quantizing them would only "
"add overhead. Layers whose in/out features are not both divisible by "
"32 still use a dequantized bf16 matmul. The training path is "
"unaffected (it always simulates W8A8 via fake-quant)."
)
else:
print_acc(
f"ConvRot: int8 matmul (torch._int_mm) is not usable on this device "
f"({device}). Inference falls back to dequantized bf16 matmuls: correct "
"output but NO speedup, and inference activations stay unquantized "
"(W8A16 numerics instead of W8A8). The training path is unaffected "
"(it always simulates W8A8 via fake-quant)."
)
return supported
_warned_no_int8 = False
# ---------------- mps int8-weight gemm ----------------
#
# Apple GPUs have no int8 tensor cores and torch._int_mm has no MPS kernel, so
# there is no W8A8 path here and no reason to want one: a quantized activation
# would buy nothing back on hardware that multiplies in fp anyway. What DOES pay
# on a unified-memory Apple GPU is not materializing the dequantized weight —
# torch._weight_int8pack_mm (int8 (out, in) weight + per-row scales, exactly the
# layout convrot8/convrotint already store) feeds the codes to the gemm directly.
#
# Measured on an M3, convrot8 3072x3072, vs the dequant+F.linear fallback it
# replaces: 18x at M=1, 5.7x at M=64, 1.2x at M=1024, ~1.0x at M=4096. The
# fallback's dequant is a fixed ~2.7ms that amortizes away as M grows, so the win
# concentrates at decode/small-batch shapes and large batches are a wash. The
# bitpacked convrotint*/bitnet backends see only 1.1-1.4x — their per-forward
# weight unpack dominates on MPS (no triton fused gemv path there). Error vs the
# unquantized bf16 reference is marginally LOWER than the fallback's. Holding no
# dequantized weight transient is a win at every M on unified memory.
_int8pack_mm_ok = None
def _mps_int8pack_device_ok(device) -> bool:
"""Whether torch._weight_int8pack_mm has a working kernel on this device."""
global _int8pack_mm_ok
if torch.device(device).type != "mps":
return False
if _int8pack_mm_ok is None:
try:
torch._weight_int8pack_mm(
torch.zeros(1, 32, dtype=torch.bfloat16, device=device),
torch.zeros(32, 32, dtype=torch.int8, device=device),
torch.ones(32, dtype=torch.bfloat16, device=device),
)
_int8pack_mm_ok = True
except Exception:
_int8pack_mm_ok = False
return _int8pack_mm_ok
def _mps_int8pack_usable(device, in_f: int, out_f: int, dtype: torch.dtype) -> bool:
# the kernel asserts N % 32 == 0 && K % 32 == 0; can_quantize only guarantees
# in % 16 and out % 8, so the odd layer still needs the dequant fallback
return (
in_f % 32 == 0
and out_f % 32 == 0
and dtype in (torch.bfloat16, torch.float16, torch.float32)
and _mps_int8pack_device_ok(device)
)
# registered as a custom op so torch.compile treats it as an opaque node with a
# known output shape (see _nvfp4_act_quant_op). Calling the aten op directly
# compiles fine on cuda but inductor's MPS lowering asserts on it as soon as its
# input is an unrealized Pointwise — which is exactly what the rotate() feeding
# it produces, so every compiled convrot8 inference on mac fell back to eager.
@torch.library.custom_op("ostris::convrot_mps_int8pack_mm", mutates_args=())
def _mps_int8pack_mm_op(
x2d: torch.Tensor, qdata: torch.Tensor, scales: torch.Tensor
) -> torch.Tensor:
# scales is cast to x2d's dtype deliberately: the kernel reads the scale
# buffer AS the activation dtype with no check, so handing it the fp32 scales
# convrot stores returns silent garbage (nan / zeros) rather than an error —
# the one sharp edge here. It also rejects non-contiguous operands outright;
# both are contiguous today, so those calls are free no-ops that keep this
# correct if a caller's layout ever changes.
return torch._weight_int8pack_mm(
x2d.contiguous(), qdata.contiguous(), scales.to(x2d.dtype)
)
@_mps_int8pack_mm_op.register_fake
def _mps_int8pack_mm_fake(x2d, qdata, scales):
return torch.empty(
x2d.shape[0], qdata.shape[0], device=x2d.device, dtype=x2d.dtype
)
def _mps_int8pack_linear(x2d, qdata, scales, bias) -> torch.Tensor:
"""out = x2d @ dequant(qdata).T (+ bias) via the MPS int8-weight gemm."""
out = _mps_int8pack_mm_op(x2d, qdata, scales)
if bias is not None:
out = out + bias
return out
# ---------------- mps: fused elementwise chains ----------------
#
# On MPS the gemm runs in bf16 whatever the weight storage is, so there is no
# matmul-level win to chase: torch._weight_int8pack_mm measures within 2% of an
# equal-shape bf16 F.linear for M >= 256 (it only pulls ahead at decode-size
# batches — 2.1x at M=1, 1.2x at M=64, measured on an M3 at 3072x3072). What
# separates convrot8 from plain bf16 at training shapes is therefore ENTIRELY the
# elementwise work wrapped around the gemm: the activation fake-quant and the
# weight dequant.
#
# Eager runs each of those as a chain of ~8 kernels over full-size fp32
# temporaries, which on unified memory costs more than the quantization saves —
# the activation fake-quant alone measures 9.1ms against a 25.8ms bf16 gemm (35%
# overhead for one of two chains). Inductor's MPS backend fuses each chain into a
# single Metal kernel and takes that 9.1ms to 1.0ms, which is what moves convrot8
# training from 0.69x of bf16 to ~0.99x.
#
# The two chains want OPPOSITE compile modes, both for measured reasons — hence
# the per-chain `dynamic` below rather than one global setting:
#
# activations (_fake_quant_rows), dynamic=True: the batch dim genuinely varies,
# and a static compile lowers the fp32 divide to a reciprocal-multiply, which
# shifts ~0.002% of codes by 1. (Towards an fp64 reference, as it happens —
# but it is still a silent numerics change, and bit-identity with the cuda
# path is this backend's whole contract.) dynamic=True keeps the divide and
# is bit-identical, one compile covers every shape, and it costs ~0.4ms on a
# 52ms step against the static build.
#
# weights (_dequant_rows), dynamic=False: here dynamic shapes destroy the win
# outright — 2.5ms vs 0.38ms static, against 2.8ms eager, i.e. no better than
# not compiling at all (inductor cannot vectorize the broadcast multiply
# without a static inner width). All modes are bit-identical for this chain,
# so static is free of the numerics objection. Weight shapes are also fixed
# per layer, so this only recompiles once per distinct layer shape.
#
# The static build does hit dynamo's recompile limit (8) on a model with more
# than 8 distinct linear shapes. That degrades gracefully and was measured, not
# assumed: the first 8 shapes keep their compiled kernels and everything past
# them runs the eager chain, which is exactly today's behaviour. Nothing silently
# breaks, the tail just stops getting faster.
_mps_fused_cache = {}
_mps_fuse_disabled = False
def _mps_fused(fn, dynamic):
"""Run `fn` through inductor, falling back to eager for good on error.
Callers reach these only on mps — cuda has the triton kernels above and keeps
its existing code path untouched. Compilation is also skipped while tracing,
so an outer torch.compile fuses the chain into its own graph rather than
meeting an opaque nested compile.
"""
def run(*args):
global _mps_fuse_disabled
if _mps_fuse_disabled or torch.compiler.is_compiling():
return fn(*args)
compiled = _mps_fused_cache.get(fn)
if compiled is None:
compiled = torch.compile(fn, dynamic=dynamic)
_mps_fused_cache[fn] = compiled
try:
return compiled(*args)
except Exception as e:
_mps_fuse_disabled = True
print_acc(
f"ConvRot: inductor fusion is unavailable on this mps build "
f"({type(e).__name__}: {e}) — falling back to the eager "
"elementwise path. Output is unchanged; expect roughly 0.7x bf16 "
"training speed instead of ~1.0x."
)
return fn(*args)
return run
def _fake_quant_rows_impl(x: torch.Tensor, qmax: int) -> torch.Tensor:
"""dequant(quant(x)) per row, returned in x's dtype — the value half of the
activation STE.
Bit-identical to quantize_int8_rows followed by the dequant it feeds, with
both full-size fp32 temporaries dropped: amax over the input equals amax over
its fp32 upcast (bf16 -> fp32 is exact and a max is a selection, not an
accumulation), and the int8 round trip is lossless on a value already rounded
and clamped into [-qmax, qmax].
"""
scales = x.abs().amax(dim=1).float() / qmax
scales = torch.where(scales > 0, scales, torch.ones_like(scales))
s = scales.unsqueeze(1)
return (torch.round(x.float() / s).clamp_(-qmax, qmax) * s).to(x.dtype)
_fake_quant_rows = _mps_fused(_fake_quant_rows_impl, dynamic=True)
def _dequant_rows_impl(
qdata: torch.Tensor, scales: torch.Tensor, dtype: torch.dtype
) -> torch.Tensor:
"""dequant of per-row-scaled integer codes: the weight-side chain."""
return (qdata.float() * scales.unsqueeze(1)).to(dtype)
_dequant_rows = _mps_fused(_dequant_rows_impl, dynamic=False)
class ConvRotInt8Quantizer(OstrisQuantizer):
"""ConvRot W8A8 backend: shared regular-Hadamard rotation + per-token /
per-output-channel symmetric int8 with torch._int_mm. One instance per qtype,
@ -1124,7 +1336,10 @@ class ConvRotInt8Quantizer(OstrisQuantizer):
return torch.round(w_rot.float() / s).clamp_(-127, 127) * s
def _dequantize_rotated(self, module, dtype: torch.dtype) -> torch.Tensor:
w = self._qdata(module).float() * self._scales(module).unsqueeze(1)
qdata, scales = self._qdata(module), self._scales(module)
if qdata.device.type == "mps":
return _dequant_rows(qdata, scales, dtype)
w = qdata.float() * scales.unsqueeze(1)
return w.to(dtype)
def dequantize(self, module) -> torch.Tensor:
@ -1176,8 +1391,13 @@ class ConvRotInt8Quantizer(OstrisQuantizer):
# no int8 hardware: straight-through fake-quant + bf16 matmul
x2d = rotate(x, rot).reshape(-1, in_f)
with torch.no_grad():
aq, a_s = quantize_int8_rows(x2d.detach(), self.act_qmax)
x_dq = (aq.float() * a_s.unsqueeze(1)).to(x.dtype)
if x.device.type == "mps":
# same arithmetic, as one fused Metal kernel instead of a
# chain over full-size fp32 temporaries (see _mps_fused)
x_dq = _fake_quant_rows(x2d.detach(), self.act_qmax)
else:
aq, a_s = quantize_int8_rows(x2d.detach(), self.act_qmax)
x_dq = (aq.float() * a_s.unsqueeze(1)).to(x.dtype)
w = self._dequantize_rotated(module, x.dtype)
x_ste = x2d + (x_dq - x2d).detach()
out = F.linear(x_ste, w, module.bias)
@ -1214,6 +1434,16 @@ class ConvRotInt8Quantizer(OstrisQuantizer):
)
return out.reshape(*x.shape[:-1], out_f)
if _mps_int8pack_usable(x.device, in_f, out_f, x.dtype):
# W8A16 on the int8 weight codes, no dequantized weight transient
out = _mps_int8pack_linear(
rotate(x, rot).reshape(-1, in_f),
self._qdata(module),
self._scales(module),
module.bias,
)
return out.reshape(*x.shape[:-1], out_f)
w = self._dequantize_rotated(module, x.dtype)
out = F.linear(rotate(x, rot).reshape(-1, in_f), w, module.bias)
return out.reshape(*x.shape[:-1], out_f)

58
toolkit/win_console.py Normal file
View File

@ -0,0 +1,58 @@
"""Keep child processes from popping console windows on Windows.
The UI worker launches training jobs with pythonw.exe and DETACHED_PROCESS so
they survive the UI shutting down (see ui/cron/actions/startJob.ts). That
leaves the job with no console at all, and Windows hands a brand new console
-- with a visible window -- to any console program launched from a process
that has none. MSVC during a torch/triton compile, git during a HF download
and ffmpeg would each flash a window on the user's desktop.
Defaulting those spawns to CREATE_NO_WINDOW suppresses the flash. This is a
no-op unless we are on Windows *and* have no console, so running run.py from a
terminal behaves exactly as it did before.
"""
import subprocess
import sys
CREATE_NEW_CONSOLE = 0x00000010
DETACHED_PROCESS = 0x00000008
CREATE_NO_WINDOW = 0x08000000
# creationflags is the 14th positional parameter of Popen.__init__ after self.
_CREATIONFLAGS_POSITION = 14
_patched = False
def _has_console():
import ctypes
return bool(ctypes.windll.kernel32.GetConsoleWindow())
def suppress_child_consoles():
"""Make CREATE_NO_WINDOW the default for subprocesses, where it matters."""
global _patched
if _patched or sys.platform != "win32":
return
try:
if _has_console():
return
except Exception:
# Never let a console tweak take down a training run.
return
original_init = subprocess.Popen.__init__
def patched_init(self, *args, **kwargs):
if len(args) >= _CREATIONFLAGS_POSITION:
# Passed positionally; leave the caller's choice alone.
return original_init(self, *args, **kwargs)
flags = kwargs.get("creationflags", 0)
if not flags & (CREATE_NEW_CONSOLE | DETACHED_PROCESS | CREATE_NO_WINDOW):
kwargs["creationflags"] = flags | CREATE_NO_WINDOW
return original_init(self, *args, **kwargs)
subprocess.Popen.__init__ = patched_init
_patched = True

232
ui/cron/actions/startJob.ts
View File

@ -1,10 +1,10 @@
import prisma from '../prisma';
import { Job } from '@prisma/client';
import { spawn } from 'child_process';
import { spawn, ChildProcess } from 'child_process';
import path from 'path';
import fs from 'fs';
import { TOOLKIT_ROOT, getTrainingFolder, getHFToken } from '../paths';
import { resolvePythonPath } from '../pythonPath';
import { resolveDetachedPythonPath } from '../pythonPath';
const isWindows = process.platform === 'win32';
const appendJobLog = (logPath: string, message: string) => {
@ -13,6 +13,168 @@ const appendJobLog = (logPath: string, message: string) => {
});
};
// Windows only. Launched as `node -e <this>` so the job ends up outside the
// worker's process tree: `taskkill /T` (the dev script's `concurrently -k`,
// or any shutdown that kills the tree) walks parent/child links and would take
// a direct child down with the UI. This relay exits immediately, orphaning the
// job, and `detached` keeps the job alive once its parent is gone. Its own
// stdout/stderr are the job log, so the job inherits them as fds 1 and 2.
// Python failing to launch at all (broken venv, missing interpreter) happens
// inside the relay, so the relay -- not the worker -- is what sees that error.
// It reports it two ways: on stderr, which is the job log, and through the pid
// file, so the worker can put the real reason in the database.
const RELAY_ERROR_PREFIX = 'error:';
const WINDOWS_RELAY_SCRIPT = `
const { spawn } = require('child_process');
const fs = require('fs');
const [pidFile, command, ...args] = process.argv.slice(1);
const child = spawn(command, args, {
detached: true,
windowsHide: true,
stdio: ['ignore', 1, 2],
});
child.once('error', error => {
process.stderr.write('Error launching job process: ' + error.message + '\\n');
try {
fs.writeFileSync(pidFile, '${RELAY_ERROR_PREFIX}' + error.message);
} catch (e) {
process.stderr.write('Could not write job pid file: ' + e.message + '\\n');
}
process.exit(1);
});
if (child.pid) {
fs.writeFileSync(pidFile, String(child.pid));
child.unref();
}
`;
const RELAY_PID_TIMEOUT_MS = 30000;
type RelayResult = { pid: number | null; error?: string };
// The relay exits as soon as it has launched the job, leaving the real pid in
// pidPath. Without this we would only ever know the (already dead) relay's pid.
const readRelayPid = (relay: ChildProcess, pidPath: string): Promise<RelayResult> => {
return new Promise(resolve => {
let settled = false;
const finish = (value: RelayResult) => {
if (settled) return;
settled = true;
clearTimeout(timer);
resolve(value);
};
const timer = setTimeout(
() => finish({ pid: null, error: 'Timed out waiting for the job process to start' }),
RELAY_PID_TIMEOUT_MS,
);
relay.once('exit', () => {
let contents: string;
try {
contents = fs.readFileSync(pidPath, 'utf8').trim();
} catch {
finish({ pid: null, error: 'Job process did not report a pid' });
return;
}
if (contents.startsWith(RELAY_ERROR_PREFIX)) {
finish({ pid: null, error: contents.slice(RELAY_ERROR_PREFIX.length) });
return;
}
const pid = Number(contents);
finish(
Number.isInteger(pid) && pid > 0
? { pid }
: { pid: null, error: 'Job process did not report a usable pid' },
);
});
relay.once('error', error => finish({ pid: null, error: error.message }));
});
};
const isProcessAlive = (pid: number): boolean => {
try {
process.kill(pid, 0);
return true;
} catch (e: any) {
// EPERM means it exists but belongs to someone else, which still counts.
return e?.code === 'EPERM';
}
};
// We cannot read an exit code off a process that is not our child, so pull the
// last thing it said instead -- for a job that dies on startup (bad venv,
// missing CUDA libs) that traceback line is the whole diagnosis.
const LOG_TAIL_BYTES = 4096;
const LOG_TAIL_MAX_CHARS = 300;
const readLogTail = (logPath: string): string | null => {
let fd: number | null = null;
try {
const size = fs.statSync(logPath).size;
const length = Math.min(size, LOG_TAIL_BYTES);
if (length === 0) return null;
const buffer = Buffer.alloc(length);
fd = fs.openSync(logPath, 'r');
fs.readSync(fd, buffer, 0, length, size - length);
const lines = buffer.toString('utf8').split(/\r?\n/).filter(line => line.trim() !== '');
const lastLine = lines[lines.length - 1];
if (!lastLine) return null;
return lastLine.length > LOG_TAIL_MAX_CHARS ? `${lastLine.slice(-LOG_TAIL_MAX_CHARS)}` : lastLine;
} catch {
return null;
} finally {
if (fd !== null) {
try {
fs.closeSync(fd);
} catch {
// nothing useful to do if the log handle will not close
}
}
}
};
// The job is not our child anymore, so there is no 'exit' event to listen for.
// Poll instead, so a job that dies without updating its own row (OOM kill,
// hard crash) still gets marked as an error rather than sitting on 'running'.
const JOB_POLL_INTERVAL_MS = 2000;
const watchDetachedJob = (pid: number, jobID: string, logPath: string) => {
const timer = setInterval(() => {
if (isProcessAlive(pid)) return;
clearInterval(timer);
// A stopped or completed job writes its own ending (status row + final
// log lines) via its KeyboardInterrupt/done handlers -- stay out of the
// way. Only a job that vanished while still marked 'running' died without
// getting to say anything; record that. There is no exit code to read off
// a process that is not our child, so report the last thing it logged.
const tail = readLogTail(logPath);
const message = tail
? `Job process exited unexpectedly. Last log line: ${tail}`
: 'Job process exited unexpectedly.';
void prisma.job
.updateMany({
where: { id: jobID, status: 'running' },
data: { status: 'error', info: message, pid: null },
})
.then(result => {
if (result.count > 0) appendJobLog(logPath, `\n${message}\n`);
})
.catch(updateError => {
console.error('Error updating job after process disappeared:', updateError);
});
}, JOB_POLL_INTERVAL_MS);
// Never hold the worker open on account of this poll.
if (timer.unref) timer.unref();
};
const startAndWatchJob = (job: Job) => {
// starts and watches the job asynchronously
return new Promise<void>(async (resolve, reject) => {
@ -59,7 +221,7 @@ const startAndWatchJob = (job: Job) => {
// write the config file
fs.writeFileSync(configPath, JSON.stringify(jobConfig, null, 2));
const pythonPath = resolvePythonPath();
const pythonPath = resolveDetachedPythonPath();
const runFilePath = path.join(TOOLKIT_ROOT, 'run.py');
if (!fs.existsSync(runFilePath)) {
@ -90,6 +252,9 @@ const startAndWatchJob = (job: Job) => {
const args = [runFilePath, configPath];
// Where the Windows relay reports the job's real pid back to us.
const relayPidPath = path.join(trainingFolder, '.job_pid');
let logFd: number | null = null;
try {
// Capture errors that occur before run.py can initialize file logging.
@ -97,14 +262,23 @@ const startAndWatchJob = (job: Job) => {
let subprocess;
if (isWindows) {
// Spawn Python directly on Windows so the process can survive parent exit
subprocess = spawn(pythonPath, args, {
// Launch through the relay (see WINDOWS_RELAY_SCRIPT) so the job is not
// a descendant of this worker and survives the UI being shut down or
// tree-killed. The relay spawns the job `detached`, which is what keeps
// it alive once the relay exits; that in turn means DETACHED_PROCESS,
// so pythonPath is pythonw.exe to avoid Windows handing the job a
// console window of its own.
try {
fs.unlinkSync(relayPidPath);
} catch {
// no stale pid file to clear
}
subprocess = spawn(process.execPath, ['-e', WINDOWS_RELAY_SCRIPT, relayPidPath, pythonPath, ...args], {
env: {
...process.env,
...additionalEnv,
},
cwd: TOOLKIT_ROOT,
detached: true,
windowsHide: true,
stdio: ['ignore', logFd, logFd], // don't tie stdio to parent; log fd passed as stdout and stderr
});
@ -136,25 +310,39 @@ const startAndWatchJob = (job: Job) => {
});
});
// Record abnormal termination and repair jobs Python could not update itself.
subprocess.once('exit', (code, signal) => {
if (code === 0) return;
let pid: number | null;
const result = signal ? `signal ${signal}` : `exit code ${code}`;
const message = `Job process terminated with ${result}.`;
appendJobLog(logPath, `\n${message}\n`);
void prisma.job
.updateMany({
where: { id: jobID, status: 'running' },
data: { status: 'error', info: message, pid: null },
})
.catch(updateError => {
console.error('Error updating job after abnormal process exit:', updateError);
});
});
if (isWindows) {
// The relay is gone within a few hundred ms; the pid it leaves behind is
// the job's. Poll that pid for liveness since we get no 'exit' event.
const relayResult = await readRelayPid(subprocess, relayPidPath);
if (relayResult.pid == null) {
throw new Error(relayResult.error ?? 'Job process did not report a pid');
}
pid = relayResult.pid;
watchDetachedJob(pid, jobID, logPath);
} else {
pid = subprocess.pid ?? null;
// Record abnormal termination and repair jobs Python could not update itself.
subprocess.once('exit', (code, signal) => {
if (code === 0) return;
const result = signal ? `signal ${signal}` : `exit code ${code}`;
const message = `Job process terminated with ${result}.`;
appendJobLog(logPath, `\n${message}\n`);
void prisma.job
.updateMany({
where: { id: jobID, status: 'running' },
data: { status: 'error', info: message, pid: null },
})
.catch(updateError => {
console.error('Error updating job after abnormal process exit:', updateError);
});
});
}
// Save the PID to the database and a file for future management (stop/inspect)
const pid = subprocess.pid ?? null;
if (pid != null) {
await prisma.job.update({
where: { id: jobID },

15
ui/cron/pythonPath.ts
View File

@ -25,3 +25,18 @@ export const resolvePythonPath = (): string => {
return isWindows ? 'python.exe' : 'python3';
};
// Interpreter for jobs we detach so they outlive the UI. On Windows detaching
// means DETACHED_PROCESS, and Windows gives any *console* app started that way
// a fresh console with a visible window. pythonw.exe is the GUI-subsystem
// build, so it never gets one; stdout/stderr still go to the handles we pass.
export const resolveDetachedPythonPath = (): string => {
const pythonPath = resolvePythonPath();
if (!isWindows) return pythonPath;
// Always take the pythonw next to the interpreter we already resolved, so
// both come from the same environment. Falling back to python.exe still
// works, it just shows a console window.
const pythonwPath = path.join(path.dirname(pythonPath), 'pythonw.exe');
return fs.existsSync(pythonwPath) ? pythonwPath : pythonPath;
};

4
ui/package.json
View File

@ -6,7 +6,9 @@
"dev": "concurrently -k -n WORKER,UI \"ts-node-dev --project tsconfig.worker.json --respawn --watch cron --transpile-only cron/worker.ts\" \"ts-node-dev --project tsconfig.worker.json --respawn --watch cron --transpile-only cron/fileServer.ts dev --port 3000\"",
"build": "tsc -p tsconfig.worker.json && next build",
"start": "concurrently --restart-tries -1 --restart-after 1000 -n WORKER,UI \"node dist/cron/worker.js\" \"node dist/cron/fileServer.js start --port 8675\"",
"build_and_start": "npm install && npm run update_db && npm run build && npm run start",
"install_deps": "npm install --no-save --no-audit --no-fund",
"db_build_start": "npm run update_db && npm run build && npm run start",
"build_and_start": "npm run install_deps && npm run db_build_start",
"lint": "next lint",
"update_db": "npx prisma generate && npx prisma db push",
"format": "prettier --write \"**/*.{js,jsx,ts,tsx,css,scss}\""

6
ui/src/app/api/cpu/route.ts
View File

@ -1,9 +1,9 @@
import { NextResponse } from 'next/server';
import si from 'systeminformation';
import { createRequire } from 'module';
import os from 'os';
import { CpuInfo } from '@/types';
import { cached } from '@/server/apiCache';
import { loadMacstats } from '@/server/macstats';
const isMac = os.platform() === 'darwin';
@ -13,8 +13,8 @@ async function getCpuInfo(): Promise<CpuInfo> {
if (isMac) {
try {
const nativeRequire = createRequire(import.meta.url);
const ms = nativeRequire('macstats') as any;
const ms = loadMacstats();
if (!ms) throw new Error('macstats unavailable');
const ramData = ms.getRAMUsageSync();
const cpuData = ms.getCpuDataSync();

11
ui/src/app/api/gpu/route.ts
View File

@ -1,9 +1,9 @@
import { NextResponse } from 'next/server';
import { exec } from 'child_process';
import { promisify } from 'util';
import { createRequire } from 'module';
import os from 'os';
import { cached } from '@/server/apiCache';
import { loadMacstats } from '@/server/macstats';
const execAsync = promisify(exec);
@ -47,11 +47,8 @@ async function getMacGpuInfo(): Promise<MacGpuResult | null> {
let memUsed = 0;
let memTotal = memoryTotal;
try {
// Use createRequire to hide from webpack static analysis so it doesn't fail on non-mac platforms
const nativeRequire = createRequire(import.meta.url);
const ms = nativeRequire('macstats') as any;
const ms = loadMacstats();
if (ms) {
try {
const gpuData = ms.getGpuDataSync();
temperature = gpuData.temperature || 0;
@ -84,8 +81,6 @@ async function getMacGpuInfo(): Promise<MacGpuResult | null> {
} catch {
// ignore
}
} catch (error) {
console.warn('macstats not available:', error);
}
return { name: gpuName, memUsed, memTotal, gpuLoad, temperature, fanSpeed, powerDraw };

46
ui/src/app/api/jobs/[jobID]/stop/route.ts
View File

@ -6,6 +6,43 @@ import { promisify } from 'util';
const execAsync = promisify(exec);
const isWindows = process.platform === 'win32';
// How long a stopping job gets to run its graceful shutdown (KeyboardInterrupt
// -> on_error -> final DB write) before we assume it is hung and force-kill it.
const GRACEFUL_STOP_TIMEOUT_MS = 60_000;
const BACKSTOP_POLL_MS = 2_000;
const isProcessAlive = (pid: number): boolean => {
try {
process.kill(pid, 0);
return true;
} catch (e: any) {
// EPERM means it exists but belongs to someone else, which still counts.
return e?.code === 'EPERM';
}
};
// Windows: the trainer's stop watcher sees the `stop` flag and raises SIGINT
// inside the process, so the graceful path runs without us sending anything.
// This backstop only exists for a job that is too hung to notice the flag.
// Poll (rather than one long timer) so we stop watching the moment the pid
// dies and never touch a recycled pid.
const scheduleForceKillBackstop = (pid: number, jobID: string) => {
const startedAt = Date.now();
const timer = setInterval(() => {
if (!isProcessAlive(pid)) {
clearInterval(timer);
return;
}
if (Date.now() - startedAt < GRACEFUL_STOP_TIMEOUT_MS) return;
clearInterval(timer);
console.warn(`Job ${jobID} (pid ${pid}) still alive ${GRACEFUL_STOP_TIMEOUT_MS / 1000}s after stop request, force killing`);
execAsync(`taskkill /PID ${pid} /T /F`, { windowsHide: true }).catch(() => {
// already gone
});
}, BACKSTOP_POLL_MS);
timer.unref?.();
};
export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) {
const { jobID } = await params;
@ -30,9 +67,12 @@ export async function GET(request: NextRequest, { params }: { params: { jobID: s
console.log(`Attempting to stop job ${jobID} with PID ${job.pid}`);
try {
if (isWindows) {
// Windows doesn't support SIGINT for arbitrary processes.
// Use taskkill with /T (tree) to send a CTRL+C-like termination.
await execAsync(`taskkill /PID ${job.pid} /T /F`, { windowsHide: true });
// No external SIGINT possible on Windows (the job runs under pythonw
// with no console), and none is needed: the `stop` flag written above
// is the signal. The trainer's stop watcher polls it and raises
// SIGINT in-process, which runs the same KeyboardInterrupt shutdown
// as Linux -- final DB write and the closing log lines.
scheduleForceKillBackstop(job.pid, jobID);
} else {
process.kill(job.pid, 'SIGINT');
}

38
ui/src/server/macstats.ts Normal file
View File

@ -0,0 +1,38 @@
import * as nodeModule from 'module';
import os from 'os';
import path from 'path';
/**
* macstats is a native, macOS-only optional dependency.
*
* It must not be bundled. Webpack special-cases `createRequire(...)`: with a literal base it
* resolves and inlines the module (a hard MODULE_NOT_FOUND throw when it can't), and with a
* computed base it replaces the whole call with `void 0` — both leave macstats unreachable at
* runtime even when it is installed, regardless of serverExternalPackages/externals. So look
* createRequire up dynamically, where the bundler can't recognize it, and base the require on
* process.cwd() (the ui/ directory the Next server runs in) so it resolves against
* ui/node_modules at runtime.
*/
// undefined = not tried yet, null = unavailable on this machine
let cachedModule: any | null | undefined;
export function loadMacstats(): any | null {
if (cachedModule !== undefined) return cachedModule;
if (os.platform() !== 'darwin') {
cachedModule = null;
return cachedModule;
}
try {
const createRequire = (nodeModule as any)['create' + 'Require'] as typeof nodeModule.createRequire;
const nativeRequire = createRequire(path.join(process.cwd(), 'package.json'));
cachedModule = nativeRequire('macstats');
} catch (error) {
console.warn('macstats not available:', error);
cachedModule = null;
}
return cachedModule;
}

2
version.py
View File

@ -1 +1 @@
VERSION = "0.11.2"
VERSION = "0.12.0"