Buckets:
| """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.") | |
Xet Storage Details
- Size:
- 25.6 kB
- Xet hash:
- 79e2bb4a6c49ce31e17e0f843e1af3e3cb6aa0ced938a5aaef0d72491f8086eb
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.