[MUSA] Harden CI dependencies and diffusion warmup (#35610)

scripts/ci/musa/musa_install_dependency.sh

@@ -1,4 +1,5 @@
 #!/bin/bash
+# shellcheck disable=SC2034  # OPTIONAL_DEPS is retained for CLI compatibility.
 set -euo pipefail
 
 # Parse command line arguments
@@ -21,8 +22,40 @@ while [[ $# -gt 0 ]]; do
   esac
 done
 
-PIP_INSTALL="python3 -m pip install --no-cache-dir"
-${PIP_INSTALL} --upgrade pip setuptools torchada --user
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
+STACK_HELPER="${SCRIPT_DIR}/musa_python_stack.py"
+# Keep installed packages isolated in PYTHONUSERBASE, but let pip reuse its
+# content-addressed wheel cache across jobs. Disabling the cache made every
+# MUSA lane cold-download hundreds of megabytes and exceed the install timeout.
+PIP_INSTALL=(python3 -m pip install)
+readonly MUSA_TRITON_VERSION="3.2.0"
+# Current MUSA CI uses the CPython 3.10 x86_64 wheel. A Python upgrade must
+# update this digest together with the pinned artifact.
+readonly MUSA_TRITON_SHA256="65b15d42fac24a2eca4c0c9f0ac68c8bd7cbe6bcc9f619c3483fb4f323391303"
+readonly MUSA_TRITON_INDEX_URL="https://dl.mthreads.com/repo/api/pypi/pypi/simple"
+readonly MUSA_TORCHADA_VERSION="0.1.82"
+readonly MUSA_TORCHADA_SHA256="472663da083ef23502f08429618a36e5e6e9b2447cf72fff2568787c815b5903"
+readonly MUSA_TORCHADA_INDEX_URL="https://pypi.org/simple"
+readonly MUSA_SETUPTOOLS_SPEC="setuptools<82"
+MUSA_CI_SCRATCH="$(mktemp -d "${RUNNER_TEMP:-${TMPDIR:-/tmp}}/sglang-musa-ci.XXXXXX")"
+MUSA_CI_ISOLATED_USERBASE=""
+
+if [ "${GITHUB_ACTIONS:-}" = "true" ]; then
+    MUSA_CI_ISOLATED_USERBASE="1"
+    PYTHONUSERBASE="$(mktemp -d "${RUNNER_TEMP:-${TMPDIR:-/tmp}}/sglang-musa-python.XXXXXX")"
+    export PYTHONUSERBASE
+    if [ -n "${GITHUB_ENV:-}" ]; then
+        echo "PYTHONUSERBASE=${PYTHONUSERBASE}" >> "$GITHUB_ENV"
+    fi
+    if [ -n "${GITHUB_PATH:-}" ]; then
+        echo "${PYTHONUSERBASE}/bin" >> "$GITHUB_PATH"
+    fi
+    echo "Using task-local Python user base: ${PYTHONUSERBASE}"
+fi
+
+# torchada has an unpinned Torch dependency. Installing it from the public
+# index before the MUSA wheel bundle can pull CUDA Torch and CUDA Triton.
+"${PIP_INSTALL[@]}" --upgrade pip "$MUSA_SETUPTOOLS_SPEC" --user
 
 echo "Checking stale torchada extension locks..."
 active_torchada_builds="$(
@@ -43,11 +76,135 @@ if [ -d "$torch_extensions_dir" ]; then
         -delete
 fi
 
-WHL_DIR="/sglang-checkout/whl"
+WHL_DIR="${WHL_DIR:-/sglang-checkout/whl}"
+MUSA_TRITON_WHEEL=""
+MUSA_TRITON_TASK_LOCAL=""
+MUSA_TORCHADA_WHEEL=""
+
+wheel_sha256() {
+    sha256sum "$1" | awk '{print $1}'
+}
+
+find_bundled_musa_triton() {
+    local candidate
+    local -a candidates
+    if [ ! -d "$WHL_DIR" ]; then
+        return
+    fi
+    mapfile -t candidates < <(
+        compgen -G "${WHL_DIR}/triton-${MUSA_TRITON_VERSION}-*.whl" || true
+    )
+    for candidate in "${candidates[@]}"; do
+        if [ "$(wheel_sha256 "$candidate")" = "$MUSA_TRITON_SHA256" ]; then
+            MUSA_TRITON_WHEEL="$candidate"
+            return
+        fi
+    done
+}
+
+download_musa_triton() {
+    local download_dir="${MUSA_CI_SCRATCH}/musa-triton"
+    local -a candidates
+    mkdir -p "$download_dir"
+    python3 -m pip --isolated download \
+        --dest "$download_dir" \
+        --index-url "$MUSA_TRITON_INDEX_URL" \
+        --no-deps \
+        --only-binary=:all: \
+        "triton==${MUSA_TRITON_VERSION}"
+    mapfile -t candidates < <(find "$download_dir" -maxdepth 1 -type f -name '*.whl')
+    if [ "${#candidates[@]}" -ne 1 ]; then
+        echo "::error::Expected one MUSA Triton wheel, found ${#candidates[@]}"
+        exit 1
+    fi
+    MUSA_TRITON_WHEEL="${candidates[0]}"
+}
+
+find_bundled_torchada() {
+    local candidate
+    local -a candidates
+    if [ ! -d "$WHL_DIR" ]; then
+        return
+    fi
+    mapfile -t candidates < <(
+        compgen -G "${WHL_DIR}/torchada-${MUSA_TORCHADA_VERSION}-*.whl" || true
+    )
+    for candidate in "${candidates[@]}"; do
+        if [ "$(wheel_sha256 "$candidate")" = "$MUSA_TORCHADA_SHA256" ]; then
+            MUSA_TORCHADA_WHEEL="$candidate"
+            return
+        fi
+    done
+}
+
+download_torchada() {
+    local download_dir="${MUSA_CI_SCRATCH}/torchada"
+    local -a candidates
+    mkdir -p "$download_dir"
+    python3 -m pip --isolated download \
+        --dest "$download_dir" \
+        --index-url "$MUSA_TORCHADA_INDEX_URL" \
+        --no-deps \
+        --only-binary=:all: \
+        "torchada==${MUSA_TORCHADA_VERSION}"
+    mapfile -t candidates < <(find "$download_dir" -maxdepth 1 -type f -name '*.whl')
+    if [ "${#candidates[@]}" -ne 1 ]; then
+        echo "::error::Expected one torchada wheel, found ${#candidates[@]}"
+        exit 1
+    fi
+    MUSA_TORCHADA_WHEEL="${candidates[0]}"
+}
+
+find_bundled_musa_triton
+if [ -z "$MUSA_TRITON_WHEEL" ]; then
+    if python3 "$STACK_HELPER" verify \
+        --expected-triton-version "$MUSA_TRITON_VERSION" \
+        --triton-only; then
+        echo "Reusing the installed MUSA Triton"
+    else
+        # Existing runner bundles do not contain Triton yet. Keep a hash-pinned
+        # vendor-index fallback until the wheel is shipped in /sglang-checkout/whl.
+        download_musa_triton
+    fi
+fi
+if [ -n "$MUSA_TRITON_WHEEL" ]; then
+    if [ "$(wheel_sha256 "$MUSA_TRITON_WHEEL")" != "$MUSA_TRITON_SHA256" ]; then
+        echo "::error::MUSA Triton wheel SHA256 does not match the pinned artifact"
+        exit 1
+    fi
+    MUSA_TRITON_TASK_LOCAL="1"
+    echo "Using MUSA Triton wheel: ${MUSA_TRITON_WHEEL}"
+fi
+
+# torchada declares an unpinned dependency on Torch. Install the exact wheel
+# with --no-deps so a fresh user site cannot resolve public CUDA Torch/Triton.
+find_bundled_torchada
+if [ -z "$MUSA_TORCHADA_WHEEL" ]; then
+    download_torchada
+fi
+if [ "$(wheel_sha256 "$MUSA_TORCHADA_WHEEL")" != "$MUSA_TORCHADA_SHA256" ]; then
+    echo "::error::torchada wheel SHA256 does not match the pinned artifact"
+    exit 1
+fi
+echo "Using torchada wheel: ${MUSA_TORCHADA_WHEEL}"
+
+VENDOR_WHEELS=("$MUSA_TORCHADA_WHEEL")
+if [ -n "$MUSA_TRITON_WHEEL" ]; then
+    VENDOR_WHEELS+=("$MUSA_TRITON_WHEEL")
+fi
 if [ -d "$WHL_DIR" ] && compgen -G "${WHL_DIR}"/*.whl > /dev/null; then
+    for whl in "${WHL_DIR}"/*.whl; do
+        case "$(basename "$whl")" in
+          triton-*.whl|torchada-*.whl) continue;;
+        esac
+        VENDOR_WHEELS+=("$whl")
+    done
+fi
+
+if [ -z "$MUSA_CI_ISOLATED_USERBASE" ]; then
     echo "Uninstall old packages based on wheel METADATA..."
     PKGS=$(
-      for whl in "${WHL_DIR}"/*.whl; do
+      for whl in "${VENDOR_WHEELS[@]}"; do
         meta_file=$(zipinfo -1 "$whl" | awk '/\.dist-info\/METADATA$/ {print; exit}')
         [ -n "$meta_file" ] || continue
         unzip -p "$whl" "$meta_file" 2>/dev/null | sed -n 's/^Name: //p' | head -n1
@@ -55,18 +212,36 @@ if [ -d "$WHL_DIR" ] && compgen -G "${WHL_DIR}"/*.whl > /dev/null; then
     )
     for pkg in $PKGS; do
       echo "Uninstalling $pkg"
-      pip uninstall -y "$pkg" || true
+      python3 -m pip uninstall -y "$pkg" || true
     done
-    echo "Installing wheel files without dependency resolution..."
-    ${PIP_INSTALL} "${WHL_DIR}"/*.whl --user
 fi
 
+if [ "${#VENDOR_WHEELS[@]}" -gt 0 ]; then
+    echo "Installing vendor wheels without dependency resolution..."
+    if [ -n "$MUSA_CI_ISOLATED_USERBASE" ]; then
+        "${PIP_INSTALL[@]}" --ignore-installed --no-deps --user "${VENDOR_WHEELS[@]}"
+    else
+        "${PIP_INSTALL[@]}" --no-deps --user "${VENDOR_WHEELS[@]}"
+    fi
+fi
+
+MUSA_CONSTRAINTS="${MUSA_CI_SCRATCH}/constraints.txt"
+python3 "$STACK_HELPER" constraints --output "$MUSA_CONSTRAINTS"
+if [ -n "${GITHUB_ENV:-}" ]; then
+    echo "MUSA_CONSTRAINTS=${MUSA_CONSTRAINTS}" >> "$GITHUB_ENV"
+fi
+python3 "$STACK_HELPER" verify \
+    --expected-triton-version "$MUSA_TRITON_VERSION" \
+    ${MUSA_TRITON_TASK_LOCAL:+--require-user-site}
+
 if [ -n "$SKIP_SGLANG_BUILD" ]; then
     echo "Didn't build checkout SGLang"
     exit 0
 else
-    pip uninstall sgl-kernel -y || true
-    pip uninstall sglang -y || true
+    if [ -z "$MUSA_CI_ISOLATED_USERBASE" ]; then
+        python3 -m pip uninstall sgl-kernel -y || true
+        python3 -m pip uninstall sglang -y || true
+    fi
     # Clear Python cache to ensure latest code is used (works for any env: venv, system, conda)
     REPO_ROOT="${GITHUB_WORKSPACE:-$(pwd)}"
     find "$REPO_ROOT" -name "*.pyc" -delete 2>/dev/null || true
@@ -83,9 +258,17 @@ else
     bash "${REPO_ROOT}/scripts/ci/utils/install_rustup.sh"
     export PATH="${CARGO_HOME:-$HOME/.cargo}/bin:${PATH}"
 
-    cd "${REPO_ROOT}" && ${PIP_INSTALL} -v -e "python[dev_musa]" --user
+    cd "${REPO_ROOT}" && "${PIP_INSTALL[@]}" \
+        --constraint "$MUSA_CONSTRAINTS" \
+        -v \
+        -e "python[dev_musa]" \
+        --user
 
     cd "${REPO_ROOT}/python/sglang/kernels/aot"
     rm -f pyproject.toml && mv pyproject_musa.toml pyproject.toml && MTGPU_TARGET=mp_31 python3 setup_musa.py install --user
-    echo "$HOME/.local/bin" >> "$GITHUB_PATH"
+    python3 "$STACK_HELPER" verify \
+        --expected-triton-version "$MUSA_TRITON_VERSION" \
+        --require-driver \
+        --require-resolved-dependencies \
+        ${MUSA_TRITON_TASK_LOCAL:+--require-user-site}
 fi

scripts/ci/musa/musa_python_stack.py

@@ -0,0 +1,247 @@
+"""Build constraints and verify the Python stack used by MUSA CI."""
+
+from __future__ import annotations
+
+import argparse
+import importlib.metadata
+import json
+import logging
+import re
+import site
+from pathlib import Path
+
+CORE_DISTRIBUTIONS = (
+    "torch",
+    "torch-musa",
+    "torchada",
+    "triton",
+)
+
+OPTIONAL_VENDOR_DISTRIBUTIONS = (
+    "apache-tvm-ffi",
+    "deep-gemm",
+    "flash-attn-3",
+    "mate",
+    "mthreads-ml-py",
+    "mt-sparse-attention",
+    "torchaudio",
+    "torchvision",
+)
+
+OPTIONAL_STACK_DISTRIBUTIONS = ("setuptools",)
+
+# compressed-tensors 0.16+ requires Torch 2.10+, while the older MUSA runner
+# stack uses Torch 2.9. Keep this mapping explicit instead of globally pinning
+# one version in pyproject_other.toml for every accelerator stack.
+COMPRESSED_TENSORS_BY_TORCH_MINOR = {
+    (2, 9): "0.15.0",
+    (2, 11): "0.17.0",
+}
+
+LOGGER = logging.getLogger(__name__)
+
+
+class StackError(RuntimeError):
+    """Raised when the installed MUSA stack violates the CI contract."""
+
+
+def distribution_version(name: str) -> str:
+    try:
+        return importlib.metadata.version(name)
+    except importlib.metadata.PackageNotFoundError as exc:
+        raise StackError(f"required distribution is not installed: {name}") from exc
+
+
+def torch_minor(version: str) -> tuple[int, int]:
+    match = re.match(r"^(\d+)\.(\d+)", version)
+    if match is None:
+        raise StackError(f"cannot parse Torch version: {version!r}")
+    return int(match.group(1)), int(match.group(2))
+
+
+def compressed_tensors_version(torch_version: str) -> str:
+    minor = torch_minor(torch_version)
+    try:
+        return COMPRESSED_TENSORS_BY_TORCH_MINOR[minor]
+    except KeyError as exc:
+        supported = ", ".join(
+            f"{major}.{minor}"
+            for major, minor in sorted(COMPRESSED_TENSORS_BY_TORCH_MINOR)
+        )
+        raise StackError(
+            f"unsupported MUSA Torch line {minor[0]}.{minor[1]}; "
+            f"supported lines: {supported}"
+        ) from exc
+
+
+def build_constraints() -> list[str]:
+    versions = {name: distribution_version(name) for name in CORE_DISTRIBUTIONS}
+    pins = [f"{name}=={version}" for name, version in versions.items()]
+
+    for name in OPTIONAL_VENDOR_DISTRIBUTIONS + OPTIONAL_STACK_DISTRIBUTIONS:
+        try:
+            version = importlib.metadata.version(name)
+        except importlib.metadata.PackageNotFoundError:
+            continue
+        pins.append(f"{name}=={version}")
+
+    pins.append(f"compressed-tensors=={compressed_tensors_version(versions['torch'])}")
+    return sorted(pins, key=str.casefold)
+
+
+def write_constraints(output: Path) -> None:
+    pins = build_constraints()
+    output.parent.mkdir(parents=True, exist_ok=True)
+    output.write_text("\n".join(pins) + "\n", encoding="utf-8")
+    LOGGER.info("Wrote MUSA constraints to %s", output)
+    for pin in pins:
+        LOGGER.info("  %s", pin)
+
+
+def validate_core_versions(versions: dict[str, str]) -> None:
+    torch_line = torch_minor(versions["torch"])
+    torch_musa_line = torch_minor(versions["torch-musa"])
+    if torch_line != torch_musa_line:
+        raise StackError(
+            "Torch and Torch-MUSA lines do not match: "
+            f"torch={versions['torch']}, torch-musa={versions['torch-musa']}"
+        )
+
+
+def triton_metadata() -> dict[str, object]:
+    import triton
+
+    distribution = importlib.metadata.distribution("triton")
+    summary = distribution.metadata.get("Summary") or ""
+    module_file = getattr(triton, "__file__", None)
+    if module_file is None:
+        raise StackError("cannot locate the imported Triton module")
+    module_path = Path(module_file).resolve()
+    backend_root = module_path.parent / "backends"
+    if not backend_root.is_dir():
+        raise StackError(f"Triton backend directory is missing: {backend_root}")
+    backends = sorted(path.name for path in backend_root.iterdir() if path.is_dir())
+    return {
+        "version": distribution.version,
+        "summary": summary,
+        "module": str(module_path),
+        "backends": backends,
+        "user_site": site.getusersitepackages(),
+    }
+
+
+def verify_stack(
+    *,
+    expected_triton_version: str,
+    require_driver: bool,
+    require_resolved_dependencies: bool,
+    require_user_site: bool,
+    triton_only: bool,
+) -> None:
+    if require_driver:
+        import torchada  # noqa: F401
+
+    info = triton_metadata()
+    if info["version"] != expected_triton_version:
+        raise StackError(
+            "unexpected Triton version: "
+            f"observed={info['version']}, expected={expected_triton_version}"
+        )
+    if "MUSA" not in str(info["summary"]):
+        raise StackError(f"Triton is not the MUSA build: {info['summary']!r}")
+    if "mtgpu" not in info["backends"]:
+        raise StackError(f"Triton has no mtgpu backend: {info['backends']}")
+    if require_user_site:
+        module_path = Path(str(info["module"]))
+        user_site = Path(str(info["user_site"])).resolve()
+        if not module_path.is_relative_to(user_site):
+            raise StackError(
+                "Triton was not imported from the task-local user site: "
+                f"module={module_path}, user_site={user_site}"
+            )
+    if triton_only:
+        LOGGER.info(json.dumps({"triton": info}, indent=2, sort_keys=True))
+        return
+
+    versions = {name: distribution_version(name) for name in CORE_DISTRIBUTIONS}
+    validate_core_versions(versions)
+    result: dict[str, object] = {
+        "versions": versions,
+        "triton": info,
+        "driver_checked": require_driver,
+    }
+
+    if require_resolved_dependencies:
+        expected_compressed_tensors = compressed_tensors_version(versions["torch"])
+        observed_compressed_tensors = distribution_version("compressed-tensors")
+        if observed_compressed_tensors != expected_compressed_tensors:
+            raise StackError(
+                "compressed-tensors does not match the MUSA Torch line: "
+                f"observed={observed_compressed_tensors}, "
+                f"expected={expected_compressed_tensors}"
+            )
+        result["compressed_tensors"] = observed_compressed_tensors
+
+    if require_driver:
+        import torch
+        from triton.runtime import driver
+
+        musa_version = getattr(torch.version, "musa", None)
+        if musa_version is None:
+            raise StackError(f"Torch is not a MUSA build: {torch.__version__}")
+        if not hasattr(torch, "musa"):
+            raise StackError("torch.musa is unavailable after importing torchada")
+        device_count = torch.musa.device_count()
+        if device_count < 1:
+            raise StackError(f"no MUSA device is visible: device_count={device_count}")
+
+        target = driver.active.get_current_target()
+        if getattr(target, "backend", None) != "musa":
+            raise StackError(f"Triton active target is not MUSA: {target}")
+        result.update(
+            {
+                "musa_version": musa_version,
+                "device_count": device_count,
+                "target": repr(target),
+            }
+        )
+
+    LOGGER.info(json.dumps(result, indent=2, sort_keys=True, default=str))
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    constraints = subparsers.add_parser("constraints")
+    constraints.add_argument("--output", type=Path, required=True)
+
+    verify = subparsers.add_parser("verify")
+    verify.add_argument("--expected-triton-version", required=True)
+    verify.add_argument("--require-driver", action="store_true")
+    verify.add_argument("--require-resolved-dependencies", action="store_true")
+    verify.add_argument("--require-user-site", action="store_true")
+    verify.add_argument("--triton-only", action="store_true")
+    return parser.parse_args()
+
+
+def main() -> None:
+    logging.basicConfig(level=logging.INFO, format="%(message)s")
+    args = parse_args()
+    try:
+        if args.command == "constraints":
+            write_constraints(args.output)
+        else:
+            verify_stack(
+                expected_triton_version=args.expected_triton_version,
+                require_driver=args.require_driver,
+                require_resolved_dependencies=args.require_resolved_dependencies,
+                require_user_site=args.require_user_site,
+                triton_only=args.triton_only,
+            )
+    except StackError as exc:
+        raise SystemExit(f"MUSA Python stack error: {exc}") from exc
+
+
+if __name__ == "__main__":
+    main()

scripts/ci/musa/test_musa_python_stack.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import importlib.util
+import unittest
+from pathlib import Path
+from unittest import mock
+
+MODULE_PATH = Path(__file__).with_name("musa_python_stack.py")
+SPEC = importlib.util.spec_from_file_location("musa_python_stack", MODULE_PATH)
+assert SPEC is not None and SPEC.loader is not None
+musa_python_stack = importlib.util.module_from_spec(SPEC)
+SPEC.loader.exec_module(musa_python_stack)
+
+
+class MusaPythonStackTest(unittest.TestCase):
+    def test_compressed_tensors_for_torch_29(self) -> None:
+        self.assertEqual(
+            musa_python_stack.compressed_tensors_version("2.9.0+musa"),
+            "0.15.0",
+        )
+
+    def test_compressed_tensors_for_torch_211(self) -> None:
+        self.assertEqual(
+            musa_python_stack.compressed_tensors_version("2.11.0.post1+musa5.2.0"),
+            "0.17.0",
+        )
+
+    def test_unknown_torch_line_fails_closed(self) -> None:
+        with self.assertRaisesRegex(
+            musa_python_stack.StackError, "unsupported MUSA Torch line 2.10"
+        ):
+            musa_python_stack.compressed_tensors_version("2.10.0")
+
+    def test_mismatched_torch_and_torch_musa_lines_fail(self) -> None:
+        with self.assertRaisesRegex(
+            musa_python_stack.StackError,
+            "Torch and Torch-MUSA lines do not match",
+        ):
+            musa_python_stack.validate_core_versions(
+                {
+                    "torch": "2.11.0.post1+musa5.2.0",
+                    "torch-musa": "2.9.0+musa4.3.0",
+                }
+            )
+
+    def test_constraints_pin_core_and_optional_vendor_packages(self) -> None:
+        versions = {
+            "torch": "2.9.0",
+            "torch-musa": "2.9.0",
+            "torchada": "0.1.82",
+            "triton": "3.2.0",
+            "apache-tvm-ffi": "0.1.9.post3+musa.1",
+            "deep-gemm": "0.2.4+musa",
+            "mate": "0.2.0+musa",
+        }
+
+        def version(name: str) -> str:
+            try:
+                return versions[name]
+            except KeyError as exc:
+                raise musa_python_stack.importlib.metadata.PackageNotFoundError(
+                    name
+                ) from exc
+
+        with mock.patch.object(
+            musa_python_stack.importlib.metadata, "version", side_effect=version
+        ):
+            pins = musa_python_stack.build_constraints()
+
+        self.assertIn("torch==2.9.0", pins)
+        self.assertIn("torch-musa==2.9.0", pins)
+        self.assertIn("torchada==0.1.82", pins)
+        self.assertIn("triton==3.2.0", pins)
+        self.assertIn("apache-tvm-ffi==0.1.9.post3+musa.1", pins)
+        self.assertIn("deep-gemm==0.2.4+musa", pins)
+        self.assertIn("mate==0.2.0+musa", pins)
+        self.assertIn("compressed-tensors==0.15.0", pins)
+
+
+if __name__ == "__main__":
+    unittest.main()