Centralize Megatron dependency install

Author: Kovbo

src/art/megatron/setup.sh

@@ -7,9 +7,6 @@ export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-9.0}"
 apt-get update
 apt-get install -y libcudnn9-headers-cuda-12 ninja-build
 
-# Python dependencies are declared in pyproject.toml extras.
-# Keep backend + megatron together so setup does not prune runtime deps (e.g. vllm).
-script_dir="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
-repo_root="$(cd -- "${script_dir}/../../.." && pwd)"
-cd "${repo_root}"
-uv sync --extra backend --extra megatron --frozen --active
+# Python dependencies are installed through art_megatron_install so
+# downstream repos can reuse the same source of truth for versions and VCS pins.
+uv run python -m art_megatron_install

src/art_megatron_install.py

@@ -0,0 +1,242 @@
+from __future__ import annotations
+
+import argparse
+import importlib.metadata as metadata
+import json
+import os
+import shutil
+import subprocess
+from dataclasses import dataclass
+from urllib.parse import parse_qs
+
+from packaging.requirements import Requirement
+from packaging.version import Version
+
+_APEX_ENV = {
+    "APEX_CPP_EXT": "1",
+    "APEX_CUDA_EXT": "1",
+    "APEX_FAST_LAYER_NORM": "1",
+    "APEX_PARALLEL_BUILD": "16",
+    "NVCC_APPEND_FLAGS": "--threads 4",
+}
+_TRANSFORMER_ENGINE_ENV = {"NVTE_NO_LOCAL_VERSION": "1"}
+_MEGATRON_REQUIREMENTS = {
+    "torch": Requirement("torch>=2.8.0"),
+    "quack-kernels": Requirement("quack-kernels==0.2.5"),
+    "apex": Requirement("apex @ git+https://github.com/NVIDIA/apex.git@25.09"),
+    "transformer-engine": Requirement("transformer-engine==2.11.0"),
+    "transformer-engine-cu12": Requirement("transformer-engine-cu12==2.11.0"),
+    "transformer-engine-torch": Requirement(
+        "transformer-engine-torch @ git+https://github.com/NVIDIA/TransformerEngine.git@v2.11#subdirectory=transformer_engine/pytorch"
+    ),
+    "megatron-core": Requirement("megatron-core==0.16.1"),
+    "pybind11": Requirement("pybind11>=2.13.6"),
+    "megatron-bridge": Requirement(
+        "megatron-bridge @ git+https://github.com/NVIDIA-NeMo/Megatron-Bridge.git@75f2c5ad4afb702b57b4781a00f5291a66bcf183"
+    ),
+    "nvidia-ml-py": Requirement("nvidia-ml-py==13.580.82"),
+    "ml-dtypes": Requirement("ml-dtypes>=0.5.0"),
+}
+_CORE_INSTALL_ORDER = (
+    "torch",
+    "pybind11",
+    "ml-dtypes",
+    "quack-kernels",
+    "transformer-engine",
+    "transformer-engine-cu12",
+    "transformer-engine-torch",
+    "megatron-core",
+)
+
+
+@dataclass(frozen=True)
+class _DirectUrlParts:
+    vcs: str | None
+    url: str
+    requested_revision: str | None
+    subdirectory: str | None
+
+
+def _megatron_requirements() -> dict[str, Requirement]:
+    return dict(_MEGATRON_REQUIREMENTS)
+
+
+def _format_requirement(requirement: Requirement) -> str:
+    extras = (
+        f"[{','.join(sorted(requirement.extras))}]"
+        if requirement.extras
+        else ""
+    )
+    if requirement.url is not None:
+        return f"{requirement.name}{extras} @ {requirement.url}"
+    return f"{requirement.name}{extras}{requirement.specifier}"
+
+
+def _run(command: list[str], *, env: dict[str, str] | None = None, dry_run: bool) -> None:
+    print("+", " ".join(command))
+    if dry_run:
+        return
+    subprocess.run(command, check=True, env=env)
+
+
+def _uv_executable() -> str:
+    uv = shutil.which("uv")
+    if uv is None:
+        raise RuntimeError("uv executable not found on PATH")
+    return uv
+
+
+def _parse_direct_url(url: str) -> _DirectUrlParts:
+    raw_url = url
+    vcs = None
+    if raw_url.startswith("git+"):
+        raw_url = raw_url.removeprefix("git+")
+        vcs = "git"
+    raw_url, _, fragment = raw_url.partition("#")
+    url_without_revision = raw_url
+    requested_revision = None
+    if "@" in raw_url:
+        url_without_revision, requested_revision = raw_url.rsplit("@", 1)
+    fragment_parts = parse_qs(fragment)
+    subdirectory = fragment_parts.get("subdirectory", [None])[0]
+    return _DirectUrlParts(
+        vcs=vcs,
+        url=url_without_revision,
+        requested_revision=requested_revision,
+        subdirectory=subdirectory,
+    )
+
+
+def _direct_url_matches(requirement: Requirement) -> bool:
+    try:
+        dist = metadata.distribution(requirement.name)
+    except metadata.PackageNotFoundError:
+        return False
+    direct_url_text = dist.read_text("direct_url.json")
+    if direct_url_text is None:
+        return False
+    installed = json.loads(direct_url_text)
+    expected = _parse_direct_url(requirement.url or "")
+    if installed.get("url") != expected.url:
+        return False
+    if installed.get("subdirectory") != expected.subdirectory:
+        return False
+    if expected.vcs is None:
+        return True
+    vcs_info = installed.get("vcs_info") or {}
+    if vcs_info.get("vcs") != expected.vcs:
+        return False
+    if expected.requested_revision is None:
+        return True
+    return vcs_info.get("requested_revision") == expected.requested_revision
+
+
+def _specifier_matches(requirement: Requirement) -> bool:
+    try:
+        installed_version = metadata.version(requirement.name)
+    except metadata.PackageNotFoundError:
+        return False
+    return requirement.specifier.contains(Version(installed_version), prereleases=True)
+
+
+def dependencies_in_sync() -> bool:
+    for requirement in _megatron_requirements().values():
+        if requirement.url is not None:
+            if not _direct_url_matches(requirement):
+                return False
+            continue
+        if not _specifier_matches(requirement):
+            return False
+    return True
+
+
+def install_megatron_dependencies(*, dry_run: bool = False) -> None:
+    requirements = _megatron_requirements()
+
+    apex_requirement = requirements.pop("apex", None)
+    if apex_requirement is not None:
+        _run(
+            [
+                _uv_executable(),
+                "pip",
+                "install",
+                "--no-build-isolation",
+                _format_requirement(apex_requirement),
+            ],
+            env={**os.environ, **_APEX_ENV},
+            dry_run=dry_run,
+        )
+
+    core_requirements = [
+        _format_requirement(requirements.pop(name))
+        for name in _CORE_INSTALL_ORDER
+        if name in requirements
+    ]
+    if core_requirements:
+        override_path = "/tmp/art-megatron-transformer-engine-override.txt"
+        if not dry_run:
+            with open(override_path, "w", encoding="utf-8") as override_file:
+                override_file.write("transformer-engine==2.11.0\n")
+        _run(
+            [
+                _uv_executable(),
+                "pip",
+                "install",
+                "--no-build-isolation",
+                "--override",
+                override_path,
+                *core_requirements,
+            ],
+            env={**os.environ, **_TRANSFORMER_ENGINE_ENV},
+            dry_run=dry_run,
+        )
+        if not dry_run and os.path.exists(override_path):
+            os.remove(override_path)
+
+    bridge_requirement = requirements.pop("megatron-bridge", None)
+    if bridge_requirement is not None:
+        _run(
+            [
+                _uv_executable(),
+                "pip",
+                "install",
+                "--no-deps",
+                _format_requirement(bridge_requirement),
+            ],
+            dry_run=dry_run,
+        )
+
+    _run([_uv_executable(), "pip", "uninstall", "-y", "pynvml"], dry_run=dry_run)
+
+    remaining_requirements = [
+        _format_requirement(requirement) for requirement in requirements.values()
+    ]
+    if remaining_requirements:
+        _run(
+            [_uv_executable(), "pip", "install", *remaining_requirements],
+            dry_run=dry_run,
+        )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Install ART Megatron dependencies")
+    parser.add_argument(
+        "--check",
+        action="store_true",
+        help="Exit 0 if installed Megatron dependencies already match ART's install plan.",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Print uv pip commands without executing them.",
+    )
+    args = parser.parse_args()
+
+    if args.check:
+        raise SystemExit(0 if dependencies_in_sync() else 1)
+
+    install_megatron_dependencies(dry_run=args.dry_run)
+
+
+if __name__ == "__main__":
+    main()