mirror of
https://github.com/radixark/miles.git
synced 2026-10-02 07:14:53 +08:00
fix(ci): preserve image cuDNN during dependency resolution (#3880)
This commit is contained in:
Jiajun Li
@@ -156,27 +156,8 @@ jobs:
|
||||
- name: Reconcile Miles dependencies
|
||||
shell: bash
|
||||
run: |
|
||||
# The image pins cuDNN deliberately: transformer_engine's fused attention needs a
|
||||
# newer build than torch's own exact pin, so the Dockerfile installs it last. This
|
||||
# resolve drags it back down to whatever torch asks for, which makes fused attention
|
||||
# backward fail with CUDNN_STATUS_BAD_PARAM. Record what the image shipped and put it
|
||||
# back afterwards; --no-deps so nothing else in the resolved set moves.
|
||||
pinned=""
|
||||
for pkg in nvidia-cudnn-cu13 nvidia-cudnn-cu12; do
|
||||
v=$(python -m pip show "$pkg" 2>/dev/null | awk '/^Version:/{print $2}' || true)
|
||||
if [ -n "$v" ]; then pinned="$pinned $pkg==$v"; fi
|
||||
done
|
||||
python -m pip install -r "$GITHUB_WORKSPACE/examples/multi_lora/requirements.txt" --break-system-packages
|
||||
python -m pip install -r "$GITHUB_WORKSPACE/requirements.txt" --break-system-packages
|
||||
for spec in $pinned; do
|
||||
pkg=${spec%%==*}
|
||||
want=${spec##*==}
|
||||
got=$(python -m pip show "$pkg" 2>/dev/null | awk '/^Version:/{print $2}' || true)
|
||||
if [ -n "$got" ] && [ "$got" != "$want" ]; then
|
||||
echo "Restoring image pin: $pkg $got -> $want"
|
||||
python -m pip install --no-deps --break-system-packages "$spec"
|
||||
fi
|
||||
done
|
||||
python tests/ci/reconcile_dependencies.py \
|
||||
examples/multi_lora/requirements.txt requirements.txt
|
||||
|
||||
- name: Sync dependency sources
|
||||
if: ${{ !inputs.skip_dependency_install }}
|
||||
|
||||
@@ -74,7 +74,9 @@ A test normally runs at its home stage. A [dispatch label](/developer/ci/01-labe
|
||||
|
||||
Without a dispatch label nothing leaves its home stage, so scheduled and called runs are unaffected. Whatever stage actually executes a run is also the stage its performance baseline is keyed on, keeping the generations' numbers apart.
|
||||
|
||||
**Dependency boundary.** CUDA stages start from dependencies baked into `radixark/miles`, install CUDA test dependencies from `examples/multi_lora/requirements.txt`, override them with Miles runtime dependencies from `requirements.txt`, restore the image’s cuDNN pins, update the SGLang and Megatron-LM checkouts to the selected refs, and expose all three source trees through `PYTHONPATH`; they do not rebuild or install the Miles, SGLang, or Megatron-LM source trees after the container starts. The hosted CPU stages install dependencies from `requirements.txt` and the fully pinned `tests/ci/requirements-ci-cpu.txt`, then expose the Miles, SGLang, and Megatron-LM source trees through `PYTHONPATH` without editable installs or inline package lists. The ROCm stage instead uses the SGLang and Megatron-LM versions baked into `rocm/sgl-dev`, unless the run names a ref for either.
|
||||
**Dependency boundary.** CUDA stages start from dependencies baked into `radixark/miles` and install `examples/multi_lora/requirements.txt` before `requirements.txt` through `tests/ci/reconcile_dependencies.py`. Installed CUDA 12/13 cuDNN versions are preserved with uv overrides; images without cuDNN use pip. Before each install, a dry-run gate rejects replacement of installed GPU runtimes or packages with at least 100 MiB of recorded installed files, including same-version reinstalls. These packages must be updated in the image.
|
||||
|
||||
CUDA stages then update the SGLang and Megatron-LM checkouts to the selected refs and expose all three source trees through `PYTHONPATH`; they do not rebuild or install those source trees after the container starts. Hosted CPU stages install `requirements.txt` and the fully pinned `tests/ci/requirements-ci-cpu.txt`, then expose the source trees through `PYTHONPATH` without editable installs or inline package lists. The ROCm stage uses the SGLang and Megatron-LM versions baked into `rocm/sgl-dev`, unless the run names a ref for either.
|
||||
|
||||
CUDA and CPU dependency refs resolve in this order: explicit dispatch input or PR-body directive, committed `release-lock.json`, then the moving `sglang-miles` / `miles-main` branch heads. A called release run therefore checks out its requested Miles `ref` and consumes the lockfile on that ref unless an explicit override exists. ROCm checks out the requested Miles ref but keeps the dependencies baked into its image unless the run names a ref for one.
|
||||
|
||||
|
||||
@@ -114,10 +114,7 @@ Official versioned releases add `radixark/miles:v<exact-version>` for the CUDA 1
|
||||
This is the part that decides whether your change needs a new image. A CUDA CI job starts
|
||||
from `radixark/miles:<tag>` and then:
|
||||
|
||||
1. Runs `pip install -r requirements.txt`, then restores the image's own cuDNN pin. The
|
||||
restore is not cosmetic: TE's fused attention needs a newer cuDNN than torch pins, a
|
||||
plain resolve drags it back down, and the symptom is a fused-attention backward failing
|
||||
with `CUDNN_STATUS_BAD_PARAM`.
|
||||
1. Installs `examples/multi_lora/requirements.txt` before `requirements.txt` through `tests/ci/reconcile_dependencies.py`. It preserves installed CUDA 12/13 cuDNN versions with uv overrides because TE needs a newer cuDNN than torch's exact pin; images without cuDNN use pip. Before each install, a dry-run gate rejects replacement of installed GPU runtimes or packages with at least 100 MiB of recorded installed files, including same-version reinstalls. Update these packages in the image.
|
||||
2. Resets both dependency checkouts and fetches the selected refs. Explicit dispatch or PR-body overrides win first, `release-lock.json` commits win when no override exists, and the moving `sglang-miles` / `miles-main` heads are the final defaults.
|
||||
3. Sets `PYTHONPATH` to the Miles workspace plus both source roots.
|
||||
|
||||
@@ -125,7 +122,7 @@ It never reinstalls the three source trees, because they are editable installs.
|
||||
|
||||
| Your change | Needs a new image? |
|
||||
|---|---|
|
||||
| `requirements.txt` | No. The next CI run installs it. |
|
||||
| `requirements.txt` | Only when replacing an installed GPU runtime or a package with at least 100 MiB of recorded installed files; other changes install in the next CI run. |
|
||||
| A Dockerfile layer: a pinned wheel, an inline commit, a TE patch, the base image | Yes |
|
||||
| SGLang or Megatron-LM code | No. Point CI at a ref instead |
|
||||
| Miles code | No |
|
||||
@@ -134,10 +131,7 @@ The ROCm stage is the exception: it takes SGLang and Megatron-LM from `rocm/sgl-
|
||||
|
||||
## Bumping principle
|
||||
|
||||
**Bump where the pin lives, exactly once.** A Python dependency moves in
|
||||
`requirements.txt`; an image layer moves in `docker/Dockerfile`; a variant-only difference
|
||||
moves in `docker/build.py`. If a bump needs edits in two of the three, one of them is in the
|
||||
wrong place.
|
||||
**Bump where the pin lives.** Python dependency requirements live in `requirements.txt`; image layers live in `docker/Dockerfile`; variant-only differences live in `docker/build.py`. A requirement change that replaces an installed GPU runtime or a package with at least 100 MiB of recorded installed files also needs an image rebuild so the dependency gate can retain the new image version.
|
||||
|
||||
**Prefer moving the branch to pinning a commit during rolling development.** `SGLANG_COMMIT` and `MEGATRON_COMMIT` are empty by default, so ordinary images follow `sglang-miles` and `miles-main` together. A versioned release is the deliberate exception: its lockfile supplies both exact commits to CI and the final image build.
|
||||
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
"""Install CI requirements while retaining the image's cuDNN runtime."""
|
||||
|
||||
import argparse
|
||||
import importlib.metadata
|
||||
import json
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
LARGE_PACKAGE_BYTES = 100 * 1024 * 1024
|
||||
RUNTIME_PACKAGES = {"torch", "triton", "apex", "sglang-kernel"}
|
||||
RUNTIME_PREFIXES = ("nvidia-", "flash-attn", "flashinfer-", "transformer-engine")
|
||||
|
||||
|
||||
def package_name(name: str) -> str:
|
||||
return re.sub(r"[-_.]+", "-", name).lower()
|
||||
|
||||
|
||||
def protected_packages() -> set[str]:
|
||||
protected = set()
|
||||
for distribution in importlib.metadata.distributions():
|
||||
name = package_name(distribution.metadata["Name"])
|
||||
size = sum(file.size or 0 for file in distribution.files or [])
|
||||
if name in RUNTIME_PACKAGES or name.startswith(RUNTIME_PREFIXES) or size >= LARGE_PACKAGE_BYTES:
|
||||
protected.add(name)
|
||||
return protected
|
||||
|
||||
|
||||
def uv_changes(output: str) -> set[str]:
|
||||
changes = re.findall(r"^ ([+-]) ([\w.-]+)(?:==| @ ).+$", output, re.MULTILINE)
|
||||
counts = re.findall(r"^Would (uninstall|install) (\d+) packages?$", output, re.MULTILINE)
|
||||
if not counts:
|
||||
if not changes and "Would make no changes" in output:
|
||||
return set()
|
||||
raise RuntimeError("Unrecognized uv install plan; refusing to install")
|
||||
expected = dict(counts)
|
||||
if len(expected) != len(counts) or any(
|
||||
sum(sign == symbol for sign, _ in changes) != int(expected.get(action, 0))
|
||||
for action, symbol in (("install", "+"), ("uninstall", "-"))
|
||||
):
|
||||
raise RuntimeError("Incomplete uv install plan; refusing to install")
|
||||
return {package_name(name) for _, name in changes}
|
||||
|
||||
|
||||
def check_install_plan(command: list[str], directory: Path, *, use_uv: bool) -> None:
|
||||
protected = protected_packages()
|
||||
if use_uv:
|
||||
result = subprocess.run([*command, "--dry-run", "--color", "never"], capture_output=True, text=True)
|
||||
print(result.stdout + result.stderr, end="", flush=True)
|
||||
result.check_returncode()
|
||||
changes = uv_changes(result.stdout + result.stderr)
|
||||
else:
|
||||
report = directory / "pip-plan.json"
|
||||
subprocess.run([*command, "--dry-run", "--report", str(report)], check=True)
|
||||
plan = json.loads(report.read_text())
|
||||
if plan["version"] != "1":
|
||||
raise RuntimeError("Unrecognized pip install plan; refusing to install")
|
||||
changes = {package_name(item["metadata"]["name"]) for item in plan["install"]}
|
||||
|
||||
# Do not remove or bypass this gate: CI must never reinstall image runtimes or large packages.
|
||||
# Update these packages in the Docker image instead, even for same-version reinstalls.
|
||||
if blocked := sorted(protected & changes):
|
||||
raise RuntimeError(f"CI dependency gate blocks reinstalling: {', '.join(blocked)}. Rebuild the image instead.")
|
||||
|
||||
|
||||
def reconcile(requirements: list[str]) -> None:
|
||||
pins = []
|
||||
for package in ("nvidia-cudnn-cu12", "nvidia-cudnn-cu13"):
|
||||
try:
|
||||
version = importlib.metadata.version(package)
|
||||
except importlib.metadata.PackageNotFoundError:
|
||||
continue
|
||||
pins.append(f"{package}=={version}")
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix="miles-ci-dependencies-") as directory:
|
||||
if pins:
|
||||
# TE needs the image's cuDNN even when torch declares an older exact pin.
|
||||
# An override replaces that pin; a constraint would only make it conflict.
|
||||
overrides = Path(directory) / "overrides.txt"
|
||||
overrides.write_text("\n".join(pins) + "\n")
|
||||
command = ["uv", "pip", "install", "--python", sys.executable, "--overrides", str(overrides)]
|
||||
print(f"Preserving image cuDNN: {', '.join(pins)}", flush=True)
|
||||
else:
|
||||
command = [sys.executable, "-m", "pip", "install"]
|
||||
|
||||
for path in requirements:
|
||||
install = [*command, "--break-system-packages", "-r", path]
|
||||
check_install_plan(install, Path(directory), use_uv=bool(pins))
|
||||
subprocess.run(install, check=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("requirements", nargs="+")
|
||||
reconcile(parser.parse_args().requirements)
|
||||
@@ -0,0 +1,159 @@
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
from tests.ci import reconcile_dependencies
|
||||
from tests.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=20, suite="stage-a-cpu", labels=[])
|
||||
|
||||
SCRIPT = Path(__file__).resolve().parents[1] / "reconcile_dependencies.py"
|
||||
|
||||
|
||||
def wheel(directory, name, version, requires=(), payload_bytes=0):
|
||||
normalized = name.replace("-", "_")
|
||||
info = f"{normalized}-{version}.dist-info"
|
||||
with zipfile.ZipFile(
|
||||
directory / f"{normalized}-{version}-py3-none-any.whl", "w", compression=zipfile.ZIP_DEFLATED
|
||||
) as archive:
|
||||
archive.writestr(f"{normalized}/__init__.py", f'__version__ = "{version}"\n')
|
||||
archive.writestr(
|
||||
f"{info}/METADATA",
|
||||
f"Metadata-Version: 2.1\nName: {name}\nVersion: {version}\n"
|
||||
+ "".join(f"Requires-Dist: {requirement}\n" for requirement in requires),
|
||||
)
|
||||
archive.writestr(f"{info}/WHEEL", "Wheel-Version: 1.0\nRoot-Is-Purelib: true\nTag: py3-none-any\n")
|
||||
if payload_bytes:
|
||||
archive.writestr(f"{normalized}/weights.bin", bytes(payload_bytes))
|
||||
archive.writestr(
|
||||
f"{info}/RECORD",
|
||||
"".join(f"{item.filename},,{item.file_size}\n" for item in archive.infolist()) + f"{info}/RECORD,,\n",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("cudnn", ["nvidia-cudnn-cu12", "nvidia-cudnn-cu13", None])
|
||||
def test_real_resolver_preserves_image_runtime_and_installs_dependencies(tmp_path, cudnn):
|
||||
wheels = tmp_path / "wheels"
|
||||
wheels.mkdir()
|
||||
wheel(wheels, "ci-leaf", "1.0")
|
||||
wheel(wheels, "ci-client", "1.0", ["torch==1.0", "ci-leaf==1.0"])
|
||||
wheel(wheels, "torch", "1.0", [f"{cudnn}==9.20.0.48"] if cudnn else [])
|
||||
if cudnn:
|
||||
wheel(wheels, cudnn, "9.20.0.48")
|
||||
wheel(wheels, cudnn, "9.22.0.52")
|
||||
|
||||
environment = tmp_path / "venv"
|
||||
subprocess.run([sys.executable, "-m", "venv", str(environment)], check=True)
|
||||
python = str(environment / "bin" / "python")
|
||||
env = {
|
||||
**os.environ,
|
||||
"PIP_NO_INDEX": "1",
|
||||
"PIP_FIND_LINKS": str(wheels),
|
||||
"PIP_DISABLE_PIP_VERSION_CHECK": "1",
|
||||
"UV_NO_INDEX": "1",
|
||||
"UV_FIND_LINKS": str(wheels),
|
||||
}
|
||||
seeded = ["torch==1.0"] + ([f"{cudnn}==9.22.0.52"] if cudnn else [])
|
||||
subprocess.run([python, "-m", "pip", "install", "--no-deps", *seeded], env=env, check=True)
|
||||
requirements = tmp_path / "requirements.txt"
|
||||
requirements.write_text("ci-client==1.0\n")
|
||||
|
||||
if cudnn:
|
||||
baseline = subprocess.run(
|
||||
[python, "-m", "pip", "install", "--dry-run", "-r", str(requirements)],
|
||||
env=env,
|
||||
check=True,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
assert f"{cudnn}-9.20.0.48" in baseline.stdout
|
||||
|
||||
subprocess.run([python, str(SCRIPT), str(requirements)], env=env, check=True)
|
||||
installed = json.loads(subprocess.check_output([python, "-m", "pip", "list", "--format=json"], env=env, text=True))
|
||||
versions = {package["name"]: package["version"] for package in installed}
|
||||
assert versions["ci-client"] == versions["ci-leaf"] == "1.0"
|
||||
if cudnn:
|
||||
assert versions[cudnn] == "9.22.0.52"
|
||||
else:
|
||||
assert not any(name.startswith("nvidia-cudnn") for name in versions)
|
||||
|
||||
wheel(wheels, "ci-leaf", "2.0")
|
||||
requirements.write_text("ci-leaf==2.0\n")
|
||||
subprocess.run([python, str(SCRIPT), str(requirements)], env=env, check=True)
|
||||
|
||||
wheel(wheels, "torch", "2.0")
|
||||
requirements.write_text("torch==2.0\n")
|
||||
blocked = subprocess.run([python, str(SCRIPT), str(requirements)], env=env, capture_output=True, text=True)
|
||||
assert blocked.returncode != 0
|
||||
assert "CI dependency gate blocks reinstalling: torch" in blocked.stderr
|
||||
assert (
|
||||
subprocess.check_output(
|
||||
[python, "-c", "import importlib.metadata as m; print(m.version('torch'))"], text=True
|
||||
).strip()
|
||||
== "1.0"
|
||||
)
|
||||
|
||||
runtime = cudnn or "torch"
|
||||
version = "9.22.0.52" if cudnn else "1.0"
|
||||
command = (
|
||||
["uv", "pip", "install", "--python", python, "--reinstall"]
|
||||
if cudnn
|
||||
else [python, "-m", "pip", "install", "--force-reinstall"]
|
||||
)
|
||||
command += ["--no-deps", f"{runtime}=={version}"]
|
||||
code = (
|
||||
f"import runpy; from pathlib import Path; m=runpy.run_path({str(SCRIPT)!r}); "
|
||||
f"m['check_install_plan']({command!r}, Path({str(tmp_path)!r}), use_uv={bool(cudnn)!r})"
|
||||
)
|
||||
blocked = subprocess.run([python, "-c", code], env=env, capture_output=True, text=True)
|
||||
assert blocked.returncode != 0
|
||||
assert f"CI dependency gate blocks reinstalling: {runtime}" in blocked.stderr
|
||||
|
||||
if not cudnn:
|
||||
wheel(wheels, "ci-unlisted-large", "1.0", payload_bytes=100 * 1024**2)
|
||||
wheel(wheels, "ci-unlisted-large", "2.0")
|
||||
subprocess.run([python, "-m", "pip", "install", "ci-unlisted-large==1.0"], env=env, check=True)
|
||||
requirements.write_text("ci-unlisted-large==2.0\n")
|
||||
blocked = subprocess.run([python, str(SCRIPT), str(requirements)], env=env, capture_output=True, text=True)
|
||||
assert blocked.returncode != 0
|
||||
assert "CI dependency gate blocks reinstalling: ci-unlisted-large" in blocked.stderr
|
||||
|
||||
requirements.write_text("ci-missing-dependency==1.0\n")
|
||||
failed = subprocess.run([python, str(SCRIPT), str(requirements)], env=env)
|
||||
assert failed.returncode != 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size, protected", [(100 * 1024**2 - 1, False), (100 * 1024**2, True)])
|
||||
def test_large_packages_are_protected_without_a_name_allowlist(monkeypatch, size, protected):
|
||||
distribution = SimpleNamespace(metadata={"Name": "Unknown_Large.Package"}, files=[SimpleNamespace(size=size)])
|
||||
monkeypatch.setattr(reconcile_dependencies.importlib.metadata, "distributions", lambda: [distribution])
|
||||
assert ("unknown-large-package" in reconcile_dependencies.protected_packages()) == protected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"output",
|
||||
["", "Resolved 1 package", "Would install 2 packages\n + torch==2.0\n", "Would uninstall 1 package\n"],
|
||||
)
|
||||
def test_unrecognized_or_truncated_plans_cannot_pass_the_gate(output):
|
||||
with pytest.raises(RuntimeError, match="refusing to install"):
|
||||
reconcile_dependencies.uv_changes(output)
|
||||
|
||||
|
||||
def test_gpu_workflow_uses_the_gated_installer():
|
||||
import yaml
|
||||
|
||||
workflow = Path(__file__).resolve().parents[3] / ".github/workflows/_run-ci.yml"
|
||||
steps = yaml.safe_load(workflow.read_text())["jobs"]["run"]["steps"]
|
||||
reconcile = next(step for step in steps if step.get("name") == "Reconcile Miles dependencies")
|
||||
assert reconcile["run"].split() == [
|
||||
"python",
|
||||
"tests/ci/reconcile_dependencies.py",
|
||||
"\\",
|
||||
"examples/multi_lora/requirements.txt",
|
||||
"requirements.txt",
|
||||
]
|
||||
Reference in New Issue
Block a user