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

708 lines
25 KiB
Python
Raw Permalink Blame History

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