minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
R0CKSTAR
2026-08-20 09:40:24 -07:00
committed by GitHub
parent be373395b4
commit 81df6f2c57
6 changed files with 542 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/nightly-test-musa.yml
+5 -6
View File
@@ -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
.github/workflows/pr-test-musa.yml
+10 -6
View File
@@ -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
python/sglang/multimodal_gen/test/server/musa/testcase_configs_musa.py
+5
View File
@@ -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,
scripts/ci/musa/musa_install_dependency.sh
+194 -11
View File
@@ -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
+247
View File
@@ -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
+81
View File
@@ -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()