diff --git a/.gitignore b/.gitignore index 558cf816..a6f4b109 100644 --- a/.gitignore +++ b/.gitignore @@ -124,6 +124,9 @@ celerybeat.pid .venv .python .node +.ffmpeg +.mingit +.uv env/ venv/ ENV/ diff --git a/README.md b/README.md index a27d9376..5ce450cf 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/dgx_instructions.md b/dgx_instructions.md index 06198aa8..2aa60a56 100644 --- a/dgx_instructions.md +++ b/dgx_instructions.md @@ -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 ``` diff --git a/extensions_built_in/captioner/AceStepCaptioner.py b/extensions_built_in/captioner/AceStepCaptioner.py index ce89a2f9..e1d77b31 100644 --- a/extensions_built_in/captioner/AceStepCaptioner.py +++ b/extensions_built_in/captioner/AceStepCaptioner.py @@ -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) diff --git a/extensions_built_in/captioner/BaseCaptioner.py b/extensions_built_in/captioner/BaseCaptioner.py index 6100fbbd..3c8f298b 100644 --- a/extensions_built_in/captioner/BaseCaptioner.py +++ b/extensions_built_in/captioner/BaseCaptioner.py @@ -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: diff --git a/extensions_built_in/sd_trainer/DiffusionTrainer.py b/extensions_built_in/sd_trainer/DiffusionTrainer.py index f7d2ad6d..ef9d58fa 100644 --- a/extensions_built_in/sd_trainer/DiffusionTrainer.py +++ b/extensions_built_in/sd_trainer/DiffusionTrainer.py @@ -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) diff --git a/extensions_built_in/sd_trainer/UITrainer.py b/extensions_built_in/sd_trainer/UITrainer.py index 8b5fa796..f290fc9e 100644 --- a/extensions_built_in/sd_trainer/UITrainer.py +++ b/extensions_built_in/sd_trainer/UITrainer.py @@ -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()) diff --git a/manager/README.md b/manager/README.md new file mode 100644 index 00000000..3e856fb0 --- /dev/null +++ b/manager/README.md @@ -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. diff --git a/manager/__init__.py b/manager/__init__.py new file mode 100644 index 00000000..c3e83ae8 --- /dev/null +++ b/manager/__init__.py @@ -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 diff --git a/manager/__main__.py b/manager/__main__.py new file mode 100644 index 00000000..ba9ce68d --- /dev/null +++ b/manager/__main__.py @@ -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()) diff --git a/manager/detect.py b/manager/detect.py new file mode 100644 index 00000000..48781572 --- /dev/null +++ b/manager/detect.py @@ -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 diff --git a/manager/doctor.py b/manager/doctor.py new file mode 100644 index 00000000..554c7e2c --- /dev/null +++ b/manager/doctor.py @@ -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 ""), + ) diff --git a/manager/env.py b/manager/env.py new file mode 100644 index 00000000..b7decb9d --- /dev/null +++ b/manager/env.py @@ -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@ + 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.") diff --git a/manager/ffmpeg.py b/manager/ffmpeg.py new file mode 100644 index 00000000..437d9145 --- /dev/null +++ b/manager/ffmpeg.py @@ -0,0 +1,179 @@ +"""Local (never global) FFmpeg provisioning into /.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 diff --git a/manager/gitops.py b/manager/gitops.py new file mode 100644 index 00000000..4007123d --- /dev/null +++ b/manager/gitops.py @@ -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." + ) diff --git a/manager/gitwin.py b/manager/gitwin.py new file mode 100644 index 00000000..1ec39dd0 --- /dev/null +++ b/manager/gitwin.py @@ -0,0 +1,76 @@ +"""Local (never global) Git provisioning for Windows into /.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) diff --git a/manager/launch.py b/manager/launch.py new file mode 100644 index 00000000..d94aca88 --- /dev/null +++ b/manager/launch.py @@ -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() diff --git a/manager/migrations.py b/manager/migrations.py new file mode 100644 index 00000000..fb420320 --- /dev/null +++ b/manager/migrations.py @@ -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) diff --git a/manager/nodejs.py b/manager/nodejs.py new file mode 100644 index 00000000..787a0388 --- /dev/null +++ b/manager/nodejs.py @@ -0,0 +1,245 @@ +"""Local (never global) Node.js provisioning into /.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 diff --git a/manager/sparkdeps.py b/manager/sparkdeps.py new file mode 100644 index 00000000..f5062ee0 --- /dev/null +++ b/manager/sparkdeps.py @@ -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() + # \bin\arm64 -> + 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).") diff --git a/manager/spec.py b/manager/spec.py new file mode 100644 index 00000000..74c39796 --- /dev/null +++ b/manager/spec.py @@ -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."], + ) diff --git a/manager/util.py b/manager/util.py new file mode 100644 index 00000000..cc944c1c --- /dev/null +++ b/manager/util.py @@ -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() diff --git a/manager/uvbin.py b/manager/uvbin.py new file mode 100644 index 00000000..4ff47640 --- /dev/null +++ b/manager/uvbin.py @@ -0,0 +1,78 @@ +"""Local (never system) uv provisioning into /.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-/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) diff --git a/requirements_base.txt b/requirements_base.txt index dcfba99c..30f94518 100644 --- a/requirements_base.txt +++ b/requirements_base.txt @@ -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 diff --git a/run.py b/run.py index a77a792a..77b6b91d 100644 --- a/run.py +++ b/run.py @@ -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 diff --git a/run_linux.sh b/run_linux.sh new file mode 100755 index 00000000..fb379604 --- /dev/null +++ b/run_linux.sh @@ -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 diff --git a/run_mac.zsh b/run_mac.zsh index 8f7b49e9..5025b38d 100755 --- a/run_mac.zsh +++ b/run_mac.zsh @@ -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 diff --git a/run_modal.py b/run_modal.py index 4675c1cb..52730456 100644 --- a/run_modal.py +++ b/run_modal.py @@ -55,7 +55,6 @@ image = ( "toml", "pydantic", "omegaconf", - "k-diffusion", "open_clip_torch", "timm", "prodigyopt", diff --git a/run_windows.bat b/run_windows.bat new file mode 100644 index 00000000..756bd017 --- /dev/null +++ b/run_windows.bat @@ -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 diff --git a/scripts/test_quantizations.py b/scripts/test_quantizations.py index 445edaaa..a997cd16 100644 --- a/scripts/test_quantizations.py +++ b/scripts/test_quantizations.py @@ -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) diff --git a/spark_requirements.txt b/spark_requirements.txt new file mode 100644 index 00000000..55c6c147 --- /dev/null +++ b/spark_requirements.txt @@ -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 diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index a0ef9717..0726f52d 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -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) diff --git a/toolkit/data_loader.py b/toolkit/data_loader.py index ee90e263..b0737601 100644 --- a/toolkit/data_loader.py +++ b/toolkit/data_loader.py @@ -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 diff --git a/toolkit/data_transfer_object/data_loader.py b/toolkit/data_transfer_object/data_loader.py index d07a6d6d..103b1733 100644 --- a/toolkit/data_transfer_object/data_loader.py +++ b/toolkit/data_transfer_object/data_loader.py @@ -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") diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py index 0c5ac31d..6adee1bf 100644 --- a/toolkit/dataloader_mixins.py +++ b/toolkit/dataloader_mixins.py @@ -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 diff --git a/toolkit/util/convrot_quant.py b/toolkit/util/convrot_quant.py index 92935075..8b7234cf 100644 --- a/toolkit/util/convrot_quant.py +++ b/toolkit/util/convrot_quant.py @@ -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) diff --git a/toolkit/win_console.py b/toolkit/win_console.py new file mode 100644 index 00000000..d7eca0f7 --- /dev/null +++ b/toolkit/win_console.py @@ -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 diff --git a/ui/cron/actions/startJob.ts b/ui/cron/actions/startJob.ts index 8beb1861..0312b038 100644 --- a/ui/cron/actions/startJob.ts +++ b/ui/cron/actions/startJob.ts @@ -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 ` 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 => { + 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(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 }, diff --git a/ui/cron/pythonPath.ts b/ui/cron/pythonPath.ts index 64ac231a..05187e50 100644 --- a/ui/cron/pythonPath.ts +++ b/ui/cron/pythonPath.ts @@ -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; +}; diff --git a/ui/package.json b/ui/package.json index b799bba4..0b72a0b8 100644 --- a/ui/package.json +++ b/ui/package.json @@ -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}\"" diff --git a/ui/src/app/api/cpu/route.ts b/ui/src/app/api/cpu/route.ts index 12eab7a1..ee9f388d 100644 --- a/ui/src/app/api/cpu/route.ts +++ b/ui/src/app/api/cpu/route.ts @@ -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 { 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(); diff --git a/ui/src/app/api/gpu/route.ts b/ui/src/app/api/gpu/route.ts index ffb47c26..643357e7 100644 --- a/ui/src/app/api/gpu/route.ts +++ b/ui/src/app/api/gpu/route.ts @@ -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 { 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 { } catch { // ignore } - } catch (error) { - console.warn('macstats not available:', error); } return { name: gpuName, memUsed, memTotal, gpuLoad, temperature, fanSpeed, powerDraw }; diff --git a/ui/src/app/api/jobs/[jobID]/stop/route.ts b/ui/src/app/api/jobs/[jobID]/stop/route.ts index cb14ad6b..abfd11b9 100644 --- a/ui/src/app/api/jobs/[jobID]/stop/route.ts +++ b/ui/src/app/api/jobs/[jobID]/stop/route.ts @@ -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'); } diff --git a/ui/src/server/macstats.ts b/ui/src/server/macstats.ts new file mode 100644 index 00000000..d8326aa2 --- /dev/null +++ b/ui/src/server/macstats.ts @@ -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; +} diff --git a/version.py b/version.py index 2073d43e..c1830e41 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.11.2" +VERSION = "0.12.0"