diff --git a/.github/workflows/nightly-test-musa.yml b/.github/workflows/nightly-test-musa.yml index 0d9087e22..23a257ae0 100644 --- a/.github/workflows/nightly-test-musa.yml +++ b/.github/workflows/nightly-test-musa.yml @@ -66,7 +66,7 @@ jobs: - name: Install dependencies if: steps.gate.outputs.run_job == 'true' - timeout-minutes: 15 + timeout-minutes: 30 run: | bash scripts/ci/musa/musa_install_dependency.sh @@ -74,7 +74,6 @@ jobs: if: steps.gate.outputs.run_job == 'true' timeout-minutes: 30 run: | - pytest python/sglang/kernels/aot/tests/test_per_token_quant_fp8.py pytest python/sglang/kernels/aot/tests/speculative/test_eagle_utils.py pytest python/sglang/kernels/aot/tests/speculative/test_ngram_utils.py pytest python/sglang/kernels/aot/tests/speculative/test_speculative_sampling.py @@ -111,7 +110,7 @@ jobs: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies - timeout-minutes: 15 + timeout-minutes: 30 if: steps.gate.outputs.run_job == 'true' run: | bash scripts/ci/musa/musa_install_dependency.sh @@ -154,7 +153,7 @@ jobs: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies - timeout-minutes: 15 + timeout-minutes: 30 if: steps.gate.outputs.run_job == 'true' run: | bash scripts/ci/musa/musa_install_dependency.sh @@ -206,7 +205,7 @@ jobs: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies - timeout-minutes: 15 + timeout-minutes: 30 if: steps.gate.outputs.run_job == 'true' run: | bash scripts/ci/musa/musa_install_dependency.sh @@ -255,7 +254,7 @@ jobs: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies - timeout-minutes: 15 + timeout-minutes: 30 if: steps.gate.outputs.run_job == 'true' run: | bash scripts/ci/musa/musa_install_dependency.sh diff --git a/.github/workflows/pr-test-musa.yml b/.github/workflows/pr-test-musa.yml index 550136fdc..c85273d3a 100644 --- a/.github/workflows/pr-test-musa.yml +++ b/.github/workflows/pr-test-musa.yml @@ -72,7 +72,7 @@ jobs: - "python/pyproject_other.toml" - "python/sglang/!(multimodal_gen)/**" - "python/sglang/srt/hardware_backend/musa/**" - - "scripts/ci/musa/*" + - "scripts/ci/musa/**" - "scripts/ci/utils/*" multimodal_gen: - "python/pyproject_other.toml" @@ -81,9 +81,14 @@ jobs: - "python/sglang/multimodal_gen/test/unit/musa/layers/test_musa_silu_and_mul.py" - "python/sglang/multimodal_gen/test/server/musa/run_suite.py" - "python/sglang/multimodal_gen/test/server/musa/**" + - "scripts/ci/musa/**" sgl_kernel: - ".github/workflows/pr-test-musa.yml" - "python/sglang/kernels/aot/csrc/musa/**" + - "scripts/ci/musa/**" + + - name: Test MUSA dependency constraints + run: python3 -m unittest scripts/ci/musa/test_musa_python_stack.py # ==================== PR Gate ==================== # pr-gate: @@ -117,7 +122,7 @@ jobs: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies - timeout-minutes: 15 + timeout-minutes: 30 run: | bash scripts/ci/musa/musa_install_dependency.sh @@ -152,7 +157,7 @@ jobs: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies - timeout-minutes: 15 + timeout-minutes: 30 run: | bash scripts/ci/musa/musa_install_dependency.sh @@ -183,7 +188,7 @@ jobs: uses: actions/checkout@v4 - name: Install dependencies - timeout-minutes: 15 + timeout-minutes: 30 run: | bash scripts/ci/musa/musa_install_dependency.sh @@ -210,14 +215,13 @@ jobs: ref: ${{ inputs.ref || github.ref }} - name: Install dependencies - timeout-minutes: 15 + timeout-minutes: 30 run: | bash scripts/ci/musa/musa_install_dependency.sh - name: Run sgl-kernel test timeout-minutes: 20 run: | - pytest python/sglang/kernels/aot/tests/test_per_token_quant_fp8.py pytest python/sglang/kernels/aot/tests/speculative/test_eagle_utils.py pytest python/sglang/kernels/aot/tests/speculative/test_ngram_utils.py pytest python/sglang/kernels/aot/tests/speculative/test_speculative_sampling.py diff --git a/python/sglang/multimodal_gen/test/server/musa/testcase_configs_musa.py b/python/sglang/multimodal_gen/test/server/musa/testcase_configs_musa.py index c3019dedb..7e2398188 100644 --- a/python/sglang/multimodal_gen/test/server/musa/testcase_configs_musa.py +++ b/python/sglang/multimodal_gen/test/server/musa/testcase_configs_musa.py @@ -46,6 +46,11 @@ ONE_GPU_MUSA_CASES: list[DiffusionTestCase] = [ model_path=hf_cached_model("Wan-AI/Wan2.1-T2V-1.3B-Diffusers"), modality="video", custom_validator="video", + # Server warmup caps videos at 17 frames, while this test's + # one-second request resolves to 24. Reuse the first real request + # shape for an unmeasured warmup so first-shape initialization + # stays outside the denoising performance metrics. + extras=["--warmup-mode", "request"], ), DiffusionSamplingParams( prompt=T2V_PROMPT, diff --git a/scripts/ci/musa/musa_install_dependency.sh b/scripts/ci/musa/musa_install_dependency.sh index 1056cc9cc..edac11a1d 100755 --- a/scripts/ci/musa/musa_install_dependency.sh +++ b/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 diff --git a/scripts/ci/musa/musa_python_stack.py b/scripts/ci/musa/musa_python_stack.py new file mode 100644 index 000000000..ad06825f2 --- /dev/null +++ b/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() diff --git a/scripts/ci/musa/test_musa_python_stack.py b/scripts/ci/musa/test_musa_python_stack.py new file mode 100644 index 000000000..53302171e --- /dev/null +++ b/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()