Spaces:
Running on Zero
Running on Zero
File size: 4,871 Bytes
6d5047c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | #!/usr/bin/env python3
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""Regenerate `docker_requirements.txt` from `docker_requirements.in` using `uv`, targeting the
Docker image runtime, and filter out `torch` + CUDA wheels so Docker doesn't try to reinstall
PyTorch.
Usage:
python3 kimodo/scripts/lock_requirements.py
Optional args:
--python-version 3.10
--python-platform x86_64-manylinux2014
--in docker_requirements.in
--out docker_requirements.txt
"""
import argparse
import shutil
import subprocess
from pathlib import Path
from typing import Iterable
DEFAULT_PYTHON_VERSION = "3.10"
DEFAULT_PYTHON_PLATFORM = "x86_64-manylinux2014"
# Packages to omit from the lockfile because the Docker base image already provides torch+CUDA.
OMIT_NAMES = {"torch", "triton", "networkx", "sympy", "mpmath"}
OMIT_PREFIXES = ("nvidia-",)
def _run(cmd: list[str]) -> None:
print("+", " ".join(cmd))
subprocess.run(cmd, check=True)
def _ensure_uv() -> None:
if shutil.which("uv") is None:
raise SystemExit(
"ERROR: `uv` is not installed or not on PATH.\n"
"Install it (one of):\n"
" - pipx install uv\n"
" - python -m pip install --user uv\n"
"Then rerun this script."
)
def _parse_req_name(line: str) -> str:
# uv emits `name==version` lines.
s = line.strip()
if "==" in s:
return s.split("==", 1)[0].strip()
# Fallback: treat the whole token before space as name.
return s.split()[0].strip()
def _iter_blocks(lines: list[str]) -> Iterable[list[str]]:
"""Split a docker_requirements.txt into blocks: [top-level req line + indented comments]."""
i = 0
n = len(lines)
while i < n:
line = lines[i]
# Header/comments/blank
if line.startswith("#") or line.strip() == "":
yield [line]
i += 1
continue
# Top-level requirement line
if not line.startswith(" "):
block = [line]
i += 1
while i < n and (lines[i].startswith(" ") or lines[i].strip() == "" or lines[i].startswith("#")):
# Stop if we hit another top-level requirement line
if not lines[i].startswith(" ") and not lines[i].startswith("#") and lines[i].strip() != "":
break
block.append(lines[i])
i += 1
yield block
continue
# Indented line without a requirement header (shouldn't happen, but keep)
yield [line]
i += 1
def _should_omit(req_line: str) -> bool:
name = _parse_req_name(req_line)
if name in OMIT_NAMES:
return True
for pfx in OMIT_PREFIXES:
if name.startswith(pfx):
return True
return False
def filter_lockfile(path: Path) -> None:
lines = path.read_text(encoding="utf-8").splitlines(True)
out: list[str] = []
inserted_note = False
for block in _iter_blocks(lines):
first = block[0]
# After the uv header lines, insert a short note once.
if (not inserted_note) and first.startswith("# This file was autogenerated by uv"):
out.extend(block)
out.append(
"# NOTE: `torch` (and its CUDA wheels) are intentionally omitted from this lockfile.\n"
"# The Docker base image (nvcr.io/nvidia/pytorch) already provides a tested PyTorch build.\n"
"#\n"
)
inserted_note = True
continue
if first.startswith("#") or first.strip() == "":
out.extend(block)
continue
if _should_omit(first):
continue
out.extend(block)
path.write_text("".join(out), encoding="utf-8")
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--in", dest="in_file", default="docker_requirements.in")
ap.add_argument("--out", dest="out_file", default="docker_requirements.txt")
ap.add_argument("--python-version", default=DEFAULT_PYTHON_VERSION)
ap.add_argument("--python-platform", default=DEFAULT_PYTHON_PLATFORM)
args = ap.parse_args()
_ensure_uv()
in_path = Path(args.in_file)
out_path = Path(args.out_file)
if not in_path.exists():
raise SystemExit(f"ERROR: missing {in_path}")
_run(
[
"uv",
"pip",
"compile",
"-U",
str(in_path),
"-o",
str(out_path),
"--python-version",
args.python_version,
"--python-platform",
args.python_platform,
]
)
filter_lockfile(out_path)
print(f"OK: wrote {out_path} (filtered torch/CUDA wheels)")
if __name__ == "__main__":
main()
|