diff --git a/.github/workflows/build-release.yml b/.github/workflows/build-release.yml index 08ea1d4..52f8c3d 100644 --- a/.github/workflows/build-release.yml +++ b/.github/workflows/build-release.yml @@ -1,6 +1,7 @@ -name: Build release +name: Build and publish release on: + workflow_dispatch: push: tags: ["v*"] @@ -16,6 +17,7 @@ jobs: with: python-version: "3.11" - name: Verify tag matches package version + if: github.event_name == 'push' env: RELEASE_TAG: ${{ github.ref_name }} run: | @@ -37,9 +39,11 @@ jobs: source "$RUNNER_TEMP/aspec-build-venv/bin/activate" echo "$VIRTUAL_ENV/bin" >> "$GITHUB_PATH" echo "VIRTUAL_ENV=$VIRTUAL_ENV" >> "$GITHUB_ENV" - python -m pip install --upgrade pip build + python -m pip install --upgrade pip build twine - name: Build source and wheel distributions run: python -m tools.build_release --outdir dist + - name: Check PyPI distribution metadata + run: python -m twine check --strict dist/*.whl dist/*.tar.gz - name: Save audited distributions and runtime requirements uses: actions/upload-artifact@v7 with: @@ -73,10 +77,46 @@ jobs: python -m pip install -r requirements/ci-cpu.txt # Deliberately test the wheel and patch CLI with a CPU dependency subset. # This is not installation/acceptance of the production GPU runtime. - python -m pip install --no-deps dist/*.whl sglang==0.5.21 + python -m pip install --no-deps dist/ling_aspec_runtime-*.whl sglang==0.5.21 - name: Exercise installed commands outside the checkout env: CUDA_VISIBLE_DEVICES: "" HF_HUB_OFFLINE: "1" TRANSFORMERS_OFFLINE: "1" run: python tools/check_wheel_install.py + + publish: + needs: [build, wheel-install] + if: >- + github.repository == 'inclusionAI/ASpec' && + github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v') + runs-on: ubuntu-latest + environment: + name: pypi + url: https://pypi.org/project/ling-aspec/ + steps: + - uses: actions/download-artifact@v8 + with: + name: aspec-distributions + path: dist/ + - name: Verify distribution checksums + working-directory: dist + run: sha256sum --check SHA256SUMS + - name: Select audited Python distributions + run: | + mkdir pypi-dist + cp dist/*.whl dist/*.tar.gz pypi-dist/ + - name: Require PyPI API token + env: + PYPI_API_TOKEN: ${{ secrets.PYPI_API_TOKEN }} + run: | + if [ -z "$PYPI_API_TOKEN" ]; then + echo "::error::Set PYPI_API_TOKEN in the repository's pypi environment secrets." + exit 1 + fi + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@7245906fb61fa33a13554f9d5fb856fe61758086 # v1.14.3 + with: + password: ${{ secrets.PYPI_API_TOKEN }} + packages-dir: pypi-dist/ + attestations: false diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e1a5480..70d31c6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -62,6 +62,8 @@ jobs: tests/test_environment_scripts.py \ tests/test_onboarding.py \ tests/test_dependency_extras.py \ + tests/test_target_environment.py \ + tests/test_target_service.py \ tests/test_cudnn_gqa_mode.py \ tests/test_dependency_audit.py \ tests/test_draft_model_config.py \ diff --git a/CHANGELOG.md b/CHANGELOG.md index 9c2cd9e..5c0fd71 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,19 @@ ## Unreleased +- Install independent Draft (`ling-aspec`) and Target (`ling-aspec-target`) environments, + sharing the version-aligned `ling-aspec-runtime` code package. Target retains SGLang's own + dependency requirements; no resolver overrides are used. Start `aspec target` and `aspec train` + in their respective environments and connect by service address. Optional managed-local + supervision uses the same service interface and retains full checkpoint/resume. +- Keep `aspec patch-sglang` limited to patching files. Report role-specific worker provenance, + and expose pip CUDA toolkit libraries through a private JIT development view. +- Set the first PyPI version to `0.1.0b0`; developers maintain release versions with the `b0` + suffix while CI checks the matching tag without changing the version. +- Use `ling-aspec` as the distribution name, retaining the `aspec` namespace and CLI. Publish + audited tag artifacts to PyPI after wheel installation checks through the GitHub `pypi` environment. +- Allow repeatable manual build/install validation without publication; reserve PyPI uploads + for version tag pushes. - Consolidate supplementary examples around distinct workflows; use recipe/config overrides for alternative lengths, backends and structure combinations. - Advance fully committed restored epochs for repeating training sources, including new diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b7a2c14..9f4876b 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -21,19 +21,25 @@ Use a dedicated Linux venv and the [pinned runtime](docs/user-guide/installation ```bash bash setup.sh --dev +bash setup.sh --target --dev source .venv/bin/activate python -m ruff check aspec tests tools setup.py python -m tools.check_docs -CUDA_VISIBLE_DEVICES="" HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 bash tools/run_in_env.sh python -m pytest -q +CUDA_VISIBLE_DEVICES="" HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 bash tools/run_in_env.sh python -m pytest -q --ignore=tests/test_sgl_engine_factory.py --ignore=tests/test_inference_concurrency.py python -m tools.build_release --outdir /absolute/path/to/new-release-dist +source .venv-target/bin/activate +CUDA_VISIBLE_DEVICES="" HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 aspec doctor --target +python -m pytest -q tests/test_sgl_engine_factory.py tests/test_inference_concurrency.py tests/test_sgl_engine_decode.py ``` Default tests are fast/offline. GPU, tokenizer, training and cuDNN gates need explicit switches declared in their test files. A skip is not GPU acceptance. Public CI checks lint, local doc links, compilation, plus a curated CPU/offline contract suite installed from `requirements/ci-cpu.txt` on pull requests and configured branch pushes. The separate [release build](.github/workflows/build-release.yml) -workflow builds and audits sdist/wheel artifacts and checks wheel installation only when a `v*` tag -is pushed; see [release instructions](docs/releasing.md). Neither workflow certifies GPU health or +workflow builds and audits sdist/wheel artifacts and checks wheel installation on manual runs and +`v*` tag pushes. Manual runs only validate the artifacts and their installation. +In `inclusionAI/ASpec`, tag pushes publish `ling-aspec` to PyPI through the `pypi` environment; +see [release instructions](docs/releasing.md). Neither workflow certifies GPU health or replaces the complete pinned-runtime suite above. Equivalent remote checks are not proof that hosted GitHub Actions has already run. diff --git a/Dockerfile b/Dockerfile index e676bf0..7b9232c 100644 --- a/Dockerfile +++ b/Dockerfile @@ -6,6 +6,7 @@ RUN apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-ins WORKDIR /opt/ASpec COPY . . RUN bash setup.sh --venv /opt/aspec-venv --python /usr/bin/python3.12 +RUN bash setup.sh --target --venv /opt/aspec-target --python /usr/bin/python3.12 ENV VIRTUAL_ENV=/opt/aspec-venv ENV PATH="/opt/aspec-venv/bin:${PATH}" CMD ["bash"] diff --git a/README.md b/README.md index a58f221..8bccd20 100644 --- a/README.md +++ b/README.md @@ -52,22 +52,36 @@ paths have explicit support boundaries, and DFlash2 remains experimental. Scale from a first multi-GPU run to static multi-node training, with NVLink, TCP and RDMA data paths. Check [supported paths and limitations](docs/support.md), then browse the [recipes](docs/user-guide/recipes.md) for starting points. -## Start here +## Installation -1. [Run source setup](docs/user-guide/installation.md); [containers](docs/user-guide/containers.md) are also available. -2. [Choose Qwen3-8B or Ling tiny](docs/user-guide/quickstart.md) for the same two-GPU Plan, train, - report, resume and export workflow. -3. [Choose a training recipe](docs/user-guide/recipes.md) and [prepare training data](docs/user-guide/data.md). +In an activated Linux virtual environment: -Recipes cover Qwen3-8B, Qwen3.5-9B Native MTP, Qwen3.8-27B long context, and Ling tiny/flash 3.0/flash 3.1. -For complete feature workflows, see [offline feature generation and replay](examples/qwen3-8b-offline/README.md) -and [DSA dense → indexer → sparse training](examples/qwen3-8b-dsa/README.md). +```bash +pip install ling-aspec +``` + +This installs the Draft training environment. In a **separate Target virtual environment**, install +the SGLang integration and apply its patch: + +```bash +pip install ling-aspec-target +aspec patch-sglang +aspec target +``` + +Keep the Target service running, then start training in the Draft environment below. +The default connection is `http://127.0.0.1:30000`; set `inference.target_service.address` +for a remote Target. Each command uses its own environment. + +See the [installation guide](docs/user-guide/installation.md) for prerequisites, source installation +and release wheels, or use the [container setup](docs/user-guide/containers.md). +Optionally run `aspec doctor` to check the environment. -From the repository root: +## Quickstart + +With two NVLink-connected GPUs: ```bash -bash setup.sh -source .venv/bin/activate CUDA_VISIBLE_DEVICES=0,1 aspec quickstart ling-tiny --output outputs/ling-tiny-001 # Or choose Qwen3-8B: CUDA_VISIBLE_DEVICES=0,1 aspec quickstart qwen3-8b --output outputs/qwen3-8b-001 @@ -75,9 +89,14 @@ CUDA_VISIBLE_DEVICES=0,1 aspec quickstart qwen3-8b --output outputs/qwen3-8b-001 Both profiles use bundled raw conversations, their own tokenizer/template and an eight-step train/resume/native-export check. Add `--target /absolute/path/to/model` to use a local checkpoint. +See [Quickstart](docs/user-guide/quickstart.md) for the full workflow, then choose a +[training recipe](docs/user-guide/recipes.md) and [prepare your data](docs/user-guide/data.md). For training with Target-generated data and held-out evaluation, continue with the [Qwen3-8B DSpark recipe](docs/user-guide/recipes.md#qwen3-8b-dspark-quality-recipe). +Additional workflows cover [offline feature generation and replay](examples/qwen3-8b-offline/README.md) +and [DSA dense → indexer → sparse training](examples/qwen3-8b-dsa/README.md). + ## Performance ASpec improves both Draft execution and the complete online training pipeline. The diff --git a/README_zh-CN.md b/README_zh-CN.md index a46b68a..d282f49 100644 --- a/README_zh-CN.md +++ b/README_zh-CN.md @@ -47,23 +47,42 @@ Qwen、Ling 上的 DSpark/DFlash 是主要优化方向和 recipe 维护重点； 从首次多卡运行扩展到静态多机训练，支持 NVLink、TCP 和 RDMA 数据路径。 先查看[支持范围与已知限制](docs/support.md)，再从 [recipe 导航](docs/user-guide/recipes.md)选择起点。 -## 开始使用 +## 安装 -1. 按[安装指南](docs/user-guide/installation.md)一键创建源码环境，也可选择[容器](docs/user-guide/containers.md)； -2. 在[快速上手](docs/user-guide/quickstart.md)中选择 Qwen3-8B 或 Ling tiny，用两张 GPU 完成 - 相同的 Plan、训练、报告、恢复和导出流程； -3. 根据[recipe 导航](docs/user-guide/recipes.md)选择算法/结构并准备实际数据。 +在已激活的 Linux 虚拟环境中安装： + +```bash +pip install ling-aspec +``` + +该命令安装 Draft 训练依赖。在独立的 Target 虚拟环境中安装 SGLang 集成并应用补丁： + +```bash +pip install ling-aspec-target +aspec patch-sglang +aspec target +``` + +随后在 Draft 环境启动训练；Target 服务保持运行。默认连接 `http://127.0.0.1:30000`，远程 Target 通过 +`inference.target_service.address` 配置地址。两个命令各自使用当前环境。 + +环境要求、源码安装和本地 wheel 安装见[安装指南](docs/user-guide/installation.md)， +也可使用[容器](docs/user-guide/containers.md)。可选执行 `aspec doctor` 检查环境。 + +## 快速上手 + +使用两张 NVLink 互联的 GPU： ```bash -bash setup.sh -source .venv/bin/activate CUDA_VISIBLE_DEVICES=0,1 aspec quickstart ling-tiny --output outputs/ling-tiny-001 # 或选择 Qwen3-8B： CUDA_VISIBLE_DEVICES=0,1 aspec quickstart qwen3-8b --output outputs/qwen3-8b-001 ``` 两个入口都使用内置原始对话和各自的 tokenizer/template，完成 8 步训练、精确恢复和 native 导出加载。 -使用已有模型时添加 `--target /absolute/path/to/model`。使用 Target 生成的数据开展训练与 +使用已有模型时添加 `--target /absolute/path/to/model`。完整流程见[快速上手](docs/user-guide/quickstart.md)， +实际训练可从 [recipe 导航](docs/user-guide/recipes.md)选择算法与结构，并[准备数据](docs/user-guide/data.md)。 +使用 Target 生成的数据开展训练与 独立评估，可继续阅读 [Qwen3-8B DSpark recipe](docs/user-guide/recipes.md#qwen3-8b-dspark-quality-recipe)。 ## 性能 diff --git a/aspec/_environment.py b/aspec/_environment.py index 00d3206..c580d42 100644 --- a/aspec/_environment.py +++ b/aspec/_environment.py @@ -13,12 +13,46 @@ # limitations under the License. """Prepare this venv's native libraries before executing a fresh process.""" +import hashlib import os import sys import sysconfig from pathlib import Path +def _toolkit_view(toolkit, env): + """Expose pip's CUDA layout to JITs expecting a development toolkit. + + NVIDIA wheels ship lib/libcudart.so.13, while SGLang links -lcudart from + CUDA_HOME/lib64. Keep the development aliases in ASpec's private cache; + the installed wheels and host toolkit remain untouched. + """ + cache = Path(env.get("ASPEC_CACHE_HOME") or + Path(env.get("XDG_CACHE_HOME") or Path.home() / ".cache") / "aspec") + identity = hashlib.sha256(str(toolkit).encode()).hexdigest()[:16] + view = cache.expanduser().resolve() / "cuda-toolkit" / identity + libraries = view / "lib64" + libraries.mkdir(parents=True, exist_ok=True) + + def link(destination, source): + try: + destination.symlink_to(source, target_is_directory=source.is_dir()) + except FileExistsError: + if not destination.is_symlink() or destination.readlink() != source: + raise ValueError(f"Unexpected file in private CUDA toolkit view: {destination}") + + for source in toolkit.iterdir(): + if source.name != "lib64": + link(view / source.name, source) + for source in sorted((toolkit / "lib").iterdir(), key=lambda p: (len(p.name), p.name)): + link(libraries / source.name, source) + if ".so." in source.name: + alias = libraries / (source.name.split(".so.", 1)[0] + ".so") + if not alias.is_symlink(): + link(alias, source) + return view + + def command_environment(): if sys.platform != "linux" or sys.version_info[:2] < (3, 11): raise ValueError("Use Linux with Python >=3.11.") @@ -40,8 +74,19 @@ def command_environment(): raise ValueError("ASPEC_CUDA_COMPAT_DIR must contain libcuda.so.1.") paths.append(str(path)) purelib = Path(sysconfig.get_paths()["purelib"]) + toolkit = purelib / "nvidia/cu13" + if not env.get("CUDA_HOME") and not env.get("CUDA_PATH") and (toolkit / "bin/nvcc").is_file(): + view = _toolkit_view(toolkit, env) + env["CUDA_HOME"] = str(view) + env["ASPEC_AUTO_CUDA_HOME"] = str(view) + env["PATH"] = str(view / "bin") + os.pathsep + env.get("PATH", "") + previous = env.get("LIBRARY_PATH", "") + env["LIBRARY_PATH"] = str(view / "lib64") + (os.pathsep + previous if previous else "") + paths.append(str(view / "lib64")) paths.extend(str(p) for p in (purelib / "nvidia/cuda_runtime/lib", purelib / "torch/lib") if p.is_dir()) + if (toolkit / "lib").is_dir(): + paths.append(str(toolkit / "lib")) if paths: previous = env.get("LD_LIBRARY_PATH", "") env["LD_LIBRARY_PATH"] = ":".join(paths + ([previous] if previous else [])) diff --git a/aspec/_target_worker.py b/aspec/_target_worker.py new file mode 100644 index 0000000..5655ad4 --- /dev/null +++ b/aspec/_target_worker.py @@ -0,0 +1,29 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +"""Bootstrap the optional managed-local Target with its own native libraries.""" +import os +import sys +from pathlib import Path + +from aspec._environment import execute + + +def main(): + os.environ["VIRTUAL_ENV"] = sys.prefix + automatic_toolkit = os.environ.pop("ASPEC_AUTO_CUDA_HOME", None) + if automatic_toolkit and os.environ.get("CUDA_HOME") == automatic_toolkit: + os.environ.pop("CUDA_HOME", None) + os.environ["PATH"] = str(Path(sys.executable).parent) + os.pathsep + os.environ.get("PATH", "") + if cache := os.environ.get("ASPEC_CACHE_HOME"): + if Path(cache).name != "target": + os.environ["ASPEC_CACHE_HOME"] = str(Path(cache) / "target") + # The supervisor belongs to Draft. Never load its CUDA libraries in Target. + os.environ.pop("LD_LIBRARY_PATH", None) + os.environ.pop("LIBRARY_PATH", None) + os.environ.pop("PYTHONPATH", None) + os.environ.pop("PYTHONHOME", None) + execute([sys.executable, *sys.argv[1:]]) + + +if __name__ == "__main__": + main() diff --git a/aspec/cli.py b/aspec/cli.py index dc7515c..d6c2c59 100644 --- a/aspec/cli.py +++ b/aspec/cli.py @@ -13,7 +13,9 @@ # limitations under the License. """Workflow shortcuts and native-library bootstrap for the public CLI.""" +import importlib.metadata import sys +from pathlib import Path from aspec._environment import execute @@ -21,6 +23,8 @@ def main(): arguments = sys.argv[1:] modules = { + "train": "aspec.train_entry", + "target": "aspec.inference.service", "quickstart": "aspec.commands.quickstart", "recipes": "aspec.commands.recipes", "doctor": "aspec.commands.doctor", @@ -45,6 +49,19 @@ def main(): module = modules[arguments.pop(0)] else: module = "aspec.train_entry" + marker = Path(sys.prefix) / ".aspec-role" + target_role = marker.is_file() and marker.read_text().strip() == "target" + try: + importlib.metadata.version("ling-aspec-target") + target_role = True + except importlib.metadata.PackageNotFoundError: + pass + help_requested = arguments in (["--help"], ["-h"]) + if not help_requested and module in {"aspec.train_entry", "aspec.commands.quickstart"} and target_role: + raise SystemExit("Run training from the Draft venv (pip install ling-aspec); " + "this interpreter belongs to Target") + if module == "aspec.inference.service" and not target_role and not help_requested: + raise SystemExit("Start aspec target in the Target environment (pip install ling-aspec-target)") try: execute([sys.executable, "-m", module, *arguments]) except (ValueError, OSError) as error: diff --git a/aspec/commands/convert_to_hf.py b/aspec/commands/convert_to_hf.py index a469064..64d3c81 100644 --- a/aspec/commands/convert_to_hf.py +++ b/aspec/commands/convert_to_hf.py @@ -100,7 +100,7 @@ def _sha256_file(path: Path) -> str: def _aspec_version() -> str: try: - return importlib.metadata.version("aspec") + return importlib.metadata.version("ling-aspec-runtime") except importlib.metadata.PackageNotFoundError: return "unknown" diff --git a/aspec/commands/doctor.py b/aspec/commands/doctor.py index 507c283..410eeaa 100644 --- a/aspec/commands/doctor.py +++ b/aspec/commands/doctor.py @@ -19,6 +19,7 @@ import json import platform import sys +from pathlib import Path from packaging.requirements import InvalidRequirement, Requirement from packaging.utils import canonicalize_name @@ -26,16 +27,9 @@ from aspec.assets import asset_path from aspec.compat.sglang_patch import installed_site, patch_state, venv_error -# These are deliberate overrides, not satisfied upstream metadata. Keep the -# upstream requirement and installed version exact so new conflicts fail closed. -APPROVED_OVERRIDES = { - ("sglang", "0.5.21", "nvidia-cutlass-dsl[cu13]==4.6.2", "4.7.1"), - ("sglang", "0.5.21", "quack-kernels==0.6.4", "0.6.5"), -} - def check_dependency_metadata(distributions=None): - """Audit installed active requirements and disclose exact approved exceptions.""" + """Audit installed active requirements; no metadata conflicts are exempted.""" if distributions is None: distributions = importlib.metadata.distributions() installed = {canonicalize_name(dist.metadata["Name"]): dist for dist in distributions} @@ -53,27 +47,24 @@ def check_dependency_metadata(distributions=None): actual = dependency.version if dependency is not None else None if actual is not None and requirement.specifier.contains(actual, prereleases=True): continue - identity = (name, dist.version, str(requirement), actual) - if identity in APPROVED_OVERRIDES: - result["overrides"].append({"package": name, "version": dist.version, - "requires": str(requirement), "installed": actual}) - else: - result["conflicts"].append( - f"{name} {dist.version} requires {requirement}; installed: {actual}" - ) + result["conflicts"].append( + f"{name} {dist.version} requires {requirement}; installed: {actual}" + ) return result -def expected_versions(): - source = asset_path("requirements/sglang-v0.5.21.txt") - return {name.split("[")[0]: version for line in source.read_text().splitlines() - if line and not line.startswith("#") and "==" in line - for name, version in [line.split("==", 1)]} +def expected_versions(target=False): + source = asset_path("requirements/sglang-v0.5.21.txt" if target else "requirements/draft.txt") + requirements = [Requirement(line) for line in source.read_text().splitlines() + if line.strip() and not line.startswith("#")] + return {requirement.name: next((spec.version for spec in requirement.specifier + if spec.operator == "=="), None) + for requirement in requirements} -def inspect_environment(*, gpu=False): +def inspect_environment(*, gpu=False, target=False): result = {"python": sys.executable, "prefix": sys.prefix, "platform": platform.platform(), - "gpu_requested": gpu, "checks": [], "errors": []} + "gpu_requested": gpu, "target_requested": target, "checks": [], "errors": []} result["dependencies"] = check_dependency_metadata() result["errors"].extend(result["dependencies"]["conflicts"]) if error := venv_error(): @@ -82,24 +73,44 @@ def inspect_environment(*, gpu=False): result["errors"].append("Python >=3.11 is required") if sys.platform != "linux": result["errors"].append("The GPU runtime requires Linux; macOS/Windows training is not validated") - for name, expected in expected_versions().items(): + installed = {canonicalize_name(dist.metadata["Name"]) for dist in importlib.metadata.distributions()} + marker = Path(sys.prefix) / ".aspec-role" + target = target or "ling-aspec-target" in installed or (marker.is_file() and marker.read_text().strip() == "target") + result["role"] = "target" if target else "draft" + if {"ling-aspec", "ling-aspec-target"} <= installed: + result["errors"].append("Install Draft and Target in separate venvs") + has_sglang = False + for name, expected in expected_versions(target).items(): try: actual = importlib.metadata.version(name) except importlib.metadata.PackageNotFoundError: actual = None - ok = actual is not None and actual.split("+", 1)[0] == expected + if name == "sglang": + has_sglang = actual is not None + if not target and not has_sglang: + continue + ok = actual is not None and (expected is None or actual.split("+", 1)[0] == expected) result["checks"].append({"package": name, "expected": expected, "installed": actual, "ok": ok}) if not ok: result["errors"].append(f"{name}: expected {expected}, found {actual}; use the pinned requirements") - try: - site = installed_site() - result["sglang_path"] = str(site / "sglang") - result["patch"] = patch_state(site) - if result["patch"] != "patched": - result["errors"].append("Apply the training patch: aspec patch-sglang") - except (RuntimeError, OSError, importlib.metadata.PackageNotFoundError) as exc: - if str(exc) not in result["errors"]: - result["errors"].append(str(exc)) + if not target: + try: + importlib.metadata.version("sglang") + except importlib.metadata.PackageNotFoundError: + pass + else: + result["errors"].append("SGLang belongs in the isolated Target environment") + result["patch"] = "unavailable" if target or has_sglang else "not_required" + if target or has_sglang: + try: + site = installed_site() + result["sglang_path"] = str(site / "sglang") + result["patch"] = patch_state(site) + if result["patch"] != "patched": + result["errors"].append("Apply the training patch: aspec patch-sglang") + except (RuntimeError, OSError, importlib.metadata.PackageNotFoundError) as exc: + if str(exc) not in result["errors"]: + result["errors"].append(str(exc)) if gpu: try: import torch @@ -125,9 +136,10 @@ def inspect_environment(*, gpu=False): def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--gpu", action="store_true", help="Also allocate a small CUDA tensor and import Mooncake") + parser.add_argument("--target", action="store_true", help="Also require SGLang and its training patch") parser.add_argument("--json", action="store_true") args = parser.parse_args() - result = inspect_environment(gpu=args.gpu) + result = inspect_environment(gpu=args.gpu, target=args.target) if args.json: print(json.dumps(result, indent=2)) else: diff --git a/aspec/commands/quickstart.py b/aspec/commands/quickstart.py index a86b883..536856f 100644 --- a/aspec/commands/quickstart.py +++ b/aspec/commands/quickstart.py @@ -54,10 +54,22 @@ def main(): parser.add_argument("model", choices=MODELS) parser.add_argument("--target", help="Optional compatible local Target directory or Hub ID") parser.add_argument("--output", required=True, type=Path, help="Fresh run directory") + parser.add_argument("--target-address", help="Address of an already running Target service") + parser.add_argument("--managed-target", help="Optionally start local Target using this absolute Python path") + parser.add_argument("--set", action="append", default=[], metavar="KEY=VALUE") parser.add_argument("--steps", type=int, default=8) parser.add_argument("--plan", action="store_true", help="Inspect CPU Plan only; no GPU services") args = parser.parse_args() config = build_profile(args.model, args.target) + if args.target_address: + config.inference.target_service = {"address": args.target_address} + if args.managed_target: + config.inference.target_service = OmegaConf.merge( + config.inference.get("target_service", {}), + {"mode": "managed-local", "python_executable": args.managed_target}, + ) + if args.set: + config = OmegaConf.merge(config, OmegaConf.from_dotlist(args.set)) output = args.output.resolve() if args.plan: with tempfile.TemporaryDirectory(prefix="aspec-plan-") as folder: diff --git a/aspec/commands/recipes.py b/aspec/commands/recipes.py index d1a6c23..5c9c4b8 100644 --- a/aspec/commands/recipes.py +++ b/aspec/commands/recipes.py @@ -265,11 +265,12 @@ def stop_memory_monitor(process, stream): stream.close() -def acceptance(recipe, output, steps): - return run_acceptance(recipe, catalog()[recipe], output, steps, inputs=check_inputs(recipe, acceptance=True)) +def acceptance(recipe, output, steps, overrides=()): + return run_acceptance(recipe, catalog()[recipe], output, steps, + inputs=check_inputs(recipe, acceptance=True), overrides=overrides) -def run_acceptance(recipe, entry, output, steps, *, inputs, prepared_config=None, progress=None): +def run_acceptance(recipe, entry, output, steps, *, inputs, prepared_config=None, progress=None, overrides=()): if steps < 2 or steps % 2: raise ValueError("--steps must be a positive even number >= 2") if entry["nodes"] == 1: @@ -355,6 +356,7 @@ def run_acceptance(recipe, entry, output, steps, *, inputs, prepared_config=None "logging.use_wandb=false", f"dataset.eval_interval={steps // 2}", f"cache_dir={output / 'cache'}", + *overrides, ] announce("Inspecting the execution Plan") execute( @@ -389,7 +391,7 @@ def run_acceptance(recipe, entry, output, steps, *, inputs, prepared_config=None ) for label, limit in (("first", steps // 2), ("resumed", steps)): run = stage / label - overrides = [*base, f"training.num_train_steps={limit}"] + stage_overrides = [*base, f"training.num_train_steps={limit}"] if label == "resumed": if recipe == "shared-target": from omegaconf import OmegaConf @@ -401,9 +403,9 @@ def run_acceptance(recipe, entry, output, steps, *, inputs, prepared_config=None experiment.setdefault("training", {})["load_path"] = str( stage / "first" / experiment["name"] / "checkpoints" ) - overrides.append("experiments=" + json.dumps(experiments)) + stage_overrides.append("experiments=" + json.dumps(experiments)) else: - overrides.append(f"training.load_path={stage / 'first' / 'checkpoints'}") + stage_overrides.append(f"training.load_path={stage / 'first' / 'checkpoints'}") command = [ sys.executable, "-m", @@ -419,7 +421,7 @@ def run_acceptance(recipe, entry, output, steps, *, inputs, prepared_config=None command += ["--cluster-address", os.environ["RAY_ADDRESS"]] announce(f"Training the first {limit} steps" if label == "first" else f"Restoring the checkpoint and continuing to step {limit}") - execute([*command, *overrides], stage / f"{label}.log") + execute([*command, *stage_overrides], stage / f"{label}.log") children = audit_run( run, limit, @@ -476,11 +478,11 @@ def main(): command = sub.add_parser(action) command.add_argument("recipe", choices=catalog()) command.add_argument("--output", required=action != "plan") + command.add_argument("--set", action="append", default=[], metavar="KEY=VALUE") if action == "accept": command.add_argument("--steps", type=int, default=8) else: command.add_argument("--stage", type=int, default=1) - command.add_argument("--set", action="append", default=[], metavar="KEY=VALUE") args = parser.parse_args() if args.action == "list": for name, entry in catalog().items(): @@ -503,7 +505,7 @@ def main(): print(path.resolve()) return if args.action == "accept": - acceptance(args.recipe, args.output, args.steps) + acceptance(args.recipe, args.output, args.steps, args.set) return if not 1 <= args.stage <= len(entry["configs"]): parser.error("Invalid stage for recipe") diff --git a/aspec/compat/sglang_patch.py b/aspec/compat/sglang_patch.py index 0f9c405..292f3b9 100644 --- a/aspec/compat/sglang_patch.py +++ b/aspec/compat/sglang_patch.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Apply/check the pinned training patch in the active, isolated SGLang wheel.""" +"""Apply/check the training patch in the active Target environment.""" import argparse import importlib.metadata @@ -104,7 +104,9 @@ def main(): print(f"Applied ASpec training patch to {site / 'sglang'}") else: print(f"ASpec training patch already present: {site / 'sglang'}") - except (RuntimeError, OSError, importlib.metadata.PackageNotFoundError, subprocess.CalledProcessError) as exc: + except importlib.metadata.PackageNotFoundError: + parser.exit(1, "Activate the Target venv and install ling-aspec-target before patching SGLang\n") + except (ValueError, RuntimeError, OSError, subprocess.CalledProcessError) as exc: parser.exit(1, f"ASpec patch check failed: {exc}\n") diff --git a/aspec/config/inference_config.py b/aspec/config/inference_config.py index ffacc0d..1f8a7b0 100644 --- a/aspec/config/inference_config.py +++ b/aspec/config/inference_config.py @@ -162,8 +162,18 @@ class OfflineTrainingConfig: num_engines: int = 1 +@dataclass +class TargetServiceConfig: + address: str = "http://127.0.0.1:30000" + addresses: Dict[str, str] = field(default_factory=dict) + mode: str = "external" + # Used only by the optional local supervisor, never by the Draft client. + python_executable: Optional[str] = None + + @dataclass class InferenceConfig: + target_service: TargetServiceConfig = field(default_factory=TargetServiceConfig) aux_hidden_states_layers: Optional[list] = None inference_batch_size: int = 1 inference_buffer_threshold: int = 32 diff --git a/aspec/config/train_config.py b/aspec/config/train_config.py index 59108f4..501558b 100644 --- a/aspec/config/train_config.py +++ b/aspec/config/train_config.py @@ -1006,6 +1006,21 @@ def load_config( configs_to_merge.append(cli_config) config = OmegaConf.merge(*configs_to_merge) + from aspec.inference.service_protocol import validate_address + + service = config.inference.target_service + validate_address(service.address) + for address in service.addresses.values(): + validate_address(address) + if service.mode not in ("external", "managed-local"): + raise ValueError("inference.target_service.mode must be external or managed-local") + if service.mode == "managed-local": + if config.cluster.mode != "single_node": + raise ValueError("managed-local requires a single-node training run") + if not service.python_executable or not Path(service.python_executable).is_absolute(): + raise ValueError("managed-local requires an absolute target_service.python_executable") + elif service.python_executable is not None: + raise ValueError("Target Python is configured only for managed-local supervision") _resolve_relative_paths(config, os.getcwd()) _validate_current_scope(config) diff --git a/aspec/controller/loop.py b/aspec/controller/loop.py index a193bc4..0ca8c52 100644 --- a/aspec/controller/loop.py +++ b/aspec/controller/loop.py @@ -615,6 +615,11 @@ def record(component: str, state: str, error: Exception | None = None) -> None: if inference_engines and force_kill_trainers and source_teardown_safe: for engine in inference_engines: try: + if getattr(args, "target_service", None): + # The facade is Draft-owned; the actual source is owned by + # Target. Acknowledged teardown follows the reader fence. + ray.get(engine.shutdown.remote(), timeout=float( + getattr(args, "inference_shutdown_timeout_seconds", 120))) ray.kill(engine, no_restart=True) record("inference_engine", "force_killed") except Exception as exc: diff --git a/aspec/evidence.py b/aspec/evidence.py index 25677fd..ccfc32e 100644 --- a/aspec/evidence.py +++ b/aspec/evidence.py @@ -40,7 +40,9 @@ EVIDENCE_SCHEMA_VERSION = "aspec.run-evidence/v0alpha8" _PACKAGE_NAMES = ( - "aspec", + "ling-aspec", + "ling-aspec-runtime", + "ling-aspec-target", "torch", "ray", "sglang", @@ -179,7 +181,7 @@ def _git(*args: str) -> subprocess.CompletedProcess[str]: build_info = package / "_build_info.json" identity["build"] = json.loads(build_info.read_text()) if build_info.is_file() else None try: - identity["version"] = importlib.metadata.version("aspec") + identity["version"] = importlib.metadata.version("ling-aspec-runtime") except importlib.metadata.PackageNotFoundError: identity["version"] = None source_files = [path for path in package.rglob("*") if path.suffix in @@ -303,6 +305,7 @@ def __init__( "plan": plan_dict["source"], "aspec": _source_identity(), "packages": _package_versions(), + "package_scope": "draft_driver", "backends": {"sglang": _module_source_identity("sglang")}, "host": { "hostname": platform.node(), @@ -427,6 +430,17 @@ def record_actor_initialization( training_actor_statuses=training_actor_statuses, inference_engine_statuses=inference_engine_statuses, ) + target_runtimes = [dict(status["runtime"]) for status in (inference_engine_statuses or ()) + if isinstance(status.get("runtime"), Mapping)] + if target_runtimes: + backend_identities = [item.get("backend_identity") for item in target_runtimes] + if (len(target_runtimes) != inference_engine_count + or any(not isinstance(item, Mapping) or not item.get("effective_tree_sha256") + or item.get("missing_files") for item in backend_identities) + or len({item["effective_tree_sha256"] for item in backend_identities}) != 1): + raise RuntimeError("Target workers reported incomplete or different SGLang identities") + self._data["source"]["backends"]["sglang"] = dict(backend_identities[0]) + self._data["source"]["target_runtimes"] = target_runtimes prior_runtime = self._data["observations"]["runtime"] if dataset_admission is not None: diff --git a/aspec/inference/engine/__init__.py b/aspec/inference/engine/__init__.py index a4eaa7a..ee5b830 100644 --- a/aspec/inference/engine/__init__.py +++ b/aspec/inference/engine/__init__.py @@ -19,10 +19,16 @@ # SOFTWARE. -from aspec.inference.engine.base import InferenceEngine -from aspec.inference.engine.sgl_engine import SglEngine +# ASpec: defer backend imports so Draft never imports the Target dependency graph. +from importlib import import_module -__all__ = [ - "InferenceEngine", - "SglEngine", -] +__all__ = ["InferenceEngine", "SglEngine"] + + +def __getattr__(name): + modules = {"InferenceEngine": "base", "SglEngine": "sgl_engine"} + if name not in modules: + raise AttributeError(name) + value = getattr(import_module("aspec.inference.engine." + modules[name]), name) + globals()[name] = value + return value diff --git a/aspec/inference/engine/target_actor.py b/aspec/inference/engine/target_actor.py new file mode 100644 index 0000000..1fdd10d --- /dev/null +++ b/aspec/inference/engine/target_actor.py @@ -0,0 +1,94 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +"""Lazy Target boundary: the submitting Draft process never imports SGLang. + +Inspired by SpecForge 629092ad's separate capture-server interpreter. Adapted +to ASpec's retained Ray actor contract so placement, storage and cancellation +continue to use the existing SglEngine rather than a second serving pipeline. +""" + +import importlib.metadata +import sys +from pathlib import Path + + +class TargetActor: + def __init__(self, *args, **kwargs): + from aspec.commands.doctor import inspect_environment + from aspec.inference.target_runtime import code_identity + + expected = kwargs.pop("expected_runtime") + self.runtime_identity = code_identity() + if self.runtime_identity != expected: + raise RuntimeError(f"Draft/Target ASpec code differs: {expected} != {self.runtime_identity}") + + result = inspect_environment(target=True) + if not result["ok"]: + raise RuntimeError("Target environment is not ready: " + "; ".join(result["errors"])) + from aspec.inference.engine.sgl_engine import SglEngine + + self.engine = SglEngine(*args, **kwargs) + + def init(self, **kwargs): + return self.engine.init(**kwargs) + + async def generate(self, *args, **kwargs): + return await self.engine.generate(*args, **kwargs) + + def generate_with_decode(self, *args, **kwargs): + return self.engine.generate_with_decode(*args, **kwargs) + + def update_weights_from_disk(self, *args, **kwargs): + return self.engine.update_weights_from_disk(*args, **kwargs) + + def get_node_ip(self): + return self.engine.get_node_ip() + + def find_free_port(self, *args, **kwargs): + return self.engine.find_free_port(*args, **kwargs) + + def get_status(self): + import psutil + import ray + import sglang + import torch + + import aspec + from aspec.runtime_identity import SGLANG_RUNTIME_RELATIVE_PATHS, module_runtime_identity + + backend = module_runtime_identity("sglang", relative_paths=SGLANG_RUNTIME_RELATIVE_PATHS) + backend_identity = { + "path": backend["path"], "git_revision": None, "git_dirty": None, + "effective_tree_sha256": backend["effective_tree_sha256"], + "source_file_count": backend["file_count"], "missing_files": backend["missing_files"], + } + result = self.engine.get_status() + result["created"] = psutil.Process().create_time() + result["runtime"] = { + **self.runtime_identity, + "python": sys.executable, "prefix": sys.prefix, + "aspec": str(Path(aspec.__file__).resolve()), + "ray": str(Path(ray.__file__).resolve()), + "sglang": str(Path(sglang.__file__).resolve()), + "backend_identity": backend_identity, + "sglang_version": importlib.metadata.version("sglang"), + "torch": torch.__version__, "cuda": torch.version.cuda, + "cutlass": importlib.metadata.version("nvidia-cutlass-dsl"), + "quack": importlib.metadata.version("quack-kernels"), + } + return result + + def get_mooncake_config(self): + return self.engine.get_mooncake_config() + + def get_feature_contract_status(self): + return self.engine.get_feature_contract_status() + + async def prepare_feature_contract_status(self): + return await self.engine.prepare_feature_contract_status() + + def health_check(self, **kwargs): + return self.engine.health_check(**kwargs) + + async def shutdown(self): + return await self.engine.shutdown() diff --git a/aspec/inference/factory.py b/aspec/inference/factory.py index 0ff48f7..9203b91 100644 --- a/aspec/inference/factory.py +++ b/aspec/inference/factory.py @@ -242,12 +242,25 @@ def _prepare_sgl_engines( f"({gpus_per_engine} GPU(s) each, nnodes={nnodes}, replicas={num_replicas})" ) - from aspec.inference.engine.sgl_engine import SglEngine + from aspec.inference.service_client import TargetServiceClient pg_obj, reordered_bundle_indices, reordered_gpu_ids = pg - SglRayActor = ray.remote(SglEngine) + SglRayActor = ray.remote(TargetServiceClient) env_vars = get_aspec_env_vars(args) + runtime_env = {"env_vars": env_vars} + service = getattr(args, "target_service", None) or {} + address = service.get("address", "http://127.0.0.1:30000") + addresses = service.get("addresses") or {} + node_addresses = {} + if addresses: + from ray.util.placement_group import placement_group_table + + table = placement_group_table(pg_obj) + nodes = {node["NodeID"]: node["NodeManagerAddress"] for node in ray.nodes() if node["Alive"]} + node_addresses = {int(bundle): addresses.get(nodes[node], address) + for bundle, node in table["bundles_to_node_id"].items() + if int(bundle) in reordered_bundle_indices} # Step 1: Create all engine actors (without calling init yet) engines = [] @@ -267,8 +280,9 @@ def _prepare_sgl_engines( num_cpus=0.2, num_gpus=0.2, scheduling_strategy=scheduling_strategy, - runtime_env={"env_vars": env_vars}, + runtime_env=runtime_env, ).remote( + address=node_addresses.get(reordered_bundle_indices[bundle_offset], address), args=args, rank=i, base_gpu_id=base_gpu_id, diff --git a/aspec/inference/managed.py b/aspec/inference/managed.py new file mode 100644 index 0000000..3d14ecc --- /dev/null +++ b/aspec/inference/managed.py @@ -0,0 +1,81 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +"""Optional local supervision of the same external Target service. + +Inspired by SpecForge f8cc3dc managed-local readiness and process ownership. +ASpec keeps its full checkpoint/resume path and its existing reader fence. +""" +import json +import logging +import os +import signal +import subprocess +import time +from contextlib import contextmanager +from pathlib import Path +from urllib.parse import urlsplit + +from aspec.inference.service_protocol import request +from aspec.runtime_control import termination_signal_handlers + +logger = logging.getLogger(__name__) + + +@contextmanager +def managed_target(args): + settings = getattr(args, "target_service", None) or {} + if settings.get("mode", "external") != "managed-local": + yield + return + address = settings["address"] + if urlsplit(address).hostname not in ("localhost", "127.0.0.1", "::1"): + raise ValueError("managed-local Target must listen on loopback") + try: + request(address, "/v1/info", timeout=2) + except OSError: + pass + else: + raise RuntimeError("managed-local refuses to replace an existing Target service") + root = Path(args.output_dir) + root.mkdir(parents=True, exist_ok=True) + config = root / "target-service.json" + config.write_text(json.dumps({"model": {"target_model_path": args.target_model_path}})) + command = [settings["python_executable"], "-m", "aspec._target_worker", "-m", + "aspec.inference.service", "--config", str(config), "--address", address] + log_path = root / "target-service.log" + with log_path.open("w") as log: + process = subprocess.Popen(command, stdout=log, stderr=subprocess.STDOUT, + start_new_session=True, env=dict(os.environ)) + ready = False + try: + deadline = time.monotonic() + float(getattr(args, "sglang_init_timeout", 300)) + while time.monotonic() < deadline: + if process.poll() is not None: + raise RuntimeError(f"Target exited before readiness; see {log_path}") + try: + info = request(address, "/v1/info", timeout=2) + if info.get("draining"): + raise RuntimeError("Target is draining during startup") + if info.get("pid") != process.pid: + raise RuntimeError("Target address belongs to another process") + ready = True + break + except OSError: + time.sleep(0.25) + if not ready: + raise TimeoutError(f"Target readiness timed out; see {log_path}") + yield + finally: + if process.poll() is None: + process.send_signal(signal.SIGTERM) + with termination_signal_handlers(defer=True): + # Target retains disconnected source storage until its + # registered read processes have observably exited. The + # supervisor must not bypass that guard with SIGKILL. + while process.poll() is None: + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + logger.info("Waiting for owned Target cleanup; see %s", log_path) + if ready and process.returncode != 0: + logger.error("Owned Target exited with %s; see %s", process.returncode, log_path) diff --git a/aspec/inference/service.py b/aspec/inference/service.py new file mode 100644 index 0000000..89c920b --- /dev/null +++ b/aspec/inference/service.py @@ -0,0 +1,402 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +"""Standalone Target control service around the retained SglEngine. + +Inspired by SpecForge f8cc3dc external capture. JSON control messages replace +cross-environment Ray calls; ASpec's GPU leases, feature contracts and tensor +transport remain in use. No SpecForge server code is copied. +""" +import argparse +import json +import logging +import os +import signal +import tempfile +import threading +import time +import uuid +from argparse import Namespace +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from urllib.parse import urlsplit + +import ray + +from aspec.inference.service_protocol import ( + MAX_BODY, + METHODS, + PROTOCOL, + decode, + encode, + validate_address, +) +from aspec.inference.target_runtime import code_identity, inventory +from aspec.utils.env import get_aspec_env_vars + +logger = logging.getLogger(__name__) + + +def readers_exited(owners, local_ip): + """Only observe local processes; unknown/remote owners retain source storage.""" + import psutil + + if not owners: + return False + for owner in owners: + if (owner.get("node_ip") != local_ip or not isinstance(owner.get("pid"), int) + or not isinstance(owner.get("created"), (int, float))): + return False + try: + process = psutil.Process(owner["pid"]) + if (process.create_time() == owner["created"] + and process.status() not in (psutil.STATUS_ZOMBIE, psutil.STATUS_DEAD)): + return False + except psutil.NoSuchProcess: + pass + return True + + +class TargetService: + def __init__(self, *, model=None, devices=None, runtime=None): + self.runtime = runtime if runtime is not None else inventory() + self.identity = code_identity() + self.model = model + self.devices = None if devices is None else frozenset(devices) + self.server_id = uuid.uuid4().hex + self.lock = threading.RLock() + self.engines = {} + self.claims = {} + self.draining = False + + def info(self): + with self.lock: + return {"protocol": PROTOCOL, "identity": self.identity, "runtime": self.runtime, + "pid": os.getpid(), + "server_id": self.server_id, "model": self.model, + "devices": sorted(self.devices) if self.devices is not None else None, + "draining": self.draining, + "engines": {key: {k: item[k] for k in ("state", "session", "gpus")} + for key, item in self.engines.items()}} + + def create(self, value): + from aspec.inference.engine.target_actor import TargetActor + + if value["identity"] != self.identity: + raise ValueError("Draft/Target ASpec code identity differs") + args = Namespace(**decode(value["arguments"])) + if self.model and args.target_model_path != self.model: + raise ValueError("Target model differs from the service configuration") + session = getattr(args, "aspec_attempt_id", None) + run_id = getattr(args, "aspec_run_id", None) + if not isinstance(session, str) or not session or not isinstance(run_id, str) or not run_id: + raise ValueError("Target requires explicit run and attempt identities") + base, count = value["base_gpu_id"], value["num_gpus_per_engine"] + if type(base) is not int or type(count) is not int or base < 0 or count <= 0: + raise ValueError("Invalid Target GPU allocation") + gpus = list(range(base, base + count)) + key = value["key"] + with self.lock: + if self.draining: + raise RuntimeError("Target service is draining") + if key in self.engines: + raise ValueError("Target engine identity already exists; creation is not replayed") + if self.devices is not None and not set(gpus).issubset(self.devices): + raise ValueError("Target allocation is outside the service GPU set") + if any(gpu in self.claims for gpu in gpus): + raise RuntimeError("Target GPU already belongs to another engine/session") + for old_key in list(self.engines): + if len(self.engines) < 128: + break + if self.engines[old_key]["state"] == "closed": + del self.engines[old_key] + # Draft's placement group reserves these physical GPUs. This private + # Ray instance only orchestrates Target processes; it owns no second + # GPU allocator and cannot schedule an overlapping training worker. + environment = get_aspec_env_vars(args) + environment.update(session_environment(args)) + actor = ray.remote(TargetActor).options( + num_cpus=0.2, num_gpus=0, runtime_env={"env_vars": environment}, + ).remote(expected_runtime=self.identity, args=args, rank=value["rank"], + base_gpu_id=base, num_gpus_per_engine=count, + node_rank=value["node_rank"], engine_group=value["engine_group"]) + self.engines[key] = {"actor": actor, "state": "created", "session": session, + "gpus": gpus, "heartbeat": time.monotonic(), + "owners": getattr(args, "aspec_read_owners", None), + "status": {}, "features": {}, "active": 1, "processes": {}} + self.claims.update({gpu: key for gpu in gpus}) + try: + status = ray.get(actor.get_status.remote(), timeout=120) + with self.lock: + self.engines[key]["status"] = status + if status.get("pid") and status.get("created"): + self.engines[key]["processes"][status["pid"]] = status["created"] + except BaseException: + # SglEngine's constructor loads no model or source buffers. Failure + # here precedes init/publication, so this creation can be rolled back. + ray.kill(actor, no_restart=True) + with self.lock: + self.engines[key]["state"] = "closed" + for gpu in gpus: + self.claims.pop(gpu) + raise + finally: + with self.lock: + self.engines[key]["active"] -= 1 + return {"key": key} + + def heartbeat(self, key): + with self.lock: + item = self.engines[key] + if item["state"] in ("retained", "closing"): + raise RuntimeError("Target session has stopped admitting work") + item["heartbeat"] = time.monotonic() + return None + + def call(self, key, method, args, kwargs): + if method not in METHODS: + raise ValueError("Unknown Target control method") + with self.lock: + item = self.engines[key] + if item["state"] == "closed": + if method == "shutdown": + return None + if method == "get_status": + return {**item["status"], "initialized": False} + if method in ("get_feature_contract_status", "prepare_feature_contract_status"): + return item["features"] + raise RuntimeError("Target engine is closed") + if item["state"] in ("retained", "closing") and method != "shutdown": + raise RuntimeError("Target engine is retained pending safe cleanup") + if method == "shutdown": + if item["active"]: + raise RuntimeError("Target still has active calls; drain before shutdown") + item["state"] = "closing" + item["heartbeat"] = time.monotonic() + item["active"] += 1 + try: + if method == "init" and kwargs.get("mooncake_config") is not None: + from aspec.config.mooncake_config import MooncakeConfig + + kwargs["mooncake_config"] = MooncakeConfig(**kwargs["mooncake_config"]) + result = ray.get(getattr(item["actor"], method).remote(*args, **kwargs), timeout=1800) + with self.lock: + if method == "get_status": + item["status"] = result + if method in ("get_feature_contract_status", "prepare_feature_contract_status"): + item["features"] = result + if method == "init": + item["state"] = "ready" + if method == "shutdown": + ray.kill(item["actor"], no_restart=True) + item["state"] = "closed" + for gpu in item["gpus"]: + self.claims.pop(gpu) + logger.info("Target session %s engine %s closed", item["session"], key) + return result + except BaseException: + if method == "shutdown": + with self.lock: + item["state"] = "retained" + raise + finally: + with self.lock: + item["active"] -= 1 + + def reap_disconnected(self, *, expiry=30): + """Fail closed on lost control; reclaim only after observed reader exit.""" + from aspec.utils.misc import get_current_node_ip + + with self.lock: + # Record descendants while their parent is still alive. A backend + # scheduler can kill its actor on initialization failure; later + # cleanup must identify only this service's own remaining children. + for item in self.engines.values(): + if item["state"] != "closed": + observe_process_tree(item["processes"]) + candidates = [(key, item) for key, item in self.engines.items() + if item["state"] not in ("closed", "closing") + and item["active"] == 0 + and time.monotonic() - item["heartbeat"] > expiry] + for _, item in candidates: + item["state"] = "retained" + for key, item in candidates: + if readers_exited(item["owners"], get_current_node_ip()): + logger.warning("Fenced disconnected session %s; cancelling Target %s", item["session"], key) + try: + self.call(key, "shutdown", [], {}) + except ray.exceptions.RayActorError: + # The read fence is established above. Observe exact source + # process identities before retiring a dead actor's lease. + if not item["processes"]: + raise + stop_owned_processes(item["processes"]) + with self.lock: + item["state"] = "closed" + item["terminal_reason"] = "actor_failed_readers_fenced_sources_exited" + for gpu in item["gpus"]: + self.claims.pop(gpu, None) + logger.warning("Retired failed Target %s after observed source exit", key) + + def idle(self): + with self.lock: + return not self.claims + + +def observe_process_tree(identities): + import psutil + + for pid, created in list(identities.items()): + try: + process = psutil.Process(pid) + if process.create_time() == created: + for child in process.children(recursive=True): + try: + identities[child.pid] = child.create_time() + except psutil.NoSuchProcess: + pass + except psutil.NoSuchProcess: + pass + + +def session_environment(args): + """Native scheduler children inherit this session, never the service's job.""" + contracts = getattr(args, "aspec_shared_consumer_contracts", None) or {} + return { + "ASPEC_RUN_ID": args.aspec_run_id, + "ASPEC_ATTEMPT_ID": args.aspec_attempt_id, + "ASPEC_SHARED_SOURCE": str(int(bool(getattr(args, "aspec_shared_source", False)))), + "ASPEC_SHARED_FEATURE_VIEWS": str(int(len(set(contracts.values())) > 1)), + } + + +def stop_owned_processes(identities): + """Called only after the read fence; never discover or stop unrelated PIDs.""" + import psutil + + processes = [] + for pid, created in identities.items(): + try: + process = psutil.Process(pid) + if process.create_time() == created: + processes.append(process) + process.terminate() + except psutil.NoSuchProcess: + pass + _, alive = psutil.wait_procs(processes, timeout=5) + for process in alive: + try: + if process.create_time() == identities[process.pid]: + process.kill() + except psutil.NoSuchProcess: + pass + _, alive = psutil.wait_procs(alive, timeout=5) + if alive: + raise RuntimeError("Owned Target source process exit remains unverified") + + +def http_server(service, host, port): + class Handler(BaseHTTPRequestHandler): + def log_message(self, *_): + pass + + def respond(self, status, result): + data = json.dumps({"protocol": PROTOCOL, "result": encode(result)}, allow_nan=False).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(data))) + self.end_headers() + try: + self.wfile.write(data) + except (BrokenPipeError, ConnectionResetError): + logger.debug("Target caller disconnected before reading its response") + + def do_GET(self): + if self.path != "/v1/info": + self.respond(404, "Unknown Target endpoint") + return + self.respond(200, service.info()) + + def do_POST(self): + try: + length = int(self.headers.get("Content-Length", "0")) + if not 0 < length <= MAX_BODY: + raise ValueError("Invalid Target request size") + value = decode(json.loads(self.rfile.read(length))) + if self.path == "/v1/engines": + result = service.create(value) + elif self.path == "/v1/heartbeat": + result = service.heartbeat(value["key"]) + elif self.path == "/v1/call": + result = service.call(value["key"], value["method"], value["args"], value["kwargs"]) + else: + raise ValueError("Unknown Target endpoint") + self.respond(200, result) + except Exception as error: + logger.exception("Target control request failed") + self.respond(409, f"{type(error).__name__}: {error}") + + server = ThreadingHTTPServer((host, port), Handler) + server.daemon_threads = True + return server + + +def main(): + from omegaconf import OmegaConf + + logging.basicConfig(level=logging.INFO) + parser = argparse.ArgumentParser(description="Run Target in this environment; Draft connects by address.") + parser.add_argument("--config", type=Path) + parser.add_argument("--address", help="Listening http://host:port; defaults to the config or localhost:30000") + parser.add_argument("--gpus", help="Allowed physical GPU IDs, comma separated") + args = parser.parse_args() + config = OmegaConf.load(args.config) if args.config else OmegaConf.create({}) + address = validate_address(args.address or OmegaConf.select( + config, "inference.target_service.address", default="http://127.0.0.1:30000")) + model = OmegaConf.select(config, "model.target_model_path") + devices = [int(gpu) for gpu in args.gpus.split(",")] if args.gpus else None + runtime = inventory() # Check this environment and patch before starting processes. + service = TargetService(model=model, devices=devices, runtime=runtime) + parsed = urlsplit(address) + server = http_server(service, parsed.hostname, parsed.port) + stop = threading.Event() + + def stopping(*_): + service.draining = True + stop.set() + + previous = {sig: signal.signal(sig, stopping) for sig in (signal.SIGINT, signal.SIGTERM)} + os.environ["RAY_USAGE_STATS_ENABLED"] = "0" + temporary = tempfile.mkdtemp(prefix="aspec-t-", dir="/tmp") + try: + ray.init(address="local", num_gpus=0, num_cpus=min(8, os.cpu_count() or 1), + include_dashboard=False, object_store_memory=256 * 1024 * 1024, + _temp_dir=temporary) + # Ray may replace the process SIGTERM handler during initialization. + for sig in previous: + signal.signal(sig, stopping) + worker = threading.Thread(target=server.serve_forever, daemon=True) + worker.start() + print(json.dumps({"state": "ready", "address": address, "protocol": PROTOCOL, + "python": runtime["python"], "pid": os.getpid()}), flush=True) + while not stop.wait(1) or not service.idle(): + service.reap_disconnected() + server.shutdown() + worker.join() + finally: + # Active sources must survive until readers are fenced, including during + # signal handling. Never let Ray's implicit driver exit reclaim them. + while not service.idle(): + try: + service.draining = True + service.reap_disconnected() + time.sleep(1) + except BaseException: + logger.exception("Target cleanup is incomplete; retaining the service") + server.server_close() + ray.shutdown() + for sig, handler in previous.items(): + signal.signal(sig, handler) + + +if __name__ == "__main__": + main() diff --git a/aspec/inference/service_client.py b/aspec/inference/service_client.py new file mode 100644 index 0000000..0b70860 --- /dev/null +++ b/aspec/inference/service_client.py @@ -0,0 +1,91 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +"""Draft-owned Ray facade for an independently launched Target service.""" +import asyncio +import threading +import uuid + +import ray + +from aspec.inference.service_protocol import PROTOCOL, engine_arguments, request +from aspec.inference.target_runtime import code_identity + + +class TargetServiceClient: + def __init__(self, *, address, args, rank, base_gpu_id, num_gpus_per_engine, + node_rank=0, engine_group=0): + self.address = address + self.key = uuid.uuid4().hex + self.timeout = float(getattr(args, "sglang_init_timeout", 1200)) + 60 + self._closed = threading.Event() + identity = code_identity() + info = request(address, "/v1/info", timeout=30) + if info["protocol"] != PROTOCOL or info["identity"] != identity: + raise RuntimeError("Target service has different ASpec code or protocol") + request(address, "/v1/engines", { + "key": self.key, "identity": identity, "arguments": engine_arguments(args), + "rank": rank, "base_gpu_id": base_gpu_id, + "num_gpus_per_engine": num_gpus_per_engine, "node_rank": node_rank, + "engine_group": engine_group, + }, timeout=self.timeout) + self._heartbeat = threading.Thread(target=self._keepalive, daemon=True) + self._heartbeat.start() + + def _keepalive(self): + while not self._closed.wait(5): + try: + request(self.address, "/v1/heartbeat", {"key": self.key}, timeout=5) + except Exception: + # Calls fail visibly; a heartbeat never retries a generation request. + pass + + def _call(self, method, *args, **kwargs): + return request(self.address, "/v1/call", { + "key": self.key, "method": method, "args": args, "kwargs": kwargs, + }, timeout=self.timeout) + + def init(self, **kwargs): + return self._call("init", **kwargs) + + async def generate(self, *args, **kwargs): + # ObjectRefs are resolved inside Draft's Ray job, never serialized to Target. + if isinstance(kwargs.get("input_ids_ref"), ray.ObjectRef): + kwargs["input_ids_ref"] = await kwargs["input_ids_ref"] + return await asyncio.to_thread(self._call, "generate", *args, **kwargs) + + def generate_with_decode(self, *args, **kwargs): + return self._call("generate_with_decode", *args, **kwargs) + + def update_weights_from_disk(self, *args, **kwargs): + return self._call("update_weights_from_disk", *args, **kwargs) + + def get_node_ip(self): + return self._call("get_node_ip") + + def find_free_port(self, *args, **kwargs): + return self._call("find_free_port", *args, **kwargs) + + def get_status(self): + status = self._call("get_status") + status["target_service"] = {"address": self.address, "protocol": PROTOCOL} + return status + + def get_mooncake_config(self): + from aspec.config.mooncake_config import MooncakeConfig + + value = self._call("get_mooncake_config") + return MooncakeConfig(**value) if value else None + + def get_feature_contract_status(self): + return self._call("get_feature_contract_status") + + async def prepare_feature_contract_status(self): + return await asyncio.to_thread(self._call, "prepare_feature_contract_status") + + def health_check(self, **kwargs): + return self._call("health_check", **kwargs) + + async def shutdown(self): + result = await asyncio.to_thread(self._call, "shutdown") + self._closed.set() + return result diff --git a/aspec/inference/service_protocol.py b/aspec/inference/service_protocol.py new file mode 100644 index 0000000..a037ef1 --- /dev/null +++ b/aspec/inference/service_protocol.py @@ -0,0 +1,113 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +"""Versioned JSON control plane; feature tensors stay in Mooncake. + +Inspired by SpecForge f8cc3dc's external capture adapter. ASpec retains its +SglEngine, feature contracts and transport rather than copying that adapter. +""" +import json +from dataclasses import asdict, is_dataclass +from urllib.error import HTTPError +from urllib.parse import urlsplit +from urllib.request import ProxyHandler, Request, build_opener + +PROTOCOL = "aspec.target-service/v1" +MAX_BODY = 64 * 1024 * 1024 +METHODS = frozenset({ + "init", "generate", "generate_with_decode", "update_weights_from_disk", + "get_node_ip", "find_free_port", "get_status", "get_mooncake_config", + "get_feature_contract_status", "prepare_feature_contract_status", "health_check", "shutdown", +}) + + +def validate_address(value): + parsed = urlsplit(value) + if (parsed.scheme != "http" or not parsed.hostname or parsed.username or parsed.password + or parsed.path not in ("", "/") or parsed.query or parsed.fragment): + raise ValueError("Target address must be an http://host:port origin") + if parsed.port is None: + raise ValueError("Target address must include a port") + return value.rstrip("/") + + +def encode(value): + """Encode only explicit data types, never Python objects or feature tensors.""" + from omegaconf import OmegaConf + + if OmegaConf.is_config(value): + value = OmegaConf.to_container(value, resolve=True) + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, dict): + return {str(k): encode(v) for k, v in value.items()} + if isinstance(value, (list, tuple)): + return [encode(v) for v in value] + if is_dataclass(value): + return encode(asdict(value)) + # Importing torch is confined to actual control payload conversion. + import torch + + if isinstance(value, torch.dtype): + return {"$dtype": str(value).removeprefix("torch.")} + if isinstance(value, torch.Tensor): + if value.device.type != "cpu" or value.dtype not in (torch.int64, torch.int32): + raise TypeError("Only CPU token IDs may cross the Target control plane") + return {"$tokens": value.tolist(), "dtype": str(value.dtype).removeprefix("torch.")} + raise TypeError(f"Unsupported Target control value: {type(value).__name__}") + + +def decode(value): + if isinstance(value, list): + return [decode(v) for v in value] + if isinstance(value, dict): + if set(value) == {"$dtype"}: + import torch + + dtype = getattr(torch, value["$dtype"], None) + if not isinstance(dtype, torch.dtype): + raise ValueError("Invalid wire dtype") + return dtype + if set(value) == {"$tokens", "dtype"}: + import torch + + if value["dtype"] not in ("int64", "int32"): + raise ValueError("Invalid token dtype") + return torch.tensor(value["$tokens"], dtype=getattr(torch, value["dtype"])) + return {k: decode(v) for k, v in value.items()} + return value + + +def request(address, path, payload=None, *, timeout=1200): + address = validate_address(address) + body = None if payload is None else json.dumps(encode(payload), allow_nan=False).encode() + if body is not None and len(body) > MAX_BODY: + raise ValueError("Target control request exceeds the size limit") + req = Request(address + path, data=body, headers={"Content-Type": "application/json"}) + # Cluster-internal control endpoints must not be sent through HTTP proxies. + opener = build_opener(ProxyHandler({})) + try: + with opener.open(req, timeout=timeout) as response: + raw = response.read(MAX_BODY + 1) + except HTTPError as error: + detail = error.read(8192).decode(errors="replace") + raise RuntimeError(f"Target {address}{path}: {detail}") from error + if len(raw) > MAX_BODY: + raise ValueError("Target control response exceeds the size limit") + result = json.loads(raw) + if result.get("protocol") != PROTOCOL: + raise RuntimeError("Target service protocol differs") + return decode(result["result"]) + + +def engine_arguments(args): + # Driver-only Python model/config objects are deliberately not sent to Target. + result = {} + for key, value in vars(args).items(): + if key.startswith("_") or key == "draft_model_config_obj": + continue + try: + result[key] = encode(value) + except TypeError: + # Unknown objects cannot become an implicit cross-environment API. + raise TypeError(f"Non-data runtime argument at Target boundary: {key}") from None + return result diff --git a/aspec/inference/target_runtime.py b/aspec/inference/target_runtime.py new file mode 100644 index 0000000..b4ebd40 --- /dev/null +++ b/aspec/inference/target_runtime.py @@ -0,0 +1,51 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +"""Identity shared across the independently installed Draft and Target roles.""" +import hashlib +import importlib.metadata +import json +import sys +from pathlib import Path + + +def code_identity(): + package = Path(__file__).resolve().parents[1] + digest = hashlib.sha256() + for path in sorted(package.rglob("*")): + if path.suffix in {".py", ".patch"}: + digest.update(path.relative_to(package).as_posix().encode()) + digest.update(b"\0") + digest.update(path.read_bytes()) + try: + version = importlib.metadata.version("ling-aspec-runtime") + except importlib.metadata.PackageNotFoundError: + version = None + return {"version": version, "code_sha256": digest.hexdigest(), + "python_minor": list(sys.version_info[:2])} + + +def dependency_versions(*names): + result = {} + for name in names: + try: + result[name] = importlib.metadata.version(name) + except importlib.metadata.PackageNotFoundError: + result[name] = None + return result + + +def inventory(): + """Read Target identity in its own interpreter, without starting GPU services.""" + from aspec.commands.doctor import inspect_environment + from aspec.runtime_identity import SGLANG_RUNTIME_RELATIVE_PATHS, module_runtime_identity + + doctor = inspect_environment(target=True) + if not doctor["ok"]: + raise RuntimeError("Target environment is not ready: " + "; ".join(doctor["errors"])) + return {**code_identity(), "python": sys.executable, "prefix": sys.prefix, + "sglang": module_runtime_identity("sglang", relative_paths=SGLANG_RUNTIME_RELATIVE_PATHS), + "dependencies": dependency_versions("torch", "ray", "nvidia-cutlass-dsl", "quack-kernels")} + + +if __name__ == "__main__": + print(json.dumps(inventory())) diff --git a/aspec/multi_draft.py b/aspec/multi_draft.py index 6ff79ca..95df810 100644 --- a/aspec/multi_draft.py +++ b/aspec/multi_draft.py @@ -492,6 +492,9 @@ def persist_session(): destination.gpu_buffer_size = build_mooncake_config(child).gpu_buffer_size destination_configs[name] = destination init_refs[name] = group.async_init(child, "training", mooncake_config=destination) + source_args.aspec_read_owners = [ + owner for statuses in actor_statuses.values() for owner in statuses + ] engines, engine_refs = prepare_inference_engines(source_args, pgs["inference"], mooncake) ray.get(engine_refs) # Target initialization failure is a shared fault. producer_placements = ray.get([engine.get_status.remote() for engine in engines]) diff --git a/aspec/offline/generate.py b/aspec/offline/generate.py index 950c26c..a902b2c 100644 --- a/aspec/offline/generate.py +++ b/aspec/offline/generate.py @@ -31,6 +31,7 @@ import json import os import time +import uuid from pathlib import Path import ray @@ -165,6 +166,8 @@ def materialize(args, output_dir: str) -> dict[str, int]: if args.draft_context_parallel_size != 1 or args.dflash_packing: raise ValueError("Materialization requires CP1 and dflash_packing=false") args.num_epochs = 1 + args.aspec_run_id = f"uuid:{uuid.uuid4()}" + args.aspec_attempt_id = f"uuid:{uuid.uuid4()}" # Materialize the complete admitted source once, independently of the # optimizer-step and producer-profiling budgets in a reused training YAML. args.num_train_steps = None @@ -209,10 +212,10 @@ def get(ref): inference_pg = create_placement_groups(args, roles={"inference"})["inference"] launch_mooncake_master(args) mooncake_config = build_mooncake_config(args) - engines = create_inference_engines(args, inference_pg, mooncake_config) - saver = OfflineSavingActor.remote(output_dir, mooncake_config, controller) saver_owner = get(saver.get_owner_status.remote()) + args.aspec_read_owners = [saver_owner] + engines = create_inference_engines(args, inference_pg, mooncake_config) train_queues, eval_queues = get( [controller.get_train_queues.remote(), controller.get_eval_queues.remote()] ) diff --git a/aspec/offline/saving_actor.py b/aspec/offline/saving_actor.py index 1ad1ae6..3ee7d80 100644 --- a/aspec/offline/saving_actor.py +++ b/aspec/offline/saving_actor.py @@ -58,10 +58,13 @@ def __init__(self, output_dir, mooncake_config, controller) -> None: self.queues = {} def get_owner_status(self): + import psutil + return { "node_id": ray.get_runtime_context().get_node_id(), "node_ip": RayActor.get_node_ip(), "pid": os.getpid(), + "created": psutil.Process().create_time(), } def set_queues(self, train_queue, eval_queue) -> None: diff --git a/aspec/planning.py b/aspec/planning.py index 7515a0e..a4b61d4 100644 --- a/aspec/planning.py +++ b/aspec/planning.py @@ -1589,6 +1589,7 @@ def _policy( "engines": (runtime["offline_num_engines"] if runtime["inference_engine_type"] == "offline" else runtime["inference_num_gpus"] // runtime["inference_num_gpus_per_engine"]), "target_group_locality": "node_local", + "target_service": _select(config, "inference.target_service"), }, "locality_domains": locality_domains, }, diff --git a/aspec/ray/cluster_preflight.py b/aspec/ray/cluster_preflight.py index a2570d6..4860c92 100644 --- a/aspec/ray/cluster_preflight.py +++ b/aspec/ray/cluster_preflight.py @@ -283,7 +283,9 @@ def _filesystem_capacity(path: Path) -> dict[str, int]: def _collect_local_inventory( required_paths: Sequence[str], marker_path: str, marker_sha256: str, - storage_mode: str = "shared", output_dir: str | None = None + storage_mode: str = "shared", output_dir: str | None = None, + target_address: str | None = None, + roles: Sequence[str] | None = None, ) -> dict[str, Any]: marker = Path(marker_path) if storage_mode == "node_local": @@ -310,6 +312,16 @@ def _collect_local_inventory( p.stat().st_size, file_digest(p)) for p in files] item["content_sha256"] = hashlib.sha256(json.dumps(identity, separators=(",", ":")).encode()).hexdigest() memlock_soft, memlock_hard = resource.getrlimit(resource.RLIMIT_MEMLOCK) + target_runtime = None + if target_address: + from aspec.inference.service_protocol import request + from aspec.inference.target_runtime import code_identity + + info = request(target_address, "/v1/info", timeout=30) + target_runtime = info["runtime"] + expected = code_identity() + if info["identity"] != expected: + raise ClusterPreflightError("Target and Draft ASpec code identity differs") return { "node_id": ray.get_runtime_context().get_node_id(), "node_ip": ray.util.get_node_ip_address(), @@ -317,8 +329,10 @@ def _collect_local_inventory( "python": sys.version.split()[0], "source_tree_sha256": _source_tree_sha256(), "packages": _package_versions(), + "roles": list(roles) if roles is not None else None, + "target_runtime": target_runtime, "runtime_artifacts": { - "sglang": module_runtime_identity( + "sglang": target_runtime["sglang"] if target_runtime else module_runtime_identity( "sglang", relative_paths=SGLANG_RUNTIME_RELATIVE_PATHS ), "mooncake": module_runtime_identity("mooncake.engine"), @@ -379,17 +393,31 @@ def validate_cluster_inventories( source_hashes = {node.get("source_tree_sha256") for node in inventories} if len(source_hashes) != 1 or None in source_hashes: raise ClusterPreflightError("ASpec source identity differs across nodes") - package_sets = { - json.dumps(node.get("packages"), sort_keys=True) for node in inventories - } - if len(package_sets) != 1: + package_sets = defaultdict(set) + for node in inventories: + roles = node.get("roles") + if roles is not None and (not roles or set(roles) - {"draft", "target"}): + raise ClusterPreflightError("Unknown or empty node runtime roles") + group = "legacy" if roles is None else "draft" if "draft" in roles else "target" + package_sets[group].add(json.dumps(node.get("packages"), sort_keys=True)) + if roles and "target" in roles and not node.get("target_runtime"): + raise ClusterPreflightError("Target node did not report its isolated runtime") + if any(len(versions) != 1 for versions in package_sets.values()): raise ClusterPreflightError("required package versions differ across nodes") - runtime_artifact_sets = set() + if len({(node.get("packages") or {}).get("ray") for node in inventories}) != 1: + raise ClusterPreflightError("Ray versions differ across runtime roles") + target_sets = {json.dumps({key: value for key, value in node["target_runtime"].items() + if key not in {"python", "prefix", "sglang"}}, sort_keys=True) + for node in inventories if node.get("target_runtime")} + if len(target_sets) > 1: + raise ClusterPreflightError("Target runtime identities differ across nodes") + runtime_artifact_sets = defaultdict(set) for node in inventories: artifacts = node.get("runtime_artifacts") or {} + names = ("sglang", "mooncake") if node.get("roles") is None or "target" in node["roles"] else ("mooncake",) incomplete = { name: identity - for name in ("sglang", "mooncake") + for name in names if not isinstance((identity := artifacts.get(name)), Mapping) or not identity.get("effective_tree_sha256") or identity.get("missing_files") @@ -399,21 +427,10 @@ def validate_cluster_inventories( "required runtime artifact identity is incomplete on node " f"{node.get('node_ip')}: {incomplete}" ) - runtime_artifact_sets.add( - json.dumps( - { - name: { - "effective_tree_sha256": identity.get( - "effective_tree_sha256" - ), - "file_count": identity.get("file_count"), - } - for name, identity in artifacts.items() - }, - sort_keys=True, - ) - ) - if len(runtime_artifact_sets) != 1: + for name in names: + identity = artifacts[name] + runtime_artifact_sets[name].add((identity["effective_tree_sha256"], identity.get("file_count"))) + if any(len(identities) != 1 for identities in runtime_artifact_sets.values()): raise ClusterPreflightError( "effective SGLang/Mooncake runtime artifacts differ across nodes" ) @@ -1153,13 +1170,20 @@ def run_cluster_preflight(args) -> dict[str, Any] | None: try: for node_ip in sorted(expected_ips): node_id = ray_nodes_by_ip[node_ip]["NodeID"] + roles = (["draft"] if node_ip in (getattr(args, "training_node_ips", None) or ()) else []) + if node_ip in (getattr(args, "inference_node_ips", None) or ()): + roles.append("target") + service = getattr(args, "target_service", None) or {} + target_address = ((service.get("addresses") or {}).get(node_ip, service.get( + "address", "http://127.0.0.1:30000")) if "target" in roles else None) refs.append( probe.options( scheduling_strategy=NodeAffinitySchedulingStrategy( node_id=node_id, soft=False ) ).remote(_required_paths(args), str(marker), marker_sha256, - getattr(args, "cluster_storage_mode", "shared"), str(output_dir)) + getattr(args, "cluster_storage_mode", "shared"), str(output_dir), + target_address, roles) ) inventories = ray.get( refs, diff --git a/aspec/ray/placement_group.py b/aspec/ray/placement_group.py index f3db4e1..80784be 100644 --- a/aspec/ray/placement_group.py +++ b/aspec/ray/placement_group.py @@ -23,6 +23,7 @@ import math import os import socket +import tempfile import time from collections.abc import Mapping, Sequence from dataclasses import dataclass @@ -334,7 +335,13 @@ def _ensure_ray_initialized(): if ray.is_initialized(): return - ray_address = os.environ.get("RAY_ADDRESS", "auto") + # Target owns a separate Ray runtime. Do not auto-attach Draft to it (or + # another user's local cluster) merely because it was started last. + ray_address = os.environ.get("RAY_ADDRESS", "local") + if ray_address == "local": + ray.init(address="local", ignore_reinit_error=True, include_dashboard=False, + _temp_dir=tempfile.mkdtemp(prefix="aspec-d-", dir="/tmp")) + return try: ray.init(address=ray_address, ignore_reinit_error=True) logger.info(f"Connected to Ray cluster at {ray_address}") diff --git a/aspec/ray/ray_actor.py b/aspec/ray/ray_actor.py index b054f3b..251b1ad 100644 --- a/aspec/ray/ray_actor.py +++ b/aspec/ray/ray_actor.py @@ -56,6 +56,12 @@ def node_affinity_for_ip(ip: str, name: str = None) -> NodeAffinitySchedulingStr class RayActor: """Base class for all aspec Ray actors.""" + def get_owner_identity(self): + import psutil + + return {"node_ip": self.get_node_ip(), "pid": os.getpid(), + "created": psutil.Process().create_time()} + @staticmethod def get_node_ip() -> str: """Get current node IP address.""" diff --git a/aspec/train_entry.py b/aspec/train_entry.py index cf0f852..5fdfe80 100644 --- a/aspec/train_entry.py +++ b/aspec/train_entry.py @@ -411,6 +411,9 @@ def train_async_no_generation(args, evidence: RunEvidenceRecorder | None = None) pg=pgs["training"], training_class=TrainerActor, ) + args.aspec_read_owners = ray.get([ + actor.get_owner_identity.remote() for actor in train_group._actor_handlers + ]) train_init_refs = train_group.async_init( args, role="training", mooncake_config=mooncake_config, with_ref=False ) @@ -528,6 +531,13 @@ def main() -> None: def _main() -> None: args, plan = _prepare_config() + from aspec.inference.managed import managed_target + + with managed_target(args): + _run_prepared(args, plan) + + +def _run_prepared(args, plan): from aspec.multi_draft import MultiDraftPlan, run_multi_draft if isinstance(plan, MultiDraftPlan): diff --git a/aspec/training/trainer_actor.py b/aspec/training/trainer_actor.py index 0460822..2587568 100644 --- a/aspec/training/trainer_actor.py +++ b/aspec/training/trainer_actor.py @@ -178,6 +178,13 @@ def get_global_step(self) -> int: def get_status(self) -> dict: """Return the rank placement actually selected by Ray and CUDA.""" + import sys + + import psutil + import torch + + import aspec + from aspec.inference.target_runtime import code_identity, dependency_versions return { "rank": self._rank, @@ -185,9 +192,14 @@ def get_status(self) -> dict: "node_ip": self.get_node_ip(), "node_id": ray.get_runtime_context().get_node_id(), "pid": os.getpid(), + "created": psutil.Process().create_time(), "assigned_gpu_ids": [int(float(value)) for value in ray.get_gpu_ids()], "local_gpu_id": int(os.environ["LOCAL_RANK"]), "initialized": hasattr(self, "_trainer"), + "runtime": {**code_identity(), "python": sys.executable, "prefix": sys.prefix, + "aspec": aspec.__file__, "ray": ray.__file__, + "torch": torch.__version__, "cuda": torch.version.cuda, + "dependencies": dependency_versions("nvidia-cutlass-dsl", "quack-kernels")}, } def save_model(self, step: int, force_sync: bool = False) -> None: diff --git a/aspec/utils/env.py b/aspec/utils/env.py index bfc8835..9e00c0c 100644 --- a/aspec/utils/env.py +++ b/aspec/utils/env.py @@ -9,6 +9,7 @@ # use its own node-local default (/tmp/torchinductor_$USER/) to avoid # cross-node triton kernel cache corruption over NFS. _ASPEC_ENV_KEYS = [ + "ASPEC_CUDA_COMPAT_DIR", "OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS", @@ -83,7 +84,7 @@ def get_aspec_env_vars(args=None) -> dict[str, str]: if getattr(args, "cluster_mode", "single_node") == "static_multi_node": # Ray is started from each node's own venv/environment. A driver's # library/cache paths must not overwrite another node's local layout. - for key in ("LD_LIBRARY_PATH", "PYTHONPATH", "HF_HOME", "CUTE_DSL_CACHE_DIR", + for key in ("LD_LIBRARY_PATH", "PYTHONPATH", "HF_HOME", "ASPEC_CUDA_COMPAT_DIR", "CUTE_DSL_CACHE_DIR", "FLASH_ATTENTION_CUTE_DSL_CACHE_DIR", "ASPEC_LOG_DIR"): env.pop(key, None) return env diff --git a/docs/framework/ARCHITECTURE.md b/docs/framework/ARCHITECTURE.md index 97f57e7..a1c979c 100644 --- a/docs/framework/ARCHITECTURE.md +++ b/docs/framework/ARCHITECTURE.md @@ -22,6 +22,13 @@ flowchart LR The diagram describes responsibilities, not a required class hierarchy or process topology. +Target and Draft start in independent environments. A versioned JSON service interface carries +model/capture identity, token requests, output descriptors and lifecycle calls; feature tensors use +the existing Mooncake data plane. The Draft Plan reserves physical GPUs, and Target validates and +exclusively leases its requested GPU set. Each side owns its runtime and native library paths. +An optional local supervisor starts this same service and retains the full checkpoint/resume path. +Lost control stops admission; source storage remains retained until reader process exit is verified. + ## Configuration, workload and Plan The configuration boundary translates supported native or TorchSpec-style inputs into canonical diff --git a/docs/framework/DECISIONS.md b/docs/framework/DECISIONS.md index f41f89e..62bc057 100644 --- a/docs/framework/DECISIONS.md +++ b/docs/framework/DECISIONS.md @@ -5,16 +5,31 @@ These decisions describe the maintained implementation and its public behavior. - **Product scope.** SGLang with DSpark/DFlash and Qwen3 is the primary multi-GPU performance path. Other implemented algorithms and structures have explicit [support boundaries](../support.md). A source release does not promise all combinations or production fault certification. -- **Naming and licensing.** Project, Python namespace and CLI use ASpec / `aspec`. - ASpec uses Apache-2.0; inherited material retains its own licenses and notices. Source release - does not establish a PyPI distribution or promise name coexistence with unrelated packages. +- **Naming and licensing.** The project uses ASpec; the Python namespace and CLI use `aspec`. + The distribution name on PyPI is `ling-aspec`; the unrelated `aspec` distribution must not be + installed alongside it because their import namespaces can overlap. ASpec uses Apache-2.0; + inherited material retains its own licenses and notices. PyPI publication requires the + [release checks and publishing configuration](../releasing.md). +- **Release versions.** Maintainers choose the major/minor/patch version and currently retain + the `b0` beta suffix for public releases. + CI validates the matching `v` tag without incrementing or promoting the version. + See [versioning](../releasing.md#versioning). - **Reuse.** Review the pinned upstream implementation before adding overlapping functionality. Adapt working code by default; explain a rewrite and verify it against the original behavior. Preserve transitive provenance. See [retention boundaries](RETENTION_BOUNDARY.md). - **Runtime baseline.** Use the [pinned installation lane](../user-guide/installation.md): official SGLang v0.5.21, commit `e00930c5489053f26d86b179cee0d087f846acbb`, with the retained training - patch. FA4 b34, CUTLASS DSL 4.7.1 and Quack 0.6.5 share a compatible dependency lane. Upstream + patch in Target. Draft uses FA4 b34, CUTLASS DSL 4.7.1 and Quack 0.6.5 independently. Upstream migration references do not constrain ASpec to their older runtime versions. +- **Installation.** `ling-aspec` and `ling-aspec-target` install separate Draft/Target dependency + environments, sharing one version of `ling-aspec-runtime`. Target follows upstream SGLang pins; + no metadata overrides are accepted. `aspec patch-sglang` only patches files. `aspec target` and + `aspec train` start in their own environments and connect through a versioned JSON control + protocol; feature tensors retain Mooncake transport. Target's private Ray runtime orchestrates + its source processes; the Draft Plan retains the physical GPU reservations. Optional single-node + managed-local supervision starts the same service, with an explicit Target interpreter only for + that supervisor. Both forms retain exact resume. This boundary is inspired by SpecForge + f8cc3dc external/managed-local capture; its newer SGLang plugin is not copied. - **Semantics and execution.** Workload, capability, Plan and observed Evidence are distinct. Required capability failures are errors; fallback/omission is explicit. No legacy fields leak across the canonical boundary. [Architecture](ARCHITECTURE.md) defines the work/ownership invariants. diff --git a/docs/releasing.md b/docs/releasing.md index aa50eba..cfd861b 100644 --- a/docs/releasing.md +++ b/docs/releasing.md @@ -1,82 +1,65 @@ -# Preparing a release +# Releasing -Wheels contain the runtime, public CLI, model JSON assets, formal recipes, Quickstart conversations, -runtime dependency manifests and the pinned SGLang patch. Source distributions also contain full -documentation, tests, setup and maintenance tools. Canonical resource files live in their existing -source locations; `setup.py` collects them into the wheel during the build. +ASpec publishes the Draft package `ling-aspec`, Target package `ling-aspec-target`, and +shared code package `ling-aspec-runtime` at the same version. Python imports and the CLI use `aspec`. -## Prepare +## Versioning -1. Select a clean, reviewed commit and set its version in `pyproject.toml`. -2. Update the changelog, user guides, supported paths and relevant examples. -3. Check source attribution, license texts and security-reporting instructions. Model and dataset - redistribution terms are separate from the source license. -4. Before publishing a new repository, check both files and Git history for private material. - A reviewed source export need not include development history. +Maintainers set `project.version` in [pyproject.toml](../pyproject.toml). +Current prereleases use `..b0`; maintainers choose the base version. +Tags must match `v` exactly. CI validates the version without changing it. +See the [Python version specification](https://packaging.python.org/en/latest/specifications/version-specifiers/). -## Verify +## Release checks -Use a dedicated Linux venv and run the checks in [CONTRIBUTING.md](../CONTRIBUTING.md). -Build into an empty output directory: +Run the checks in [CONTRIBUTING.md](../CONTRIBUTING.md) and the relevant +[validation gates](contributing/validation.md) for the changes being released. +Build from a clean, reviewed commit: ```bash -python -m tools.build_release --outdir /absolute/path/to/new-release-dist +python -m tools.build_release --outdir dist ``` -This command builds the sdist and then its wheel, audits both, collects the two runtime requirement -files under `runtime/`, and writes `SHA256SUMS`. It rejects a nonempty output directory. -`python -m build` and `python -m tools.audit_release DIST` remain available separately. +The output directory must be empty. The command builds and audits the sdist and wheel, +includes the runtime requirement files, and records their checksums in `SHA256SUMS`. +Verify source attribution and license texts, and exclude credentials, datasets, weights, +checkpoints and local environment files. -[GitHub CI](../.github/workflows/ci.yml) runs source quality checks and CPU/offline contracts on pull -requests and configured branch pushes. It does not build distributions. Packaging-related source -contracts remain in that suite so resource and audit regressions can be found before tagging. +The [release workflow](../.github/workflows/build-release.yml) checks package metadata, +builds the distributions, and validates the installed wheel outside the source checkout. +Manual runs perform these checks without publishing. CPU installation checks do not replace +the GPU validation required for runtime changes. -The separate [release build workflow](../.github/workflows/build-release.yml) runs when a `v*` tag is -pushed to GitHub. Creating a local tag alone does not trigger it. The tag must exactly match `v` plus -the version in `pyproject.toml`, for example `v0.1.0rc1` for `0.1.0rc1`; a mismatch fails before -building. The workflow runs the same build command, saves the audited release bundle as the -`aspec-distributions` workflow artifact, and runs a dependent `wheel-install` job on Python 3.11. -That job checks SHA256s, installs the built wheel in a dedicated CPU venv, and executes: +Wheels include the CLI, model configurations, recipes, Quickstart data, runtime requirement +files and the SGLang patch. Source distributions also include documentation, tests and +maintenance tools. `setup.py` collects resources from their canonical source locations. -```bash -python tools/check_wheel_install.py -``` +## Publishing configuration -The check changes to a temporary working directory and verifies installed module origins, CLI help, -recipe copying and Plan generation, both Quickstart Plans, read-only doctor, package/build identities, -and actual SGLang patch application, repetition and conflict rejection. It uses the CPU CI dependency -subset and the official SGLang wheel without its GPU dependencies. Missing full-runtime requirements -must keep doctor non-ready; this fixture is not production dependency or GPU/training acceptance. - -- Include runtime/model assets, patches, requirements, guides, recipes, source tools, source - headers and applicable license texts. -- Exclude credentials, private endpoints, local environments, datasets, weights, checkpoints, - profiler traces and generated logs. Build from the reviewed tree, not a dirty working directory. -- Unpack the sdist into a fresh directory and check documentation and tool entry points there. -- Install the wheel in a fresh venv and run `aspec patch-sglang` from outside the source tree - against the pinned SGLang installation. Verify application, repeat application and read-only - `--check`; wrong versions or conflicting hunks must fail without overwriting files. -- In a fresh environment, follow installation and Quickstart through training, recovery and export. - For runtime changes, run the relevant [GPU, numerical and fault checks](contributing/validation.md). - A documentation-only change does not require a new all-algorithm training matrix. - -Benchmark pages should contain results, methods, relevant limitations and reproduction entry points. -Historical raw records and per-artifact hash inventories are not required release assets. -Evidence records the effective installed package hash (including the bundled patch and assets), -distribution version and build provenance when available. A wheel does not borrow Git identity from -an enclosing checkout; unavailable build revision remains unknown. +In the release repository, configure a GitHub environment named `pypi`: -## Publish +- Add the environment secret `PYPI_API_TOKEN` with a PyPI token authorized for all three distributions. + Creating the project on its first upload requires an account-wide token; afterward, use a + project-scoped token. +- Allow deployments from `v*` tags. +- Add required reviewers if releases need approval before upload. -Before pushing a version tag, verify that its exact commit has passed the CI source checks and the -applicable release checks above. The tag workflow does not rerun the source test suite. After pushing, -require both `build` and `wheel-install` to pass; an uploaded artifact alone is not installation -acceptance. A failed run can be rerun for the same tag after an infrastructure issue is resolved; -source fixes require a new reviewed commit and version/tag. +See [PyPI token setup](https://packaging.python.org/en/latest/guides/distributing-packages-using-setuptools/#create-an-account) +and [GitHub environments](https://docs.github.com/en/actions/how-tos/deploy/configure-and-manage-deployments/manage-environments). -The release build only saves Actions artifacts. It does not create a GitHub Release or publish to -PyPI; publication is a separate, explicitly authorized step. +## Publish -Review the exact artifacts, repository destination, version/tag and release notes. Summarize checks -actually performed and any known limitations. Publish only the reviewed source and artifacts; -retain the corresponding build and validation records separately. +1. Merge the release changes, including the selected version and changelog, into `main`. + Require source CI and the applicable release checks to pass for that commit. +2. Create and push the matching `v` tag in `inclusionAI/ASpec`. + Mark a corresponding GitHub Release as a prerelease when the package version is a prerelease. +3. In **Actions → Build and publish release**, require `build` and `wheel-install` to pass. + Review the `aspec-distributions` artifact and approve the `pypi` deployment if approval is configured. +4. The `publish` job verifies checksums and uploads the same wheel and sdist to PyPI. + It runs only for version-tag pushes in `inclusionAI/ASpec`. +5. Check the [PyPI release](https://pypi.org/project/ling-aspec/) and follow the + [installation guide](user-guide/installation.md) in a fresh Linux environment. + +PyPI does not allow replacing an uploaded distribution file. Inspect an interrupted upload +before retrying; source changes require a new version. Retain the build and validation records. +The workflow does not create a GitHub Release. diff --git a/docs/user-guide/containers.md b/docs/user-guide/containers.md index 1e012d2..36a6a11 100644 --- a/docs/user-guide/containers.md +++ b/docs/user-guide/containers.md @@ -27,7 +27,10 @@ CUDA_VISIBLE_DEVICES=0,1 aspec quickstart ling-tiny --output /outputs/ling-tiny- # Or use qwen3-8b with a different fresh output. ``` -The image activates its venv through `PATH`/`VIRTUAL_ENV`; no additional setup is needed inside it. +The image activates Draft through `PATH`/`VIRTUAL_ENV`. In another container terminal, activate +`/opt/aspec-target/bin/activate` and run `aspec target` before training. Both environments are +already installed. Optional [local supervision](installation.md#optional-local-supervision) can +use `/opt/aspec-target/bin/python` instead of starting Target manually. The selected GPUs are numbered within the container. Mount an existing model explicitly and pass `--target /container/path/to/model` if downloading is undesirable. The image preserves NVIDIA's entrypoint and does not install a host driver, start a cluster, or assume SSH/shared storage. diff --git a/docs/user-guide/data.md b/docs/user-guide/data.md index 9ba8b00..dc6f96a 100644 --- a/docs/user-guide/data.md +++ b/docs/user-guide/data.md @@ -3,6 +3,9 @@ Target generation, training rows and decode prefixes have different boundaries. Bundled synthetic data is a smoke fixture, not a quality recipe. +Run `aspec prepare-target-data` in the Target environment. +Other local preprocessing commands run in Draft. + ## Conversations One JSON object per JSONL line: diff --git a/docs/user-guide/evaluation.md b/docs/user-guide/evaluation.md index 0a0ac49..8053853 100644 --- a/docs/user-guide/evaluation.md +++ b/docs/user-guide/evaluation.md @@ -2,6 +2,9 @@ Training health, held-out learning, asset compatibility and decode usefulness are separate questions. +Run serving-based `aspec evaluate` in the Target environment. +Checkpoint export runs in Draft. Held-out training evaluation uses the configured Target workers. + ## Held-out evaluation Set `dataset.eval_data_path` and `dataset.eval_interval`. Target observations are cached with diff --git a/docs/user-guide/installation-advanced.md b/docs/user-guide/installation-advanced.md index d5ad5ae..7c877ea 100644 --- a/docs/user-guide/installation-advanced.md +++ b/docs/user-guide/installation-advanced.md @@ -1,29 +1,43 @@ # Advanced installation and native prerequisites -Start with [one-command source setup](installation.md). This page explains its pinned lane, -native prerequisites, compatibility exceptions and manual recovery procedures. +Start with the [installation guide](installation.md). This page explains the supported runtime, +native prerequisites and manual recovery procedures. The [release wheel](installation.md#install-a-release-wheel) includes public commands, formal recipes, Quickstart data and the pinned patch. Source setup and developer `python -m tools.*` commands below -run from the repository root. Both installation forms use the same runtime lane and overrides. +run from the repository root. Both installation forms use the same role dependency manifests. + +`pip install ling-aspec` installs the Draft training dependencies. In a separate Target venv, +`pip install ling-aspec-target` installs SGLang and `aspec patch-sglang` applies its training patch. +Both role packages depend on the same version of `ling-aspec-runtime`, which owns the code/CLI. +`aspec doctor` selects the installed role; `--target` explicitly checks the Target setup. +Start `aspec target` in Target and `aspec train --config ...` in Draft. The control service uses +`inference.target_service.address`; metadata and token requests cross this interface, while feature +tensors retain Mooncake transport. Use services on trusted cluster networks. `--gpus 0,1,...` +restricts the service's allowed physical GPU IDs; it does not remap device ordinals. +For static multi-node runs, start a service on each Target node and set +`inference.target_service.addresses` as a map from planned node IP to service URL. ## Requirements - Linux x86-64 with **glibc 2.34 or newer** (required by the pinned SGLang wheel); Python >=3.11. Check the host with `getconf GNU_LIBC_VERSION`. -- A CUDA-13-compatible NVIDIA driver/runtime. This lane pins Torch2.13.0, SGLang0.5.21, - Transformers5.12.1, FA4 4.0.0b34, CUTLASS DSL 4.7.1 and Quack 0.6.5. +- An NVIDIA driver compatible with the CUDA runtime selected by the + [dependency manifest](../../requirements/sglang-v0.5.21.txt). - Git, a C++20 compiler **and standard library supporting `std::bit_cast`**, and sufficient disk/host memory for CUDA wheels, models and checkpoints. The default SGLang JIT requires both compiler and library support; the presence of `g++` alone is insufficient. - The quickstart requires **two NVLink-connected GPUs**; larger recipes differ. -Direct pins are in [requirements/sglang-v0.5.21.txt](../../requirements/sglang-v0.5.21.txt). They -start from the [official pinned SGLang manifest](https://github.com/sgl-project/sglang/blob/e00930c5489053f26d86b179cee0d087f846acbb/python/pyproject.toml) -and apply the ASpec-validated CUTLASS DSL 4.7.1 / Quack 0.6.5 compatibility override. +[Draft pins](../../requirements/draft.txt) retain CUTLASS DSL 4.7.1 / Quack 0.6.5. +[Target pins](../../requirements/sglang-v0.5.21.txt) use the +[official SGLang requirements](https://github.com/sgl-project/sglang/blob/e00930c5489053f26d86b179cee0d087f846acbb/python/pyproject.toml), +including CUTLASS DSL 4.6.2 / Quack 0.6.4, in their own environment. No dependency overrides are used. +The default Linux training stack uses CUDA 13. Older supported data-center drivers may need the +private compatibility libraries described below. CPU planning does not imply CPU training support. On Python 3.13, `outlines-core==0.1.26`, inherited through the pinned SGLang/Outlines dependency, may require a source build. Install a Rust toolchain (`rustc` and `cargo`) on `PATH` when no -matching wheel is available; uv builds the extension during setup. Setup does not install host -toolchains. The runtime lane pins `nvidia-cudnn-frontend==1.31.0` to retain CUTLASS DSL 4.7.1 +matching wheel is available; pip builds the extension during setup. Setup does not install host +toolchains. The Draft role pins `nvidia-cudnn-frontend==1.31.0` to retain CUTLASS DSL 4.7.1 compatibility; frontend 1.32 requires CUTLASS >=4.8. The C++20 requirement comes from the pinned SGLang JIT: its @@ -40,8 +54,9 @@ printf '#include \nint main() { return std::bit_cast(1.0f) == 0; ``` SGLang also invokes the CUDA toolkit's `nvcc` for JIT kernels. Setup checks the executable selected -by `CUDA_HOME`/`CUDA_PATH`, or `PATH` when those variables are unset. A CUDA-enabled PyTorch wheel -alone does not establish that this compiler is discoverable. The optional container supplies a CUDA +by `CUDA_HOME`/`CUDA_PATH`, or `PATH` when those variables are unset. The Target package installs its matching CUDA compiler; the CLI exposes a private development +toolkit view for JITs that expect `lib64` and unversioned linker names. An explicit host toolkit +selection is preserved. A CUDA-enabled PyTorch wheel alone does not establish compiler availability. The optional container supplies a CUDA 13.0.2 development toolkit. Setup preserves an explicit host compiler/toolkit selection. Use `python -m tools.doctor --gpu` below to check the native Mooncake import in the selected venv. A missing RDMA library should be resolved in the environment before training. @@ -55,29 +70,10 @@ bash tools/install.sh # Contributors use: bash tools/install.sh --dev ``` -The installer checks venv isolation, Linux/Python/glibc and C++20 `std::bit_cast`, installs the -pinned requirements with their explicit uv overrides, applies/verifies the SGLang patch, then runs -CPU doctor. `--gpu` additionally probes CUDA and Mooncake; select an idle GPU first. It does not -create a venv or modify system packages/drivers. For each later command use -`bash tools/run_in_env.sh COMMAND ...`; the example launcher uses this wrapper automatically. -It discovers native libraries inside the active venv and rejects `PYTHONPATH`/`PYTHONHOME` injection. - -The underlying install command, for inspecting or managing the same environment manually, is: - -```bash -python -m uv pip install --link-mode copy --python "$VIRTUAL_ENV/bin/python" \ - --override requirements/cute-overrides.txt -r requirements/sglang-v0.5.21.txt -e ".[fa]" -aspec patch-sglang -``` - -The official SGLang 0.5.21 metadata requires CUTLASS DSL 4.6.2 and Quack 0.6.4. ASpec's -shared FA4/Q/K lane uses 4.7.1 and 0.6.5 through the explicit -[override file](../../requirements/cute-overrides.txt). Ordinary pip resolution rejects this -combination. The override leaves upstream metadata intact: `python -m pip check` reports these -two mismatches. `python -m tools.doctor` discloses the exact exceptions and fails for any other -active missing/conflicting dependency. It does not label the original metadata satisfied. -`cuda-tile==1.6.0rc5`, required by SGLang, is pinned explicitly so resolution does not require -allowing prereleases for unrelated dependencies. +The installer checks venv isolation and native prerequisites, installs the role dependencies +with normal pip resolution, and runs doctor. Use `bash setup.sh --target` for the separate Target +venv; that role also applies/verifies the patch. `--gpu` probes CUDA and Mooncake on allocated GPUs. +Installers do not modify host packages or drivers. CUDA dependencies occupy many gigabytes and contain many small files. Installation and cold JIT can be slow on shared filesystems; use local SSD storage for the venv/caches when available. Models and @@ -102,14 +98,13 @@ ASpec ships the DSpark Q/K RMSNorm + indexed-RoPE forward and backward under `aspec.kernels.cutedsl`, without an external cuDNN source checkout or custom build. The implementation supports separate context/draft K planes, asymmetric GQA heads, deterministic FP32 norm-weight reduction, and the same BF16/FP16 operand precision as the eager reference. It requires -`nvidia-cutlass-dsl[cu13]==4.7.1`; the default `fa` extra already provides that pin and retains +`nvidia-cutlass-dsl[cu13]==4.7.1`; the Draft role provides that pin and retains `quack-kernels[cu13]==0.6.5` for the other fused training operations. Install the normal runtime, or only the kernel dependency for a focused environment: ```bash -uv pip install --python "$VIRTUAL_ENV/bin/python" --override requirements/cute-overrides.txt \ - -r requirements/sglang-v0.5.21.txt -e ".[fa,dspark-qk]" +python -m pip install -r requirements/draft.txt -e . python -m tools.doctor ``` @@ -181,7 +176,7 @@ actor processes must inherit the library paths. | Step | Automation boundary | | --- | --- | -| Python dependencies, two declared overrides, SGLang patch | `tools/install.sh`; patch check is idempotent | +| Role dependencies and the SGLang patch | `tools/install.sh`; patch check is idempotent | | Venv-native library paths and optional private compatibility/cache paths | `tools/run_in_env.sh`; example launcher delegates to it | | Base Python/venv, C++ compiler, glibc, RDMA host libraries and driver | Provision the host first; checks do not install system packages | | Private CUDA compatibility package | Obtain/extract an appropriate NVIDIA package explicitly, then set the directory above | diff --git a/docs/user-guide/installation.md b/docs/user-guide/installation.md index 7864b6f..7ca9e59 100644 --- a/docs/user-guide/installation.md +++ b/docs/user-guide/installation.md @@ -1,109 +1,113 @@ # Installation -Use Linux x86-64 with Python >=3.11, Git and a CUDA development -toolkit containing `nvcc`. Training requires NVIDIA GPUs and a CUDA-13-compatible driver; -Quickstart uses two NVLink-connected GPUs. -Source setup is the easiest complete installation. Release bundles also provide a wheel with -the public CLI, formal recipes and Quickstart data; developer validation tools remain in the source. +Use Linux x86-64, Python >=3.11 and an isolated virtual environment. Training requires NVIDIA +GPUs and a CUDA 13-compatible driver. The Target package supplies its matching CUDA compiler; +the host still needs the C++20 compiler and native libraries listed below. +See [native prerequisites](installation-advanced.md#requirements) for the compiler, glibc and RDMA +library requirements, or use the [container](containers.md). -On Python 3.13, the pinned grammar dependency may build from source and needs Rust/Cargo on -`PATH` when no matching wheel is available; see [native prerequisites](installation-advanced.md#requirements). - -## One-command setup +## Install from PyPI ```bash -bash setup.sh +python -m venv .venv source .venv/bin/activate +pip install ling-aspec ``` -Setup creates an isolated `.venv`, installs the pinned runtime, applies the SGLang training patch, -and then runs CPU doctor. The ASpec-owned DSpark Q/K CuTeDSL kernel ships in the package. Setup checks the host's glibc, C++ -toolchain, CUDA compiler and RDMA libraries before downloading the runtime and gives repair guidance when a native -prerequisite is missing. It does not change system packages/drivers, download models or start Ray. +This installs ASpec and its Draft training dependencies. The Python package and command are named +`aspec`. Keep the unrelated PyPI package `aspec` out of this environment. -Choose a different environment location or install contributor tools when needed: +## Target nodes + +In a separate virtual environment on nodes that run the Target model: ```bash -bash setup.sh --venv /local/aspec-venv --python python3.12 --dev -source /local/aspec-venv/bin/activate +python -m venv .venv-target +source .venv-target/bin/activate +pip install ling-aspec-target +aspec patch-sglang +aspec target ``` -Rerunning setup is supported for an environment created by the same checkout. Existing unowned venvs -are rejected. The Q/K kernel is selected only for compatible CUDA shapes and otherwise preserves the -eager path. `--gpu` additionally runs GPU doctor; -select an allocated idle device with `CUDA_VISIBLE_DEVICES` before using that option. +Target dependencies follow SGLang's requirements; Draft acceleration packages stay in the Draft +venv. The patch command only changes the supported SGLang source files, never dependencies. -`aspec` commands prepare this venv's native-library paths automatically. The -`tools/run_in_env.sh` wrapper remains available for arbitrary Python/torchrun commands. +Keep this service running. In the Draft environment, start training: -## First run +```bash +aspec train --config /path/to/run.yaml +``` + +The default Target address is `http://127.0.0.1:30000`. Set `inference.target_service.address` +for a remote service, or pass `--target-address` to Quickstart. Use the same ASpec release +and Python minor version in both environments. No interpreter path is exchanged. +`aspec target --config /path/to/run.yaml` reads its address and Target model from the same config. +The service remains available after a training session ends; each session releases its own workers. -Choose either model; both use bundled conversations and the same two-GPU train/resume/export flow: +Recheck the patch after reinstalling SGLang: ```bash -CUDA_VISIBLE_DEVICES=0,1 aspec quickstart ling-tiny --output outputs/ling-tiny-001 -# Or: -CUDA_VISIBLE_DEVICES=0,1 aspec quickstart qwen3-8b --output outputs/qwen3-8b-001 +aspec patch-sglang --check ``` -Use `--target /absolute/path/to/model` to reuse a compatible local checkpoint. Without it, the selected -model's public Hub ID is used. Continue with [Quickstart](quickstart.md) for Plan, results and assets. +Applying the patch again verifies an already patched installation. Conflicting changes or an +unsupported SGLang version are rejected. `--check` never installs dependencies or modifies SGLang files. -## Other environments +## Check the environment -- [Container setup](containers.md) supplies a CUDA development base, Python, C++ and RDMA libraries; - it runs the same source setup and still uses the host's NVIDIA driver. -- [Advanced installation](installation-advanced.md) covers native prerequisites, exact version pins, - dependency override disclosures, private CUDA compatibility libraries, caches and manual setup. -- [Troubleshooting](troubleshooting.md) explains runtime/transport failures. +Optionally run `aspec doctor` to check the training environment, or `aspec doctor --target` to +also require the Target runtime and patch. Add `--gpu` to check CUDA and native libraries on an +allocated GPU. These checks do not run training. -The pinned runtime is Torch 2.13.0 / SGLang 0.5.21 / CUDA 13. Doctor reports the -[two dependency metadata overrides](installation-advanced.md#dedicated-environment) explicitly. -CPU checks do not exercise GPU capacity or NVLink. - -## Install a release wheel +## First run -Use the wheel and `runtime/` requirements from the same audited release bundle. Prepare the -[native prerequisites](installation-advanced.md#requirements), then install into a fresh venv: +Quickstart requires two NVLink-connected GPUs: ```bash -python3.12 -m venv /local/aspec-venv -source /local/aspec-venv/bin/activate -python -m pip install uv==0.13.5 -python -m uv pip install --link-mode copy --python "$VIRTUAL_ENV/bin/python" \ - --override /path/to/release/runtime/cute-overrides.txt \ - -r /path/to/release/runtime/sglang-v0.5.21.txt \ - '/path/to/release/aspec-0.1.0rc1-py3-none-any.whl[fa]' +CUDA_VISIBLE_DEVICES=0,1 aspec quickstart ling-tiny --output outputs/ling-tiny-001 +# Or: +CUDA_VISIBLE_DEVICES=0,1 aspec quickstart qwen3-8b --output outputs/qwen3-8b-001 ``` -The versioned runtime requirements and explicit overrides are also bundled as package resources -for doctor. The wheel does not include Torch/SGLang/CUDA binaries or model weights. -Apply and verify the pinned SGLang training patch explicitly: +Add `--target /absolute/path/to/model` to use a local checkpoint. +Continue with [Quickstart](quickstart.md) or choose a [training recipe](recipes.md). + +## Install from source + +From a source checkout: ```bash -aspec patch-sglang -aspec patch-sglang --check -aspec doctor +bash setup.sh +source .venv/bin/activate ``` -Application is explicit: installing the wheel does not modify SGLang. Repeating the first command -verifies an already applied patch; `--check` never writes and exits nonzero if the complete patch -is absent. Wrong versions, conflicting changes and writes outside the active venv are rejected. -Recheck after reinstalling SGLang. Git is required for exact patch application and verification. +Setup installs the Draft environment. Run `bash setup.sh --target` to prepare `.venv-target`, +then activate it and run `aspec target`. Each setup runs the appropriate doctor. Contributors can use +`bash setup.sh --dev`. Run `bash setup.sh --help` for environment and interpreter options. + +## Install a release wheel -The installed CLI supports Quickstart, recipe listing/copying/planning/training, checkpoint recovery, -Report, export, Native MTP stitching, data preparation and decode evaluation outside the source tree. -For example: +Use `pip install --find-links /path/to/release ling-aspec` in Draft, and +`pip install --find-links /path/to/release ling-aspec-target` in Target. Each role installs +the same `ling-aspec-runtime` wheel, which owns the `aspec` code and CLI. Then prepare Target nodes and optionally check the environment as above. The installed commands, +recipes and Quickstart data work outside the source checkout. -```bash -aspec recipes list -aspec recipes copy dspark --output my-dspark-config -aspec quickstart ling-tiny --plan --output outputs/ling-plan -aspec export --help -aspec evaluate --help +See [advanced installation](installation-advanced.md) for native-library paths and manual +dependency management, or [troubleshooting](troubleshooting.md) if a check fails. + +## Optional local supervision + +To start and stop Target automatically for a single-node run, add: + +```yaml +inference: + target_service: + mode: managed-local + python_executable: /absolute/path/to/target-venv/bin/python ``` -Copied YAMLs refer to input assets in this installation. You can edit or replace those paths; -run/cache outputs remain relative to your working directory or explicit overrides. Use a new -copy/output directory to preserve existing work. Full docs, benchmark comparisons, fault injection, -tests and source maintenance commands remain in the source distribution. +Then run the same `aspec train --config ...` command. This optional supervisor needs the +Target environment location; independent startup only needs a service address. Both forms use +the same Target interface and support checkpoint restoration. Logs go to `target-service.log` +inside the run directory. Existing services are never adopted or stopped by the supervisor. +Quickstart also accepts `--managed-target /absolute/path/to/target-venv/bin/python`. diff --git a/pyproject.toml b/pyproject.toml index 4787caa..1b9f913 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,39 +3,28 @@ requires = ["setuptools>=77.0", "wheel"] build-backend = "setuptools.build_meta" [project] -name = "aspec" -version = "0.1.0rc1" +name = "ling-aspec-runtime" +version = "0.1.0b0" description = "A Torch-native framework for speculative decoding training" readme = "README.md" license = "Apache-2.0" license-files = ["LICENSE", "licenses/*"] requires-python = ">=3.11" dependencies = [ - "accelerate", - "cmake", - "datasets>=5.0.1,<6", - "mooncake-transfer-engine>=0.3.10.post1", - "ninja", - "numba", - "numpy<2.4", - "omegaconf", - "openai", - "openai-harmony", - "packaging", - "psutil", - "pydantic", - "pyzmq", - "qwen-vl-utils", - "ray", - "setuptools", - "tqdm", - "transformers", - "wandb>=0.19.10", + "accelerate", "datasets>=5.0.1,<6", "numpy<2.4", "omegaconf", "openai", + "openai-harmony", "packaging", "psutil", "pydantic", "pyzmq", "qwen-vl-utils", + "ray==2.59.0", "setuptools", "tqdm", "transformers", "wandb>=0.19.10", + "mooncake-transfer-engine==0.3.13.post1", "nvidia-cuda-runtime-cu12==12.9.79", ] [project.scripts] aspec = "aspec.cli:main" +[project.urls] +Homepage = "https://github.com/inclusionAI/ASpec" +Documentation = "https://github.com/inclusionAI/ASpec/blob/main/docs/README.md" +Issues = "https://github.com/inclusionAI/ASpec/issues" + [project.optional-dependencies] dev = ["build", "pre-commit", "pytest", "ruff", "tensorboard"] tracking = ["tensorboard"] diff --git a/requirements/ci-cpu.txt b/requirements/ci-cpu.txt index a68f5c7..58d020b 100644 --- a/requirements/ci-cpu.txt +++ b/requirements/ci-cpu.txt @@ -1,5 +1,5 @@ # CPU-only public CI dependencies. The production CUDA lane remains pinned in -# requirements/sglang-v0.5.21.txt; this smaller set exercises planning, data, +# requirements/draft.txt and requirements/sglang-v0.5.21.txt; this smaller set exercises planning, data, # controller and configuration contracts without installing SGLang/CUDA wheels. datasets>=5.0.1,<6 numba diff --git a/requirements/cute-overrides.txt b/requirements/cute-overrides.txt deleted file mode 100644 index d09d199..0000000 --- a/requirements/cute-overrides.txt +++ /dev/null @@ -1,4 +0,0 @@ -# Intentional, GPU-validated overrides of SGLang 0.5.21 metadata. -# The environment doctor reports these exact exceptions and rejects other conflicts. -nvidia-cutlass-dsl[cu13]==4.7.1 -quack-kernels[cu13]==0.6.5 diff --git a/requirements/draft.txt b/requirements/draft.txt new file mode 100644 index 0000000..380a850 --- /dev/null +++ b/requirements/draft.txt @@ -0,0 +1,12 @@ +# Draft-only CUDA 13 training dependencies. SGLang owns the Target environment. +torch==2.13.0 +torchvision==0.28.0 +transformers==5.12.1 +apache-tvm-ffi==0.1.11 +flash-attn-4==4.0.0b34 +nvidia-cutlass-dsl[cu13]==4.7.1 +quack-kernels[cu13]==0.6.5 +nvidia-cudnn-frontend==1.31.0 +cmake +ninja +numba diff --git a/requirements/sglang-v0.5.21.txt b/requirements/sglang-v0.5.21.txt index fd0e9c8..eb37893 100644 --- a/requirements/sglang-v0.5.21.txt +++ b/requirements/sglang-v0.5.21.txt @@ -1,23 +1,7 @@ -# Official v0.5.21 CUDA 13 dependency lane with the ASpec-validated -# CUTLASS DSL 4.7.1 / Quack 0.6.5 compatibility override. +# Target-only environment: retain SGLang's declared dependency versions. sglang==0.5.21 -torch==2.13.0 -transformers==5.12.1 -sglang-kernel==0.4.7 -flashinfer-python[cu13]==0.6.18 -apache-tvm-ffi==0.1.11 flash-attn-4==4.0.0b34 -nvidia-cutlass-dsl[cu13]==4.7.1 -quack-kernels[cu13]==0.6.5 -# Frontend 1.32 requires CUTLASS >=4.8; keep the validated 4.7.1 lane. -nvidia-cudnn-frontend==1.31.0 +cuda-toolkit[nvcc]==13.0.3.0 +# Mooncake's native wheel needs libcudart.so.12 alongside Torch's CUDA 13. mooncake-transfer-engine==0.3.13.post1 -ray==2.59.0 -# The retained Mooncake wheel links libcudart.so.12, independently of Torch CUDA 13. nvidia-cuda-runtime-cu12==12.9.79 - -# Explicit prerelease required by the official SGLang runtime. -cuda-tile==1.6.0rc5 - -# Match the Torch 2.13 runtime used by this lane. -torchvision==0.28.0 diff --git a/setup.py b/setup.py index 4fd87aa..62324ac 100644 --- a/setup.py +++ b/setup.py @@ -49,7 +49,7 @@ def run(self): ROOT / "examples/data/sample_conversations.jsonl", ROOT / "tools/recipe_catalog.json", ROOT / "requirements/sglang-v0.5.21.txt", - ROOT / "requirements/cute-overrides.txt"] + ROOT / "requirements/draft.txt"] for source in assets: destination = package / "resources" / source.relative_to(ROOT) destination.parent.mkdir(parents=True, exist_ok=True) diff --git a/setup.sh b/setup.sh index a611266..67d7d4e 100755 --- a/setup.sh +++ b/setup.sh @@ -4,15 +4,18 @@ repo_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" venv="$repo_root/.venv" base_python="" options=() +role=draft +explicit_venv=0 while [[ $# -gt 0 ]]; do case "$1" in --venv|--python) [[ $# -ge 2 && -n $2 && $2 != --* ]] || { echo "$1 requires a value" >&2; exit 2; } - if [[ $1 == --venv ]]; then venv=$2; else base_python=$2; fi + if [[ $1 == --venv ]]; then venv=$2; explicit_venv=1; else base_python=$2; fi shift 2 ;; + --target) role=target; options+=("$1"); shift ;; --dev|--gpu) options+=("$1"); shift ;; --help) - echo 'Usage: bash setup.sh [--venv PATH] [--python EXECUTABLE] [--dev] [--gpu]' + echo 'Usage: bash setup.sh [--venv PATH] [--python EXECUTABLE] [--target] [--dev] [--gpu]' echo 'Creates an isolated Linux Python >=3.11 venv and installs pinned dependencies.' echo 'Defaults to .venv. Existing venvs require this checkout owner.' echo 'Does not install host packages/drivers, download models, start training or clean clusters.' @@ -20,6 +23,7 @@ while [[ $# -gt 0 ]]; do *) echo "Unknown argument: $1" >&2; exit 2 ;; esac done +if [[ $role == target && $explicit_venv == 0 ]]; then venv="$repo_root/.venv-target"; fi [[ $(uname -s) == Linux && $(uname -m) == x86_64 ]] || { echo 'Use Linux x86-64; see docs/user-guide/installation.md.' >&2; exit 2; } @@ -51,4 +55,7 @@ fi source "$venv/bin/activate" cd "$repo_root" bash tools/install.sh "${options[@]}" -printf '\nSetup complete. Next:\n source %q\n aspec doctor --gpu\n CUDA_VISIBLE_DEVICES=0,1 aspec quickstart ling-tiny --output outputs/ling-tiny-001\n' "$venv/bin/activate" +printf '\nSetup complete. Activate with:\n source %q\n' "$venv/bin/activate" +if [[ $role == draft ]]; then + printf 'Prepare Target separately: bash setup.sh --target\nThen set ASPEC_TARGET_PYTHON to its bin/python before training.\n' +fi diff --git a/tests/test_cluster_preflight.py b/tests/test_cluster_preflight.py index e44a848..2d9fb64 100644 --- a/tests/test_cluster_preflight.py +++ b/tests/test_cluster_preflight.py @@ -123,6 +123,23 @@ def test_static_cluster_inventory_accepts_homogeneous_shared_nodes(): _validate([_inventory("10.0.0.1"), _inventory("10.0.0.2")]) +def test_separate_runtime_roles_do_not_require_sglang_on_draft_nodes(): + draft, target = _inventory("10.0.0.1"), _inventory("10.0.0.2") + draft["roles"], target["roles"] = ["draft"], ["target"] + draft["runtime_artifacts"]["sglang"] = {"path": None} + draft["packages"].update({"sglang": None, "nvidia-cutlass-dsl": "4.7.1"}) + target["packages"].update({"sglang": "0.5.21", "nvidia-cutlass-dsl": "4.6.2"}) + target["target_runtime"] = {"version": "0.1.0b0", "code_sha256": "same"} + _validate([draft, target]) + + +def test_split_target_node_cannot_omit_its_environment_observation(): + draft, target = _inventory("10.0.0.1"), _inventory("10.0.0.2") + draft["roles"], target["roles"] = ["draft"], ["target"] + with pytest.raises(ClusterPreflightError, match="isolated runtime"): + _validate([draft, target]) + + def test_locality_domain_geometry_resolves_live_node_identity(): args = SimpleNamespace( training_node_ips=["10.0.0.1", "10.0.0.2"], diff --git a/tests/test_convert_to_hf.py b/tests/test_convert_to_hf.py index e1afcea..0823a5d 100644 --- a/tests/test_convert_to_hf.py +++ b/tests/test_convert_to_hf.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import importlib.metadata import json import pytest @@ -53,6 +54,10 @@ def _model(): def test_convert_checkpoint_writes_strict_loadable_serving_artifact(tmp_path, monkeypatch): + versions = {"ling-aspec-runtime": "0.1.0b0", "aspec": "0.0.16"} + original_version = importlib.metadata.version + monkeypatch.setattr(importlib.metadata, "version", + lambda name: versions[name] if name in versions else original_version(name)) torch.manual_seed(7) source = _model() checkpoint = tmp_path / "iter_0000002" / "model" @@ -79,6 +84,7 @@ def test_convert_checkpoint_writes_strict_loadable_serving_artifact(tmp_path, mo for key, value in expected.items(): torch.testing.assert_close(exported[key], value) exported_config = json.loads((output / "config.json").read_text()) + assert exported_config["_aspec_version"] == "0.1.0b0" assert exported_config["model_type"] == "qwen3" assert exported_config["num_target_layers"] == 2 assert exported_config["dflash_config"] == { diff --git a/tests/test_dependency_audit.py b/tests/test_dependency_audit.py index 6080cbe..5736735 100644 --- a/tests/test_dependency_audit.py +++ b/tests/test_dependency_audit.py @@ -23,14 +23,13 @@ def distribution(name, version, requires=()): return SimpleNamespace(metadata={"Name": name}, version=version, requires=requires) -def test_declared_overrides_remain_visible(): +def test_previous_overrides_are_now_rejected(): result = check_dependency_metadata([ distribution("sglang", "0.5.21", ["nvidia-cutlass-dsl[cu13]==4.6.2", "quack-kernels==0.6.4"]), distribution("nvidia-cutlass-dsl", "4.7.1"), distribution("quack-kernels", "0.6.5"), ]) - assert not result["conflicts"] - assert len(result["overrides"]) == 2 - assert {row["installed"] for row in result["overrides"]} == {"4.7.1", "0.6.5"} + assert len(result["conflicts"]) == 2 + assert not result["overrides"] @pytest.mark.parametrize("owner,version,requires,installed", [ @@ -39,7 +38,7 @@ def test_declared_overrides_remain_visible(): ("sglang", "0.5.21", "quack-kernels==0.6.4", "0.6.6"), ("another-package", "0.5.21", "quack-kernels==0.6.4", "0.6.5"), ]) -def test_other_conflicts_cannot_borrow_the_override(owner, version, requires, installed): +def test_all_dependency_conflicts_are_rejected(owner, version, requires, installed): result = check_dependency_metadata([ distribution(owner, version, [requires]), distribution("quack-kernels", installed), ]) @@ -55,13 +54,32 @@ def test_missing_dependency_fails_and_inactive_extra_is_ignored(): assert "missing" in result["conflicts"][0] +@pytest.mark.parametrize("target", [False, True]) +def test_doctor_requires_sglang_only_for_target_nodes(monkeypatch, target): + from tools import doctor + + monkeypatch.setattr(doctor, "check_dependency_metadata", lambda: {"conflicts": [], "overrides": []}) + monkeypatch.setattr(doctor, "venv_error", lambda: None) + monkeypatch.setattr(doctor, "expected_versions", lambda target=False: {"sglang": "0.5.21"}) + + def missing(_=None): + raise doctor.importlib.metadata.PackageNotFoundError("sglang") + + monkeypatch.setattr(doctor.importlib.metadata, "version", missing) + monkeypatch.setattr(doctor, "installed_site", missing) + result = doctor.inspect_environment(target=target) + assert result["ok"] == (not target) + assert result["patch"] == ("unavailable" if target else "not_required") + assert result["target_requested"] == target + + @pytest.mark.parametrize("version", [(3, 10), (3, 11), (3, 12), (3, 13), (3, 14)]) def test_doctor_enforces_only_the_python_minimum(monkeypatch, tmp_path, version): from tools import doctor monkeypatch.setattr(doctor, "check_dependency_metadata", lambda: {"conflicts": [], "overrides": []}) monkeypatch.setattr(doctor, "venv_error", lambda: None) - monkeypatch.setattr(doctor, "expected_versions", lambda: {}) + monkeypatch.setattr(doctor, "expected_versions", lambda target=False: {}) monkeypatch.setattr(doctor, "installed_site", lambda: tmp_path) monkeypatch.setattr(doctor, "patch_state", lambda _: "patched") monkeypatch.setattr(doctor.sys, "version_info", (*version, *doctor.sys.version_info[2:])) diff --git a/tests/test_dependency_extras.py b/tests/test_dependency_extras.py index a22928b..7b9e86c 100644 --- a/tests/test_dependency_extras.py +++ b/tests/test_dependency_extras.py @@ -48,15 +48,14 @@ def test_dspark_qk_uses_the_fa_cutlass_runtime_without_requiring_quack() -> None assert any(requirement.startswith("quack-kernels") for requirement in fa_runtime) -def test_sglang_requirements_match_the_shared_cute_runtime() -> None: - project = tomllib.loads((ROOT / "pyproject.toml").read_text()) - expected = _shared_cute_requirements( - project["project"]["optional-dependencies"]["fa"] - ) - requirements = { - line.strip() - for line in (ROOT / "requirements" / "sglang-v0.5.21.txt").read_text().splitlines() - if line.strip() and not line.lstrip().startswith("#") - } - - assert expected <= requirements +def test_draft_and_target_dependencies_are_separate(): + project = tomllib.loads((ROOT / "pyproject.toml").read_text())["project"] + assert project["name"] == "ling-aspec-runtime" + core = project["dependencies"] + assert not any(x.startswith(("sglang", "torch==", "nvidia-cutlass", "quack", "flash-attn")) for x in core) + draft = (ROOT / "requirements/draft.txt").read_text() + target = (ROOT / "requirements/sglang-v0.5.21.txt").read_text() + assert "nvidia-cutlass-dsl[cu13]==4.7.1" in draft + assert "sglang==0.5.21" in target + assert "nvidia-cutlass" not in target and "quack-kernels" not in target + assert not (ROOT / "requirements/cute-overrides.txt").exists() diff --git a/tests/test_environment_scripts.py b/tests/test_environment_scripts.py index 1ced720..d87e579 100644 --- a/tests/test_environment_scripts.py +++ b/tests/test_environment_scripts.py @@ -76,6 +76,32 @@ def test_installer_rejects_unknown_options_without_installing(): assert "Unknown argument" in result.stderr +def test_pip_cuda_toolkit_links_like_a_development_install_without_changing_wheels(tmp_path, monkeypatch): + from aspec import _environment + + purelib = tmp_path / "site-packages" + toolkit = purelib / "nvidia/cu13" + (toolkit / "bin").mkdir(parents=True) + (toolkit / "bin/nvcc").write_text("compiler fixture") + (toolkit / "lib").mkdir() + library = toolkit / "lib/libcudart.so.13" + library.write_bytes(b"runtime fixture") + monkeypatch.setattr(_environment.sysconfig, "get_paths", lambda: {"purelib": str(purelib)}) + monkeypatch.setenv("VIRTUAL_ENV", sys.prefix) + monkeypatch.setenv("ASPEC_CACHE_HOME", str(tmp_path / "cache")) + for key in ("CUDA_HOME", "CUDA_PATH", "PYTHONPATH", "PYTHONHOME", "ASPEC_CUDA_COMPAT_DIR"): + monkeypatch.delenv(key, raising=False) + environment = _environment.command_environment() + view = Path(environment["CUDA_HOME"]) + assert (view / "lib64/libcudart.so").resolve() == library + assert (view / "bin/nvcc").resolve() == toolkit / "bin/nvcc" + assert environment["LIBRARY_PATH"].split(":")[0] == str(view / "lib64") + assert not (toolkit / "lib64").exists() + assert not (toolkit / "lib/libcudart.so").exists() + assert library.read_bytes() == b"runtime fixture" + assert _environment.command_environment()["CUDA_HOME"] == str(view) + + @pytest.mark.parametrize("version", [(3, 10), (3, 11), (3, 12), (3, 13), (3, 14)]) def test_environment_wrapper_enforces_only_the_python_minimum(monkeypatch, version): from aspec import _environment diff --git a/tests/test_evidence.py b/tests/test_evidence.py index f213a15..997f9ab 100644 --- a/tests/test_evidence.py +++ b/tests/test_evidence.py @@ -231,6 +231,8 @@ def test_host_staged_transport_does_not_hide_endpoint_without_metadata(tmp_path) def _feature_observation(plan): + policy = plan.workload.models["draft"]["attention"] + parallel = plan.policy.parallelism["training"] return { "state": "validated", "contract_id": plan.workload.features["contract_id"], @@ -239,7 +241,20 @@ def _feature_observation(plan): "producer_metadata_validations": 8, "consumer_metadata_validations": 8, "consumer_tensor_validations": 8, - "participants": [], + "participants": [{ + "role": "consumer", + "execution_metadata": { + "draft_attention_backend": [parallel["attention_backend"]], + "draft_context_parallel_size": [parallel["context_parallel"]], + }, + "draft_attention": { + "head_layout": policy["head_layout"], + "context_windows": policy["context_windows"], + "window_semantics": policy.get("window_semantics", "anchor_context"), + "draft_block_attention": policy.get("draft_block_attention", "bidirectional"), + "forward_observed": True, + }, + }], } @@ -411,6 +426,33 @@ def _inference_engine_statuses(plan): ] +@pytest.mark.parametrize("different", [False, True]) +def test_backend_source_identity_is_reported_by_target_workers(tmp_path, different): + plan, recorder = _recorder(tmp_path) + statuses = _inference_engine_statuses(plan) + backend = {"path": "/target/site-packages/sglang", "git_revision": None, + "git_dirty": None, "effective_tree_sha256": "a" * 64, + "source_file_count": 3, "missing_files": []} + for status in statuses: + status["runtime"] = {"backend_identity": dict(backend), "sglang_version": "0.5.21"} + if different: + statuses[-1]["runtime"]["backend_identity"]["effective_tree_sha256"] = "b" * 64 + arguments = dict( + training_actor_count=plan.policy.resources["training"]["world_size"], + inference_engine_count=len(statuses), inference_engine_statuses=statuses, + dataset_size=128, eval_dataset_size=0, cluster_resources={"CPU": 16, "GPU": 8}, + alive_node_count=1, + ) + if different: + with pytest.raises(RuntimeError, match="different SGLang identities"): + recorder.record_actor_initialization(**arguments) + else: + recorder.record_actor_initialization(**arguments) + assert recorder.data["source"]["backends"]["sglang"] == backend + assert recorder.data["source"]["package_scope"] == "draft_driver" + assert len(recorder.data["source"]["target_runtimes"]) == len(statuses) + + def test_static_multi_node_evidence_validates_actual_actor_placement(tmp_path): plan = _static_multi_node_plan(tmp_path) recorder = _recorder_for_plan(tmp_path, plan) @@ -669,6 +711,10 @@ def test_evidence_separates_initialized_and_end_to_end_capabilities(tmp_path): "final_metrics": {"train/avg_loss": 1.25}, "checkpoint": {"enabled": False, "exists": False}, "feature_contract": _feature_observation(plan), + "execution": { + "attention_backend": plan.policy.parallelism["training"]["attention_backend"], + "context_parallel_size": plan.policy.parallelism["training"]["context_parallel"], + }, } ) recorder.complete() @@ -690,7 +736,6 @@ def test_evidence_separates_initialized_and_end_to_end_capabilities(tmp_path): "gpu_direct_transport", "target_prefill_output", "producer_consumer_pipeline", - "draft_attention_execution", "target_logits", "mixed_length_packing", "long_context_execution", @@ -731,6 +776,10 @@ def test_required_performance_and_numerics_need_runtime_evidence(tmp_path): "optimizer_steps_observed": 3, "checkpoint": {"enabled": False, "exists": False}, "feature_contract": _feature_observation(plan), + "execution": { + "attention_backend": plan.policy.parallelism["training"]["attention_backend"], + "context_parallel_size": plan.policy.parallelism["training"]["context_parallel"], + }, "final_metrics": {"train/avg_loss": 1.25, "train/grad_norm": 2.5}, "performance": { "state": "measured", @@ -884,6 +933,7 @@ def test_optimizer_step_without_contract_observation_is_incomplete(tmp_path): assert recorder.data["required_capabilities_pending"] == [ "target_features", "tensor_transport", + "draft_attention_execution", "work_scheduler", "data_lifecycle", "capacity_control", diff --git a/tests/test_open_source_entrypoints.py b/tests/test_open_source_entrypoints.py index f8d5843..b5d3ba8 100644 --- a/tests/test_open_source_entrypoints.py +++ b/tests/test_open_source_entrypoints.py @@ -67,6 +67,19 @@ def test_unknown_wheel_changes_are_not_overwritten(wheel): assert file.read_text() == "value = 99\n" +@pytest.mark.parametrize("check", [False, True]) +def test_missing_target_runtime_never_triggers_dependency_install(monkeypatch, check): + def missing(**kwargs): + raise importlib.metadata.PackageNotFoundError("sglang") + + monkeypatch.setattr(patcher, "installed_site", missing) + monkeypatch.setattr(patcher.subprocess, "run", lambda *a, **k: pytest.fail("must not install")) + monkeypatch.setattr(sys, "argv", ["patch", *(["--check"] if check else [])]) + with pytest.raises(SystemExit) as error: + patcher.main() + assert error.value.code == 1 + + def test_patch_check_works_inside_an_unrelated_git_checkout(wheel): site, file, patch = wheel subprocess.run(["git", "init", "-q", str(site.parent)], check=True) @@ -162,7 +175,7 @@ def test_doctor_reports_system_site_packages_even_when_versions_match(tmp_path, monkeypatch.setattr(sys, "prefix", str(tmp_path)) (tmp_path / "pyvenv.cfg").write_text("include-system-site-packages = true\n") - monkeypatch.setattr(doctor, "expected_versions", lambda: {}) + monkeypatch.setattr(doctor, "expected_versions", lambda target=False: {}) monkeypatch.setattr(doctor, "installed_site", lambda: tmp_path) monkeypatch.setattr(doctor, "patch_state", lambda _: "patched") result = doctor.inspect_environment() diff --git a/tests/test_release_audit.py b/tests/test_release_audit.py index eb8ab62..0d79272 100644 --- a/tests/test_release_audit.py +++ b/tests/test_release_audit.py @@ -28,48 +28,70 @@ def _artifacts( wheel_extra=(), wheel_omit=(), sdist_extra=(), - wheel_version="0.1.0rc1", - sdist_version="0.1.0rc1", + wheel_version="0.1.0b0", + sdist_version="0.1.0b0", sdist_link=False, + distribution="ling-aspec-runtime", ): - wheel = directory / "aspec-0.1.0rc1-py3-none-any.whl" + archive_name = distribution.replace("-", "_") + wheel = directory / f"{archive_name}-0.1.0b0-py3-none-any.whl" with zipfile.ZipFile(wheel, "w") as output: for name in sorted(WHEEL_REQUIRED - set(wheel_omit)): output.writestr(name, "") output.writestr( - "aspec-0.1.0rc1.dist-info/METADATA", - f"Name: aspec\nVersion: {wheel_version}\n", + f"{archive_name}-0.1.0b0.dist-info/METADATA", + f"Name: {distribution}\nVersion: {wheel_version}\n", ) for name in wheel_extra: output.writestr(name, "") - source = directory / "aspec-0.1.0rc1.tar.gz" + source = directory / f"{archive_name}-0.1.0b0.tar.gz" with tarfile.open(source, "w:gz") as output: for relative in sorted(SDIST_REQUIRED | set(sdist_extra)): data = ( - f"Name: aspec\nVersion: {sdist_version}\n".encode() + f"Name: {distribution}\nVersion: {sdist_version}\n".encode() if relative == "PKG-INFO" else relative.encode() ) - info = tarfile.TarInfo(f"aspec-0.1.0rc1/{relative}") + info = tarfile.TarInfo(f"{archive_name}-0.1.0b0/{relative}") info.size = len(data) output.addfile(info, io.BytesIO(data)) if sdist_link: - link = tarfile.TarInfo("aspec-0.1.0rc1/aspec/link") + link = tarfile.TarInfo(f"{archive_name}-0.1.0b0/aspec/link") link.type = tarfile.SYMTYPE link.linkname = "target" output.addfile(link) + for role in ("ling-aspec", "ling-aspec-target"): + normalized = role.replace("-", "_") + metadata = (f"Name: {role}\nVersion: {wheel_version}\n" + f"Requires-Dist: ling-aspec-runtime=={wheel_version}\n") + with zipfile.ZipFile(directory / f"{normalized}-0.1.0b0-py3-none-any.whl", "w") as output: + output.writestr(f"{normalized}-0.1.0b0.dist-info/METADATA", metadata) + with tarfile.open(directory / f"{normalized}-0.1.0b0.tar.gz", "w:gz") as output: + for relative in ("PKG-INFO", "LICENSE", "README.md", "pyproject.toml"): + data = metadata.encode() if relative == "PKG-INFO" else b"fixture" + info = tarfile.TarInfo(f"{normalized}-0.1.0b0/{relative}") + info.size = len(data) + output.addfile(info, io.BytesIO(data)) return wheel, source def test_release_audit_accepts_complete_artifact_pair(tmp_path): wheel, source = _artifacts(tmp_path) result = audit(tmp_path) - assert result == {wheel.name: len(WHEEL_REQUIRED) + 1, source.name: len(SDIST_REQUIRED)} + assert result[wheel.name] == len(WHEEL_REQUIRED) + 1 + assert result[source.name] == len(SDIST_REQUIRED) + assert len(result) == 6 @pytest.mark.parametrize("member", ["aspec/compat/sglang_patch.py", - "aspec/compat/patches/sglang-v0.5.21.patch"]) + "aspec/compat/patches/sglang-v0.5.21.patch", + "aspec/_target_worker.py", + "aspec/inference/engine/target_actor.py", + "aspec/inference/service.py", + "aspec/inference/service_protocol.py", + "aspec/inference/service_client.py", + "aspec/inference/managed.py"]) def test_release_audit_requires_bundled_patch_command_and_resource(tmp_path, member): _artifacts(tmp_path, wheel_omit=(member,)) with pytest.raises(ValueError, match="missing wheel members"): @@ -97,16 +119,22 @@ def test_release_audit_rejects_non_source_members(tmp_path, member): def test_release_audit_requires_one_wheel_and_one_sdist(tmp_path): - with pytest.raises(ValueError, match="exactly one wheel and one sdist"): + with pytest.raises(ValueError, match="exactly three wheels and three sdists"): audit(tmp_path) def test_release_audit_rejects_mismatched_package_identity(tmp_path): - _artifacts(tmp_path, sdist_version="0.1.0rc2") + _artifacts(tmp_path, sdist_version="0.2.0b0") with pytest.raises(ValueError, match="does not match"): audit(tmp_path) +def test_release_audit_rejects_the_unrelated_aspec_distribution(tmp_path): + _artifacts(tmp_path, distribution="aspec") + with pytest.raises(ValueError, match="unexpected distribution name"): + audit(tmp_path) + + def test_release_audit_rejects_duplicate_wheel_member(tmp_path): wheel, _ = _artifacts(tmp_path) with pytest.warns(UserWarning, match="Duplicate name"): diff --git a/tests/test_target_environment.py b/tests/test_target_environment.py new file mode 100644 index 0000000..c8c3aea --- /dev/null +++ b/tests/test_target_environment.py @@ -0,0 +1,80 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +"""Isolation regressions; real Ray/venv and GPU checks run on Linux separately.""" +import subprocess +import sys + +import pytest + + +def test_factory_import_does_not_import_sglang(): + result = subprocess.run([sys.executable, "-c", """ +import sys +class Reject: + def find_spec(self, fullname, *args): + if fullname == 'sglang' or fullname.startswith('sglang.'): + raise AssertionError('Draft imported SGLang') +sys.meta_path.insert(0, Reject()) +import aspec.inference.factory +import aspec.inference.engine.target_actor +"""], capture_output=True, text=True) + assert result.returncode == 0, result.stderr + + +def test_target_bootstrap_rebuilds_environment_before_starting_worker(monkeypatch): + import os + + from aspec import _target_worker + + for name in ("PYTHONPATH", "PYTHONHOME", "LD_LIBRARY_PATH", "LIBRARY_PATH"): + monkeypatch.setenv(name, "/draft/foreign") + monkeypatch.setenv("VIRTUAL_ENV", "/draft") + monkeypatch.setenv("PATH", "/draft/bin:/usr/bin") + monkeypatch.setenv("ASPEC_CACHE_HOME", "/tmp/aspec-cache") + monkeypatch.setattr(sys, "argv", ["bootstrap", "ray-worker.py", "--node-id", "123"]) + calls = [] + monkeypatch.setattr(_target_worker, "execute", calls.append) + _target_worker.main() + assert calls == [[sys.executable, "ray-worker.py", "--node-id", "123"]] + assert os.environ["VIRTUAL_ENV"] == sys.prefix + assert os.environ["PATH"].startswith(os.path.dirname(sys.executable) + os.pathsep) + assert os.environ["ASPEC_CACHE_HOME"] == "/tmp/aspec-cache/target" + assert not any(name in os.environ for name in ("PYTHONPATH", "PYTHONHOME", "LD_LIBRARY_PATH", "LIBRARY_PATH")) + + +def test_target_rejects_different_code_before_loading_engine(monkeypatch): + from aspec.inference.engine.target_actor import TargetActor + + with pytest.raises(RuntimeError, match="code differs"): + TargetActor(expected_runtime={"version": "different", "code_sha256": "different"}) + + +@pytest.mark.parametrize("command", ["evaluate", "prepare-target-data", "patch-sglang"]) +def test_target_cli_commands_run_in_the_active_environment(monkeypatch, command): + from aspec import cli + + def missing(_): + raise cli.importlib.metadata.PackageNotFoundError("ling-aspec-target") + + monkeypatch.setattr(cli.importlib.metadata, "version", missing) + monkeypatch.setenv("ASPEC_TARGET_PYTHON", "/target/bin/python") + monkeypatch.setattr(sys, "argv", ["aspec", command]) + calls = [] + monkeypatch.setattr(cli, "execute", calls.append) + cli.main() + assert calls[0][:2] == [sys.executable, "-m"] + + +def test_target_address_is_recorded_in_plan_without_opposing_interpreter(monkeypatch): + from aspec.commands.quickstart import build_profile + from aspec.config.train_config import load_config + from aspec.planning import compile_execution_plan + + monkeypatch.setenv("ASPEC_TARGET_PYTHON", "/environment-target/bin/python") + config = build_profile("qwen3-8b") + config.inference.target_service = {"address": "http://target:30000"} + plan = compile_execution_plan(load_config(base_config=config)) + assert plan.to_runtime_args().target_service["address"] == "http://target:30000" + inference = plan.to_dict()["policy"]["resources"]["inference"] + assert "python_executable" not in inference + assert inference["target_service"]["mode"] == "external" diff --git a/tests/test_target_service.py b/tests/test_target_service.py new file mode 100644 index 0000000..5fbdf25 --- /dev/null +++ b/tests/test_target_service.py @@ -0,0 +1,242 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +"""Control protocol, source ownership and independent launch regressions.""" +import os +import threading +from argparse import Namespace +from types import SimpleNamespace +from unittest.mock import MagicMock + +import psutil +import pytest +import torch + +from aspec.inference.service import TargetService, http_server, readers_exited +from aspec.inference.service_protocol import ( + decode, + encode, + engine_arguments, + request, + validate_address, +) + + +def test_control_payload_roundtrip_and_feature_tensor_rejection(): + payload = {"input_ids_ref": [torch.tensor([1, 3, 5])], + "result": {"tensor_dtypes": {"hidden_states": torch.bfloat16}}} + result = decode(encode(payload)) + assert result["input_ids_ref"][0].tolist() == [1, 3, 5] + assert result["result"]["tensor_dtypes"]["hidden_states"] == torch.bfloat16 + with pytest.raises(TypeError, match="Only CPU token IDs"): + encode(torch.ones(3, dtype=torch.bfloat16)) + with pytest.raises(ValueError, match="dtype"): + decode({"$dtype": "not_a_dtype"}) + + +@pytest.mark.parametrize("address", ["file:///tmp/a", "https://target:30000", "http://target", + "http://user:secret@target:30000", "http://target:30000/run"]) +def test_invalid_target_address_is_rejected(address): + with pytest.raises(ValueError): + validate_address(address) + + +def test_no_driver_objects_cross_control_boundary(): + args = Namespace(model="qwen", draft_model_config_obj=object(), _mooncake_master_actor=object()) + assert engine_arguments(args) == {"model": "qwen"} + args.mystery = object() + with pytest.raises(TypeError, match="mystery"): + engine_arguments(args) + + +def test_complete_training_plan_has_an_explicit_data_only_target_boundary(): + import json + + from aspec.commands.quickstart import build_profile + from aspec.config.train_config import load_config + from aspec.planning import compile_execution_plan + + args = compile_execution_plan(load_config(base_config=build_profile("qwen3-8b"))).to_runtime_args() + payload = engine_arguments(args) + json.dumps(payload, allow_nan=False) + assert payload["aspec_feature_contract"]["contract_id"] + assert payload["target_service"]["mode"] == "external" + + +@pytest.fixture +def service(monkeypatch): + import aspec.inference.service as module + + handle = MagicMock() + handle.get_status.remote.return_value = {"initialized": True} + handle.shutdown.remote.return_value = None + actor = MagicMock() + actor.options.return_value.remote.return_value = handle + import ray + + monkeypatch.setattr(module, "ray", SimpleNamespace(remote=lambda _: actor, + get=lambda x, **_: x, kill=MagicMock(), + exceptions=ray.exceptions)) + instance = TargetService(model="qwen", devices=[0, 1], runtime={}) + instance.test_handle = handle + return instance + + +def creation(service, key="engine", gpu=0): + return {"key": key, "identity": service.identity, + "arguments": {"target_model_path": "qwen", "aspec_attempt_id": "attempt-1", + "aspec_run_id": "run-1"}, + "rank": 0, "base_gpu_id": gpu, "num_gpus_per_engine": 1, + "node_rank": 0, "engine_group": 0} + + +def test_gpu_ownership_and_closed_engine_idempotence(service): + service.create(creation(service)) + with pytest.raises(RuntimeError, match="already belongs"): + service.create(creation(service, "other")) + service.call("engine", "shutdown", [], {}) + service.call("engine", "shutdown", [], {}) + assert service.test_handle.shutdown.remote.call_count == 1 + assert service.idle() + service.create(creation(service, "next")) + + +def test_model_identity_and_allowed_gpu_fail_before_creation(service): + wrong = creation(service) + wrong["identity"] = {} + with pytest.raises(ValueError, match="identity"): + service.create(wrong) + wrong = creation(service) + wrong["arguments"]["target_model_path"] = "other-model" + with pytest.raises(ValueError, match="model"): + service.create(wrong) + with pytest.raises(ValueError, match="GPU set"): + service.create(creation(service, gpu=3)) + assert service.idle() + + +def test_pending_creation_is_not_expired_before_heartbeat_can_start(service, monkeypatch): + import aspec.inference.service as module + + def initializing(value, **_): + service.reap_disconnected(expiry=-1) + assert service.engines["engine"]["state"] == "created" + return value + + monkeypatch.setattr(module.ray, "get", initializing) + service.create(creation(service)) + assert service.engines["engine"]["active"] == 0 + + +def test_native_scheduler_uses_the_callers_run_identity(service, monkeypatch): + import aspec.inference.service as module + + monkeypatch.setenv("ASPEC_RUN_ID", "stale-service-run") + monkeypatch.setenv("ASPEC_ATTEMPT_ID", "stale-service-attempt") + actor = module.ray.remote(None) + service.create(creation(service)) + environment = actor.options.call_args.kwargs["runtime_env"]["env_vars"] + assert environment["ASPEC_RUN_ID"] == "run-1" + assert environment["ASPEC_ATTEMPT_ID"] == "attempt-1" + assert environment["ASPEC_SHARED_SOURCE"] == "0" + + +@pytest.mark.parametrize("second,views", [("contract-a", "0"), ("contract-b", "1")]) +def test_shared_source_view_flags_follow_the_consumer_contract_ids(second, views): + from aspec.inference.service import session_environment + + environment = session_environment(Namespace( + aspec_run_id="run", aspec_attempt_id="attempt", aspec_shared_source=True, + aspec_shared_consumer_contracts={"first": "contract-a", "second": second}, + )) + assert environment["ASPEC_SHARED_SOURCE"] == "1" + assert environment["ASPEC_SHARED_FEATURE_VIEWS"] == views + + +def test_disconnect_retains_unverified_reader_storage(service, monkeypatch): + service.create(creation(service)) + service.reap_disconnected(expiry=-1) + assert service.engines["engine"]["state"] == "retained" + assert not service.idle() + service.test_handle.shutdown.remote.assert_not_called() + monkeypatch.setattr("aspec.inference.service.readers_exited", lambda *_: True) + service.reap_disconnected(expiry=-1) + assert service.idle() + + +def test_failed_source_actor_is_retired_only_after_reader_fence(service, monkeypatch): + import ray + + service.create(creation(service)) + item = service.engines["engine"] + item["processes"] = {999999999: 1.0} + service.test_handle.shutdown.remote.side_effect = ray.exceptions.RayActorError("actor exited") + monkeypatch.setattr("aspec.inference.service.readers_exited", lambda *_: False) + stop = MagicMock() + monkeypatch.setattr("aspec.inference.service.stop_owned_processes", stop) + service.reap_disconnected(expiry=-1) + stop.assert_not_called() + assert not service.idle() + monkeypatch.setattr("aspec.inference.service.readers_exited", lambda *_: True) + service.reap_disconnected(expiry=-1) + stop.assert_called_once_with(item["processes"]) + assert service.idle() + assert item["terminal_reason"] == "actor_failed_readers_fenced_sources_exited" + + +def test_shutdown_does_not_race_active_generation(service): + service.create(creation(service)) + service.engines["engine"]["active"] = 1 + with pytest.raises(RuntimeError, match="active calls"): + service.call("engine", "shutdown", [], {}) + service.test_handle.shutdown.remote.assert_not_called() + + +def test_reader_fence_requires_process_identity_and_local_observation(): + owner = {"pid": os.getpid(), "created": psutil.Process().create_time(), "node_ip": "local"} + assert not readers_exited([owner], "local") + assert not readers_exited([owner], "another-host") + assert not readers_exited([], "local") + assert readers_exited([{**owner, "created": owner["created"] - 100}], "local") + + +def test_real_http_control_roundtrip_and_unknown_method(service): + server = http_server(service, "127.0.0.1", 0) + worker = threading.Thread(target=server.serve_forever) + worker.start() + address = f"http://127.0.0.1:{server.server_port}" + try: + assert request(address, "/v1/info")["identity"] == service.identity + request(address, "/v1/engines", creation(service)) + assert request(address, "/v1/call", {"key": "engine", "method": "get_status", + "args": [], "kwargs": {}}) == {"initialized": True} + with pytest.raises(RuntimeError, match="Unknown Target control method"): + request(address, "/v1/call", {"key": "engine", "method": "__dict__", + "args": [], "kwargs": {}}) + finally: + server.shutdown() + server.server_close() + worker.join() + + +def test_managed_local_keeps_resume_configuration(tmp_path): + from aspec.commands.quickstart import build_profile + from aspec.config.train_config import load_config + + config = build_profile("qwen3-8b") + config.inference.target_service = {"mode": "managed-local", "python_executable": "/target/bin/python"} + config.training.load_path = str(tmp_path / "checkpoint") + loaded = load_config(base_config=config) + assert loaded.training.load_path == str(tmp_path / "checkpoint") + + +def test_managed_supervisor_never_stops_an_existing_service(monkeypatch, tmp_path): + from aspec.inference import managed + + monkeypatch.setattr(managed, "request", lambda *_a, **_k: {}) + popen = MagicMock() + monkeypatch.setattr(managed.subprocess, "Popen", popen) + with pytest.raises(RuntimeError, match="existing Target"): + with managed.managed_target(Namespace(target_service={"mode": "managed-local", + "address": "http://127.0.0.1:30000"})): + pytest.fail("existing service was adopted") + popen.assert_not_called() diff --git a/tests/test_wheel_assets.py b/tests/test_wheel_assets.py index c5f68c8..1bf17da 100644 --- a/tests/test_wheel_assets.py +++ b/tests/test_wheel_assets.py @@ -9,6 +9,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import importlib.metadata import json import subprocess import sys @@ -20,7 +21,7 @@ from aspec.assets import asset_path, recipe_config from aspec.commands import recipes from aspec.config.train_config import load_config -from aspec.evidence import _source_identity +from aspec.evidence import _package_versions, _source_identity from aspec.planning import compile_execution_plan @@ -56,7 +57,18 @@ def test_installed_command_runner_preserves_cwd_and_arguments(tmp_path, monkeypa assert log.read_text().splitlines() == [str(tmp_path), value] -def test_package_identity_hashes_bundled_patch_and_does_not_borrow_parent_git(tmp_path): +def test_package_identity_hashes_bundled_patch_and_does_not_borrow_parent_git(tmp_path, monkeypatch): + versions = {"ling-aspec-runtime": "0.1.0b0", "aspec": "0.0.16"} + + def version(name): + if name not in versions: + raise importlib.metadata.PackageNotFoundError(name) + return versions[name] + + monkeypatch.setattr(importlib.metadata, "version", version) + observed = _package_versions() + assert observed["ling-aspec-runtime"] == "0.1.0b0" + assert "aspec" not in observed subprocess.run(["git", "init", "-q", str(tmp_path)], check=True) package = tmp_path / "venv/lib/site-packages/aspec" package.mkdir(parents=True) @@ -66,6 +78,7 @@ def test_package_identity_hashes_bundled_patch_and_does_not_borrow_parent_git(tm patch.parent.mkdir(parents=True) patch.write_text("first patch\n") first = _source_identity(package) + assert first["version"] == "0.1.0b0" assert first["revision"] is None and first["dirty"] is None assert first["build"] == {"revision": "a" * 40, "dirty": False} patch.write_text("second patch\n") diff --git a/tools/audit_release.py b/tools/audit_release.py index 70a009d..663a75d 100644 --- a/tools/audit_release.py +++ b/tools/audit_release.py @@ -74,6 +74,7 @@ "aspec/models/target/target_utils.py", "aspec/compat/patches/sglang-v0.5.21.patch", "requirements/sglang-v0.5.21.txt", + "requirements/draft.txt", "tools/audit_release.py", "tools/recipe_catalog.json", "tools/build_release.py", @@ -81,6 +82,13 @@ "setup.py", } WHEEL_REQUIRED = { + "aspec/_target_worker.py", + "aspec/inference/target_runtime.py", + "aspec/inference/engine/target_actor.py", + "aspec/inference/service.py", + "aspec/inference/service_client.py", + "aspec/inference/service_protocol.py", + "aspec/inference/managed.py", "aspec/compat/fa4_cutlass.py", "aspec/compat/sglang_patch.py", "aspec/compat/patches/sglang-v0.5.21.patch", @@ -101,7 +109,7 @@ "aspec/commands/_checkpoint_audit.py", "aspec/commands/_recipe_capacity.py", "aspec/resources/requirements/sglang-v0.5.21.txt", - "aspec/resources/requirements/cute-overrides.txt", + "aspec/resources/requirements/draft.txt", "aspec/resources/tools/recipe_catalog.json", "aspec/resources/examples/data/sample_conversations.jsonl", } @@ -138,6 +146,12 @@ def _identity(raw: bytes, artifact: Path) -> tuple[str, str]: version = metadata.get("Version") if not name or not version: raise ValueError(f"{artifact.name}: package metadata lacks Name or Version") + if name not in {"ling-aspec", "ling-aspec-target", "ling-aspec-runtime"}: + raise ValueError(f"{artifact.name}: unexpected distribution name {name!r}") + if name != "ling-aspec-runtime": + requirements = {value.replace(" ", "") for value in metadata.get_all("Requires-Dist", [])} + if f"ling-aspec-runtime=={version}" not in requirements: + raise ValueError(f"{artifact.name}: role must require the matching runtime version") return name, version @@ -153,7 +167,9 @@ def _wheel_members(path: Path) -> tuple[set[str], tuple[str, str]]: if len(metadata_names) != 1: raise ValueError(f"{path.name}: expected one dist-info/METADATA") identity = _identity(wheel.read(metadata_names[0]), path) - missing = WHEEL_REQUIRED - names + missing = WHEEL_REQUIRED - names if identity[0] == "ling-aspec-runtime" else set() + if identity[0] != "ling-aspec-runtime" and any(name.startswith("aspec/") for name in names): + raise ValueError(f"{path.name}: role packages must not own runtime files") if missing: raise ValueError(f"{path.name}: missing wheel members {sorted(missing)}") return names, identity @@ -179,10 +195,10 @@ def _sdist_members(path: Path) -> tuple[set[str], tuple[str, str]]: for member in members if len(PurePosixPath(member.name).parts) > 1 } - missing = SDIST_REQUIRED - names + missing = SDIST_REQUIRED - names if root.startswith("ling_aspec_runtime-") else {"PKG-INFO", "LICENSE", "README.md", "pyproject.toml"} - names if missing: raise ValueError(f"{path.name}: missing source members {sorted(missing)}") - if not root.startswith("aspec-"): + if not root.startswith(("ling_aspec-", "ling_aspec_target-", "ling_aspec_runtime-")): raise ValueError(f"{path.name}: unexpected source root {root!r}") with tarfile.open(path, mode="r:gz") as source: metadata_member = source.getmember(f"{root}/PKG-INFO") @@ -196,18 +212,20 @@ def _sdist_members(path: Path) -> tuple[set[str], tuple[str, str]]: def audit(directory: Path) -> dict[str, int]: wheels = sorted(directory.glob("*.whl")) sdists = sorted(directory.glob("*.tar.gz")) - if len(wheels) != 1 or len(sdists) != 1: - raise ValueError( - f"{directory}: expected exactly one wheel and one sdist, " - f"found {len(wheels)} wheel(s) and {len(sdists)} sdist(s)" - ) - wheel_members, wheel_identity = _wheel_members(wheels[0]) - sdist_members, sdist_identity = _sdist_members(sdists[0]) - if wheel_identity != sdist_identity: - raise ValueError( - f"wheel identity {wheel_identity!r} does not match sdist identity {sdist_identity!r}" - ) - return {wheels[0].name: len(wheel_members), sdists[0].name: len(sdist_members)} + if len(wheels) != 3 or len(sdists) != 3: + raise ValueError(f"{directory}: expected exactly three wheels and three sdists") + result, identities = {}, {} + for artifact in [*wheels, *sdists]: + members, identity = (_wheel_members(artifact) if artifact.suffix == ".whl" + else _sdist_members(artifact)) + result[artifact.name] = len(members) + identities.setdefault(identity, []).append(artifact.suffix == ".whl") + expected = {"ling-aspec", "ling-aspec-target", "ling-aspec-runtime"} + if ({name for name, version in identities} != expected + or len({version for name, version in identities}) != 1 + or any(sorted(kinds) != [False, True] for kinds in identities.values())): + raise ValueError("wheel identity does not match the three version-aligned source identities") + return result def main() -> None: diff --git a/tools/build_release.py b/tools/build_release.py index bc9aedb..ee00aef 100644 --- a/tools/build_release.py +++ b/tools/build_release.py @@ -13,9 +13,12 @@ import argparse import hashlib +import json import shutil import subprocess import sys +import tempfile +import tomllib from pathlib import Path from tools.audit_release import audit @@ -29,11 +32,34 @@ def build_release(directory): raise ValueError("Use an empty release output directory") subprocess.run([sys.executable, "-m", "build", "--outdir", str(directory)], cwd=root, check=True) + project = tomllib.loads((root / "pyproject.toml").read_text())["project"] + for role, manifest in (("draft", "draft.txt"), ("target", "sglang-v0.5.21.txt")): + with tempfile.TemporaryDirectory(prefix="aspec-role-") as staging: + stage = Path(staging) + name = "ling-aspec" if role == "draft" else "ling-aspec-target" + dependencies = [f"ling-aspec-runtime=={project['version']}"] + dependencies += [line.strip() for line in (root / "requirements" / manifest).read_text().splitlines() + if line.strip() and not line.lstrip().startswith("#")] + metadata = ("[build-system]\nrequires = [\"setuptools>=77\", \"wheel\"]\n" + "build-backend = \"setuptools.build_meta\"\n\n[project]\n" + f"name = {json.dumps(name)}\nversion = {json.dumps(project['version'])}\n" + f"description = {json.dumps('ASpec ' + role + ' environment')}\n" + "readme = \"README.md\"\nlicense = \"Apache-2.0\"\n" + "license-files = [\"LICENSE\"]\nrequires-python = \">=3.11\"\n" + f"dependencies = {json.dumps(dependencies)}\n\n" + "[project.urls]\nHomepage = \"https://github.com/inclusionAI/ASpec\"\n" + "\n[tool.setuptools]\npackages = []\n") + (stage / "pyproject.toml").write_text(metadata) + (stage / "README.md").write_text(f"# ASpec {role} environment\n\n" + "See https://github.com/inclusionAI/ASpec for installation and usage.\n") + shutil.copyfile(root / "LICENSE", stage / "LICENSE") + subprocess.run([sys.executable, "-m", "build", "--outdir", str(directory)], + cwd=stage, check=True) for artifact, count in audit(directory).items(): print(f"PASS {artifact}: {count} reviewed members", flush=True) runtime = directory / "runtime" runtime.mkdir() - for name in ("sglang-v0.5.21.txt", "cute-overrides.txt"): + for name in ("sglang-v0.5.21.txt", "draft.txt"): shutil.copyfile(root / "requirements" / name, runtime / name) files = sorted(path for path in directory.rglob("*") if path.is_file()) (directory / "SHA256SUMS").write_text("".join( diff --git a/tools/check_wheel_install.py b/tools/check_wheel_install.py index 1d0f956..da3767a 100644 --- a/tools/check_wheel_install.py +++ b/tools/check_wheel_install.py @@ -44,12 +44,13 @@ def run(arguments, expected=0, contains=None): return result.stdout observed = json.loads(run([sys.executable, "-I", "-c", """ -import importlib.util,json,sys +import importlib.metadata,importlib.util,json,sys from pathlib import Path import aspec from aspec.assets import asset_path from aspec.evidence import _source_identity from aspec.compat.sglang_patch import PATCH +from aspec.commands.convert_to_hf import _aspec_version prefix=Path(sys.prefix) assert Path(aspec.__file__).is_relative_to(prefix),aspec.__file__ assert importlib.util.find_spec('tools') is None @@ -57,7 +58,9 @@ def run(arguments, expected=0, contains=None): assert asset_path('requirements/sglang-v0.5.21.txt').is_relative_to(prefix) identity=_source_identity() assert identity['revision'] is None and identity['dirty'] is None -assert identity['source_file_count'] > 0 and identity['version'] +assert identity['source_file_count'] > 0 +assert identity['version'] == importlib.metadata.version('ling-aspec-runtime') +assert _aspec_version() == identity['version'] assert identity['build'] is not None spec=importlib.util.find_spec('sglang') assert Path(spec.origin).is_relative_to(prefix) @@ -83,6 +86,9 @@ def run(arguments, expected=0, contains=None): doctor = json.loads(run([cli, "doctor", "--json"], expected=1)) if doctor["ok"] or not doctor["checks"] or doctor["gpu_requested"]: raise RuntimeError("CPU fixture must not claim full runtime readiness") + target_doctor = json.loads(run([cli, "doctor", "--target", "--json"], expected=1)) + if not target_doctor["target_requested"]: + raise RuntimeError("Target doctor must require the Target runtime") package = Path(observed["sglang"]) def hashes(): diff --git a/tools/install.sh b/tools/install.sh index beb6809..0a088ab 100755 --- a/tools/install.sh +++ b/tools/install.sh @@ -1,13 +1,15 @@ #!/usr/bin/env bash set -euo pipefail dev=0 +role=draft gpu=() for argument in "$@"; do case "$argument" in --dev) dev=1 ;; + --target) role=target ;; --gpu) gpu=(--gpu) ;; --help) - echo 'Usage: bash tools/install.sh [--dev] [--gpu]' + echo 'Usage: bash tools/install.sh [--target] [--dev] [--gpu]' echo 'For an already activated isolated venv. New users: bash setup.sh' exit 0 ;; *) printf 'Unknown argument: %s\n' "$argument" >&2; exit 2 ;; @@ -19,20 +21,12 @@ bash tools/run_in_env.sh true command -v git >/dev/null || { echo 'Install Git before setup.' >&2; exit 2; } python - <<'CHECK' import ctypes -import os import platform -import shutil -from pathlib import Path name, version = platform.libc_ver() if name != 'glibc' or tuple(map(int, version.split('.')[:2])) < (2, 34): raise SystemExit('Pinned SGLang wheels require glibc >= 2.34; use a newer Linux base image.') if platform.machine() != 'x86_64': raise SystemExit('The pinned runtime requires Linux x86-64.') -cuda_root = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") -nvcc = str(Path(cuda_root) / "bin/nvcc") if cuda_root else shutil.which("nvcc") -if not nvcc or not os.access(nvcc, os.X_OK): - raise SystemExit("SGLang JIT requires a CUDA development toolkit with nvcc. " - "Set CUDA_HOME to that toolkit or use the development container.") for name in ('libibverbs.so.1', 'libmlx5.so.1'): try: ctypes.CDLL(name) @@ -46,11 +40,18 @@ if ! printf '#include \nint main() { return std::bit_cast(1.0f) = echo 'Install a modern C++ toolchain (Ubuntu: g++), or set CXX to one; see installation-advanced.md.' >&2 exit 2 fi -extras=fa -(( dev == 0 )) || extras=dev,$extras -python -m pip install uv==0.13.5 -python -m uv pip install --link-mode copy --python "$VIRTUAL_ENV/bin/python" \ - --override requirements/cute-overrides.txt -r requirements/sglang-v0.5.21.txt -e ".[${extras}]" -aspec patch-sglang -aspec patch-sglang --check -bash tools/run_in_env.sh python -m tools.doctor "${gpu[@]}" +manifest=requirements/draft.txt +check=() +if [[ $role == target ]]; then + manifest=requirements/sglang-v0.5.21.txt + check=(--target) +fi +extras="." +(( dev == 0 )) || extras=".[dev]" +python -m pip install -r "$manifest" -e "$extras" +printf '%s\n' "$role" > "$VIRTUAL_ENV/.aspec-role" +if [[ $role == target ]]; then + aspec patch-sglang + aspec patch-sglang --check +fi +aspec doctor "${check[@]}" "${gpu[@]}"