From fe7bc45474eec7c23bb5bfb35f100ffaef0c3942 Mon Sep 17 00:00:00 2001 From: zhouaihui Date: Thu, 12 Nov 2026 10:29:23 +0800 Subject: [PATCH 01/13] feat(release): publish ling-aspec beta packages to PyPI --- .github/workflows/build-release.yml | 44 +++++++++++- CHANGELOG.md | 6 ++ CONTRIBUTING.md | 6 +- README.md | 3 + README_zh-CN.md | 3 + aspec/commands/convert_to_hf.py | 2 +- aspec/evidence.py | 4 +- docs/framework/DECISIONS.md | 12 +++- docs/releasing.md | 107 ++++++++++++++++++++++++++-- docs/user-guide/installation.md | 9 ++- pyproject.toml | 9 ++- tests/test_convert_to_hf.py | 6 ++ tests/test_release_audit.py | 28 +++++--- tests/test_wheel_assets.py | 17 ++++- tools/audit_release.py | 4 +- tools/check_wheel_install.py | 7 +- 16 files changed, 233 insertions(+), 34 deletions(-) diff --git a/.github/workflows/build-release.yml b/.github/workflows/build-release.yml index 08ea1d4..5568df8 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: @@ -80,3 +84,39 @@ jobs: 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/CHANGELOG.md b/CHANGELOG.md index 9c2cd9e..b4c7fa2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,12 @@ ## Unreleased +- 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..2b081a3 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -32,8 +32,10 @@ Default tests are fast/offline. GPU, tokenizer, training and cuDNN gates need ex 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/README.md b/README.md index a58f221..e01d274 100644 --- a/README.md +++ b/README.md @@ -54,6 +54,9 @@ Check [supported paths and limitations](docs/support.md), then browse the [recip ## Start here +The PyPI distribution is named `ling-aspec`; the Python import and command remain `aspec`. +Use the [installation guide](docs/user-guide/installation.md) for the complete pinned GPU runtime. + 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. diff --git a/README_zh-CN.md b/README_zh-CN.md index a46b68a..16a3d85 100644 --- a/README_zh-CN.md +++ b/README_zh-CN.md @@ -49,6 +49,9 @@ Qwen、Ling 上的 DSpark/DFlash 是主要优化方向和 recipe 维护重点； ## 开始使用 +PyPI 发行包名为 `ling-aspec`，Python 导入名和命令仍为 `aspec`。 +完整的固定版本 GPU 运行环境请按[安装指南](docs/user-guide/installation.md)配置。 + 1. 按[安装指南](docs/user-guide/installation.md)一键创建源码环境，也可选择[容器](docs/user-guide/containers.md)； 2. 在[快速上手](docs/user-guide/quickstart.md)中选择 Qwen3-8B 或 Ling tiny，用两张 GPU 完成 相同的 Plan、训练、报告、恢复和导出流程； diff --git a/aspec/commands/convert_to_hf.py b/aspec/commands/convert_to_hf.py index a469064..bc42451 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") except importlib.metadata.PackageNotFoundError: return "unknown" diff --git a/aspec/evidence.py b/aspec/evidence.py index 25677fd..3e2ad1a 100644 --- a/aspec/evidence.py +++ b/aspec/evidence.py @@ -40,7 +40,7 @@ EVIDENCE_SCHEMA_VERSION = "aspec.run-evidence/v0alpha8" _PACKAGE_NAMES = ( - "aspec", + "ling-aspec", "torch", "ray", "sglang", @@ -179,7 +179,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") except importlib.metadata.PackageNotFoundError: identity["version"] = None source_files = [path for path in package.rglob("*") if path.suffix in diff --git a/docs/framework/DECISIONS.md b/docs/framework/DECISIONS.md index f41f89e..4f754cd 100644 --- a/docs/framework/DECISIONS.md +++ b/docs/framework/DECISIONS.md @@ -5,9 +5,15 @@ 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.** Developers maintain the major/minor/patch version and currently retain + the `b0` beta suffix for public releases, for example `0.1.0b0`, `0.2.0b0` and `0.2.1b0`. + 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). diff --git a/docs/releasing.md b/docs/releasing.md index aa50eba..d16095e 100644 --- a/docs/releasing.md +++ b/docs/releasing.md @@ -1,10 +1,28 @@ # Preparing a release +The distribution name is `ling-aspec`; the project name, Python import and CLI remain ASpec / `aspec`. +Setuptools normalizes archive filenames to `ling_aspec-...`. Do not upload under the +unrelated PyPI project `aspec`. + 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. +## Versioning + +Current public PyPI releases use the `b0` beta suffix, starting with `0.1.0b0`. +Developers maintain the major/minor/patch version in `pyproject.toml` and retain that suffix, +for example `0.1.0b0`, `0.2.0b0`, then `0.2.1b0`. Each upload containing changed source or +artifacts needs a new version selected by the developer. The current policy does not advance +the beta suffix to `b1`/`b2` or automatically promote a release to an unsuffixed version. + +The release tag is exactly `v` plus that complete version, such as `v0.1.0b0`. +CI reads and validates the selected version; it does not choose or increment it. +The team's current "debug" release designation uses the standard +[Python beta version syntax](https://packaging.python.org/en/latest/specifications/version-specifiers/#pre-releases). +An exact install pin such as `ling-aspec==0.1.0b0` selects that prerelease explicitly. + ## Prepare 1. Select a clean, reviewed commit and set its version in `pyproject.toml`. @@ -31,10 +49,12 @@ files under `runtime/`, and writes `SHA256SUMS`. It rejects a nonempty output di 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 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 +The separate [release build workflow](../.github/workflows/build-release.yml) supports manual +validation and runs when a `v*` tag is pushed to GitHub. A default manual run builds and checks +artifacts without publishing. Creating a local tag alone does not trigger it. On tag pushes, the tag must exactly match `v` plus +the version in `pyproject.toml`, for example `v0.1.0b0` for `0.1.0b0`; a mismatch fails before +building. The workflow runs the same build command, checks distribution metadata with +`twine check --strict`, 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: @@ -66,16 +86,91 @@ Evidence records the effective installed package hash (including the bundled pat distribution version and build provenance when available. A wheel does not borrow Git identity from an enclosing checkout; unavailable build revision remains unknown. +## Rehearse before publishing + +### Build and install without uploading + +First run the source CI for the same commit and require its quality and CPU contract jobs to pass. +The release workflow checks packaging and wheel installation; it does not rerun that source suite. + +Once this workflow configuration is on the repository's default branch, open **Actions → Build and +publish release → Run workflow**, select the branch to check, and run the workflow. +GitHub requires the manual trigger to be present on the default branch before exposing this entry; +see [manual workflow runs](https://docs.github.com/en/actions/how-tos/manage-workflow-runs/manually-run-a-workflow). + +This runs `build` and `wheel-install`, retaining the audited bundle in `aspec-distributions`. +The publishing job is skipped. The build and installation jobs receive no publishing secrets. +Manual runs never publish to production PyPI, even when the selected ref is a version tag. +Repeat these checks without changing the version or creating a release tag. + +The installation check covers installed command origins, version identity, assets, recipe/Quickstart +Plans and SGLang patch behavior. Before the first production release, download the audited bundle +and follow the [wheel installation guide](user-guide/installation.md#install-a-release-wheel) in a +fresh Linux venv on allocated GPUs. Complete Quickstart training, recovery and export; use the +relevant [validation gates](contributing/validation.md) for runtime changes. The CPU fixture does not +establish GPU or complete production dependency acceptance. + ## Publish +### One-time PyPI and GitHub setup + +The workflow uses your PyPI API token. Configure it in the repository that publishes releases, +[`inclusionAI/ASpec`](https://github.com/inclusionAI/ASpec), rather than the development fork: + +1. Sign in to [PyPI account settings](https://pypi.org/manage/account/) and create or verify an API + token. The first upload creates `ling-aspec`, so it needs an account-wide token when the project + does not exist yet. A token scoped to another project cannot create it. After the first successful + release, replace it with a token scoped to `ling-aspec` and revoke the bootstrap token. +2. Open **Settings → Environments → New environment** in the GitHub repository and name it **`pypi`**. + Under **Environment secrets**, add **`PYPI_API_TOKEN`** with the entire token value, including + its `pypi-` prefix. Enter the token directly in GitHub; keep it out of source files and logs. +3. Configure **Required reviewers** on the environment to require approval before each upload. + Choose maintainers who can approve releases. If the tag creator must approve their own release, + leave **Prevent self-review** disabled. Environment approval happens after the build/install jobs. +4. Under **Deployment branches and tags**, select **Selected branches and tags** and add a rule + with ref type **Tag** and pattern **`v*`**. A branch-only rule will not permit version tags. +5. Ensure GitHub Actions is enabled and any organization/repository action allowlist permits the + actions referenced by the workflow, including `pypa/gh-action-pypi-publish`. + +See PyPA's [first-upload instructions](https://packaging.python.org/en/latest/guides/distributing-packages-using-setuptools/#create-an-account), +GitHub's [environment configuration guide](https://docs.github.com/en/actions/how-tos/deploy/configure-and-manage-deployments/manage-environments) +and the PyPA action's [API token usage](https://github.com/pypa/gh-action-pypi-publish#specifying-a-different-username). +This workflow explicitly uses token authentication and disables attestations, which require Trusted +Publishing. No PyPI Trusted Publisher or GitHub `id-token: write` permission is needed for this setup. + +### Release a version + 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. -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. +1. Merge the reviewed release changes, including the version in `pyproject.toml`, into the main + repository. Choose a version that has not already been uploaded or tagged; `0.1.0b0` requires + exactly `v0.1.0b0`. Creating a GitHub Release for an existing tag does not send a new tag-push event. +2. Push that tag to `inclusionAI/ASpec` at the reviewed commit. Alternatively, in GitHub's + **Releases → Draft a new release**, create the new tag at that commit and publish the release. + Mark a GitHub Release for a `b0` version as a pre-release. + Pushing a version tag authorizes the publishing workflow; with required reviewers configured, + uploading still waits for environment approval. +3. Open **Actions → Build and publish release**. After `build` and `wheel-install` pass, review the + `aspec-distributions` artifact and approve the pending `pypi` deployment. +4. The `publish` job downloads that same artifact, verifies `SHA256SUMS`, and uploads only its wheel + and sdist. It does not rebuild them or upload `runtime/` and `SHA256SUMS` as Python packages. +5. Verify the files and version on [PyPI](https://pypi.org/project/ling-aspec/), then follow the + [pinned runtime installation guide](user-guide/installation.md) in a fresh Linux venv. + +The publishing job runs only in `inclusionAI/ASpec`; forks still build and check artifacts but skip +publication. Production publishing requires a `v*` tag push; manual runs cannot enable it. The +workflow does not create a GitHub Release itself. The `0.1.0b0` +version is a prerelease; installing that version requires an exact pin or `--pre` until a stable +release is available. + +Missing tokens, rejected metadata, failed checks and duplicate uploads fail the workflow. PyPI +does not allow replacing an existing distribution file. Before retrying an upload failure, inspect +PyPI for partially uploaded files; a rerun does not silently skip existing files. Keep failures and +their logs, and use a new version for source changes. 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; diff --git a/docs/user-guide/installation.md b/docs/user-guide/installation.md index 7864b6f..c1f5293 100644 --- a/docs/user-guide/installation.md +++ b/docs/user-guide/installation.md @@ -5,6 +5,8 @@ toolkit containing `nvcc`. Training requires NVIDIA GPUs and a CUDA-13-compatibl 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. +The distribution name on PyPI is `ling-aspec`; the Python import and CLI remain `aspec`. +The PyPI project named `aspec` is unrelated. Use a fresh venv so overlapping packages cannot mix. 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). @@ -73,9 +75,14 @@ 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]' + '/path/to/release/ling_aspec-0.1.0b0-py3-none-any.whl[fa]' ``` +To use a published PyPI release, replace the wheel argument in that command with +`'ling-aspec[fa]==0.1.0b0'`, choosing the version that matches the audited runtime requirement +files. Installing `ling-aspec` alone does not install or validate the complete pinned GPU runtime; +keep the requirements and explicit overrides above. + 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: diff --git a/pyproject.toml b/pyproject.toml index 4787caa..5c0ec29 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,8 +3,8 @@ requires = ["setuptools>=77.0", "wheel"] build-backend = "setuptools.build_meta" [project] -name = "aspec" -version = "0.1.0rc1" +name = "ling-aspec" +version = "0.1.0b0" description = "A Torch-native framework for speculative decoding training" readme = "README.md" license = "Apache-2.0" @@ -36,6 +36,11 @@ dependencies = [ [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/tests/test_convert_to_hf.py b/tests/test_convert_to_hf.py index e1afcea..a3a5a67 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": "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_release_audit.py b/tests/test_release_audit.py index eb8ab62..ca2cfc6 100644 --- a/tests/test_release_audit.py +++ b/tests/test_release_audit.py @@ -28,34 +28,36 @@ 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", ): - 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) @@ -102,11 +104,17 @@ def test_release_audit_requires_one_wheel_and_one_sdist(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_wheel_assets.py b/tests/test_wheel_assets.py index c5f68c8..260c8b3 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": "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"] == "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..c2b3696 100644 --- a/tools/audit_release.py +++ b/tools/audit_release.py @@ -138,6 +138,8 @@ 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 != "ling-aspec": + raise ValueError(f"{artifact.name}: unexpected distribution name {name!r}; expected 'ling-aspec'") return name, version @@ -182,7 +184,7 @@ def _sdist_members(path: Path) -> tuple[set[str], tuple[str, str]]: missing = SDIST_REQUIRED - names if missing: raise ValueError(f"{path.name}: missing source members {sorted(missing)}") - if not root.startswith("aspec-"): + if not root.startswith("ling_aspec-"): 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") diff --git a/tools/check_wheel_install.py b/tools/check_wheel_install.py index 1d0f956..2f46761 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') +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) From eeec38fcdc2c019d4af63af594c5eecef33283ad Mon Sep 17 00:00:00 2001 From: zhouaihui Date: Thu, 12 Nov 2026 11:19:23 +0800 Subject: [PATCH 02/13] feat(install): package training dependencies and prepare Target runtime --- CHANGELOG.md | 2 + README.md | 36 ++-- README_zh-CN.md | 31 ++-- aspec/commands/doctor.py | 32 ++-- aspec/compat/sglang_patch.py | 31 +++- docs/framework/DECISIONS.md | 8 +- docs/releasing.md | 207 +++++------------------ docs/user-guide/installation-advanced.md | 14 +- docs/user-guide/installation.md | 127 +++++--------- pyproject.toml | 18 +- tests/test_dependency_audit.py | 19 +++ tests/test_dependency_extras.py | 14 ++ tests/test_open_source_entrypoints.py | 48 ++++++ tools/check_wheel_install.py | 3 + 14 files changed, 301 insertions(+), 289 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b4c7fa2..f81ed0d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ## Unreleased +- Install training dependencies with the default distribution. Target nodes can prepare SGLang + and its training patch with `aspec patch-sglang`; `aspec doctor --target` checks that runtime. - 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 diff --git a/README.md b/README.md index e01d274..943e925 100644 --- a/README.md +++ b/README.md @@ -52,25 +52,30 @@ 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 -The PyPI distribution is named `ling-aspec`; the Python import and command remain `aspec`. -Use the [installation guide](docs/user-guide/installation.md) for the complete pinned GPU runtime. +In an activated Linux virtual environment: -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). +```bash +pip install ling-aspec +``` -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). +This installs the training dependencies. On nodes that run the Target model, prepare SGLang +and apply the training patch: + +```bash +aspec patch-sglang +``` + +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. + +## Quickstart -From the repository root: +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 @@ -78,9 +83,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 16a3d85..5e15ad0 100644 --- a/README_zh-CN.md +++ b/README_zh-CN.md @@ -47,26 +47,37 @@ Qwen、Ling 上的 DSpark/DFlash 是主要优化方向和 recipe 维护重点； 从首次多卡运行扩展到静态多机训练，支持 NVLink、TCP 和 RDMA 数据路径。 先查看[支持范围与已知限制](docs/support.md)，再从 [recipe 导航](docs/user-guide/recipes.md)选择起点。 -## 开始使用 +## 安装 -PyPI 发行包名为 `ling-aspec`，Python 导入名和命令仍为 `aspec`。 -完整的固定版本 GPU 运行环境请按[安装指南](docs/user-guide/installation.md)配置。 +在已激活的 Linux 虚拟环境中安装： -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)选择算法/结构并准备实际数据。 +```bash +pip install ling-aspec +``` + +该命令会安装训练依赖。需要运行 Target 模型的节点，再准备 SGLang 并应用训练补丁： + +```bash +aspec patch-sglang +``` + +环境要求、源码安装和本地 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/commands/doctor.py b/aspec/commands/doctor.py index 507c283..679a9da 100644 --- a/aspec/commands/doctor.py +++ b/aspec/commands/doctor.py @@ -71,9 +71,9 @@ def expected_versions(): for name, version in [line.split("==", 1)]} -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 +82,31 @@ 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") + has_sglang = False for name, expected in expected_versions().items(): try: actual = importlib.metadata.version(name) except importlib.metadata.PackageNotFoundError: actual = None + if name == "sglang": + has_sglang = actual is not None + if not target and not has_sglang: + continue ok = actual is not None and 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)) + 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 +132,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/compat/sglang_patch.py b/aspec/compat/sglang_patch.py index 0f9c405..2ddd081 100644 --- a/aspec/compat/sglang_patch.py +++ b/aspec/compat/sglang_patch.py @@ -12,16 +12,20 @@ # 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.""" +"""Prepare the pinned Target runtime and apply/check its training patch.""" import argparse import importlib.metadata import importlib.util import os +import shutil import subprocess import sys from pathlib import Path +from aspec._environment import command_environment +from aspec.assets import asset_path + VERSION = "0.5.21" PATCH = Path(__file__).with_name("patches") / "sglang-v0.5.21.patch" @@ -88,12 +92,33 @@ def installed_site(*, writable=False): return package.parent +def install_target_runtime(): + """Reuse tools/install.sh's explicit overrides in this isolated environment.""" + env = command_environment() + if shutil.which("git", path=env.get("PATH")) is None: + raise RuntimeError("Install Git before preparing the SGLang training patch") + print(f"Installing the pinned SGLang Target runtime in {sys.prefix}", flush=True) + subprocess.run( + [sys.executable, "-m", "uv", "pip", "install", "--link-mode", "copy", + "--python", sys.executable, + "--override", str(asset_path("requirements/cute-overrides.txt")), + "-r", str(asset_path("requirements/sglang-v0.5.21.txt"))], + env=env, check=True, + ) + + def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--check", action="store_true", help="Require the complete patch; do not change files") args = parser.parse_args() try: - site = installed_site(writable=not args.check) + try: + site = installed_site(writable=not args.check) + except importlib.metadata.PackageNotFoundError: + if args.check: + raise + install_target_runtime() + site = installed_site(writable=True) state = patch_state(site) if state == "unpatched": if args.check: @@ -104,7 +129,7 @@ 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 (ValueError, RuntimeError, OSError, importlib.metadata.PackageNotFoundError, subprocess.CalledProcessError) as exc: parser.exit(1, f"ASpec patch check failed: {exc}\n") diff --git a/docs/framework/DECISIONS.md b/docs/framework/DECISIONS.md index 4f754cd..c59bbf2 100644 --- a/docs/framework/DECISIONS.md +++ b/docs/framework/DECISIONS.md @@ -10,8 +10,8 @@ These decisions describe the maintained implementation and its public behavior. 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.** Developers maintain the major/minor/patch version and currently retain - the `b0` beta suffix for public releases, for example `0.1.0b0`, `0.2.0b0` and `0.2.1b0`. +- **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. @@ -21,6 +21,10 @@ These decisions describe the maintained implementation and its public behavior. 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 migration references do not constrain ASpec to their older runtime versions. +- **Installation.** The default distribution installs the Draft runtime. Target nodes explicitly + run `aspec patch-sglang`, which prepares a missing SGLang runtime with the documented overrides + and applies the patch. Read-only patch checks never install dependencies. Doctor checks the + Target runtime when present or explicitly requested with `--target`. - **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 d16095e..77ebbf5 100644 --- a/docs/releasing.md +++ b/docs/releasing.md @@ -1,177 +1,64 @@ -# Preparing a release +# Releasing -The distribution name is `ling-aspec`; the project name, Python import and CLI remain ASpec / `aspec`. -Setuptools normalizes archive filenames to `ling_aspec-...`. Do not upload under the -unrelated PyPI project `aspec`. - -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 is distributed on PyPI as `ling-aspec`. Python imports and the CLI use `aspec`. ## Versioning -Current public PyPI releases use the `b0` beta suffix, starting with `0.1.0b0`. -Developers maintain the major/minor/patch version in `pyproject.toml` and retain that suffix, -for example `0.1.0b0`, `0.2.0b0`, then `0.2.1b0`. Each upload containing changed source or -artifacts needs a new version selected by the developer. The current policy does not advance -the beta suffix to `b1`/`b2` or automatically promote a release to an unsuffixed version. - -The release tag is exactly `v` plus that complete version, such as `v0.1.0b0`. -CI reads and validates the selected version; it does not choose or increment it. -The team's current "debug" release designation uses the standard -[Python beta version syntax](https://packaging.python.org/en/latest/specifications/version-specifiers/#pre-releases). -An exact install pin such as `ling-aspec==0.1.0b0` selects that prerelease explicitly. - -## Prepare - -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) supports manual -validation and runs when a `v*` tag is pushed to GitHub. A default manual run builds and checks -artifacts without publishing. Creating a local tag alone does not trigger it. On tag pushes, the tag must exactly match `v` plus -the version in `pyproject.toml`, for example `v0.1.0b0` for `0.1.0b0`; a mismatch fails before -building. The workflow runs the same build command, checks distribution metadata with -`twine check --strict`, 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 + +In the release repository, configure a GitHub environment named `pypi`: + +- Add the environment secret `PYPI_API_TOKEN` with a PyPI token authorized for `ling-aspec`. + 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. -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. - -## Rehearse before publishing - -### Build and install without uploading - -First run the source CI for the same commit and require its quality and CPU contract jobs to pass. -The release workflow checks packaging and wheel installation; it does not rerun that source suite. - -Once this workflow configuration is on the repository's default branch, open **Actions → Build and -publish release → Run workflow**, select the branch to check, and run the workflow. -GitHub requires the manual trigger to be present on the default branch before exposing this entry; -see [manual workflow runs](https://docs.github.com/en/actions/how-tos/manage-workflow-runs/manually-run-a-workflow). - -This runs `build` and `wheel-install`, retaining the audited bundle in `aspec-distributions`. -The publishing job is skipped. The build and installation jobs receive no publishing secrets. -Manual runs never publish to production PyPI, even when the selected ref is a version tag. -Repeat these checks without changing the version or creating a release tag. - -The installation check covers installed command origins, version identity, assets, recipe/Quickstart -Plans and SGLang patch behavior. Before the first production release, download the audited bundle -and follow the [wheel installation guide](user-guide/installation.md#install-a-release-wheel) in a -fresh Linux venv on allocated GPUs. Complete Quickstart training, recovery and export; use the -relevant [validation gates](contributing/validation.md) for runtime changes. The CPU fixture does not -establish GPU or complete production dependency acceptance. +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). ## Publish -### One-time PyPI and GitHub setup - -The workflow uses your PyPI API token. Configure it in the repository that publishes releases, -[`inclusionAI/ASpec`](https://github.com/inclusionAI/ASpec), rather than the development fork: - -1. Sign in to [PyPI account settings](https://pypi.org/manage/account/) and create or verify an API - token. The first upload creates `ling-aspec`, so it needs an account-wide token when the project - does not exist yet. A token scoped to another project cannot create it. After the first successful - release, replace it with a token scoped to `ling-aspec` and revoke the bootstrap token. -2. Open **Settings → Environments → New environment** in the GitHub repository and name it **`pypi`**. - Under **Environment secrets**, add **`PYPI_API_TOKEN`** with the entire token value, including - its `pypi-` prefix. Enter the token directly in GitHub; keep it out of source files and logs. -3. Configure **Required reviewers** on the environment to require approval before each upload. - Choose maintainers who can approve releases. If the tag creator must approve their own release, - leave **Prevent self-review** disabled. Environment approval happens after the build/install jobs. -4. Under **Deployment branches and tags**, select **Selected branches and tags** and add a rule - with ref type **Tag** and pattern **`v*`**. A branch-only rule will not permit version tags. -5. Ensure GitHub Actions is enabled and any organization/repository action allowlist permits the - actions referenced by the workflow, including `pypa/gh-action-pypi-publish`. - -See PyPA's [first-upload instructions](https://packaging.python.org/en/latest/guides/distributing-packages-using-setuptools/#create-an-account), -GitHub's [environment configuration guide](https://docs.github.com/en/actions/how-tos/deploy/configure-and-manage-deployments/manage-environments) -and the PyPA action's [API token usage](https://github.com/pypa/gh-action-pypi-publish#specifying-a-different-username). -This workflow explicitly uses token authentication and disables attestations, which require Trusted -Publishing. No PyPI Trusted Publisher or GitHub `id-token: write` permission is needed for this setup. - -### Release a version - -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. - -1. Merge the reviewed release changes, including the version in `pyproject.toml`, into the main - repository. Choose a version that has not already been uploaded or tagged; `0.1.0b0` requires - exactly `v0.1.0b0`. Creating a GitHub Release for an existing tag does not send a new tag-push event. -2. Push that tag to `inclusionAI/ASpec` at the reviewed commit. Alternatively, in GitHub's - **Releases → Draft a new release**, create the new tag at that commit and publish the release. - Mark a GitHub Release for a `b0` version as a pre-release. - Pushing a version tag authorizes the publishing workflow; with required reviewers configured, - uploading still waits for environment approval. -3. Open **Actions → Build and publish release**. After `build` and `wheel-install` pass, review the - `aspec-distributions` artifact and approve the pending `pypi` deployment. -4. The `publish` job downloads that same artifact, verifies `SHA256SUMS`, and uploads only its wheel - and sdist. It does not rebuild them or upload `runtime/` and `SHA256SUMS` as Python packages. -5. Verify the files and version on [PyPI](https://pypi.org/project/ling-aspec/), then follow the - [pinned runtime installation guide](user-guide/installation.md) in a fresh Linux venv. - -The publishing job runs only in `inclusionAI/ASpec`; forks still build and check artifacts but skip -publication. Production publishing requires a `v*` tag push; manual runs cannot enable it. The -workflow does not create a GitHub Release itself. The `0.1.0b0` -version is a prerelease; installing that version requires an exact pin or `--pre` until a stable -release is available. - -Missing tokens, rejected metadata, failed checks and duplicate uploads fail the workflow. PyPI -does not allow replacing an existing distribution file. Before retrying an upload failure, inspect -PyPI for partially uploaded files; a rerun does not silently skip existing files. Keep failures and -their logs, and use a new version for source changes. - -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/installation-advanced.md b/docs/user-guide/installation-advanced.md index d5ad5ae..6c604fb 100644 --- a/docs/user-guide/installation-advanced.md +++ b/docs/user-guide/installation-advanced.md @@ -1,17 +1,22 @@ # Advanced installation and native prerequisites -Start with [one-command source setup](installation.md). This page explains its pinned lane, +Start with the [installation guide](installation.md). This page explains the supported runtime, native prerequisites, compatibility exceptions 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. +`pip install ling-aspec` installs the Draft training runtime. On Target nodes, +`aspec patch-sglang` also prepares a missing SGLang installation using the bundled dependency +manifest and compatibility overrides before applying the patch. `aspec doctor --target` checks +that setup; `aspec doctor` can check a training environment without SGLang. + ## 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. @@ -73,7 +78,8 @@ 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 +combination together with upstream SGLang. Target setup handles this explicit override separately +from the default pip installation. 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 diff --git a/docs/user-guide/installation.md b/docs/user-guide/installation.md index c1f5293..cf13afa 100644 --- a/docs/user-guide/installation.md +++ b/docs/user-guide/installation.md @@ -1,116 +1,79 @@ # 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. -The distribution name on PyPI is `ling-aspec`; the Python import and CLI remain `aspec`. -The PyPI project named `aspec` is unrelated. Use a fresh venv so overlapping packages cannot mix. +Use Linux x86-64, Python >=3.11 and an isolated virtual environment. Training requires NVIDIA +GPUs, a CUDA-compatible driver and a CUDA development toolkit with `nvcc`. +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 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 + +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 +aspec patch-sglang ``` -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. - -`aspec` commands prepare this venv's native-library paths automatically. The -`tools/run_in_env.sh` wrapper remains available for arbitrary Python/torchrun commands. - -## First run +The command installs the compatible SGLang runtime if it is missing, then applies the training +patch. Runtime versions and compatibility overrides are shipped with ASpec; no manual dependency +selection is needed. See [advanced installation](installation-advanced.md#dedicated-environment) +for the override details. -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. - -## Other environments +Applying the patch again verifies an already patched installation. Conflicting changes or an +unsupported SGLang version are rejected. `--check` never installs dependencies or modifies files. -- [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. +## Check the environment -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. +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. -## 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/ling_aspec-0.1.0b0-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 ``` -To use a published PyPI release, replace the wheel argument in that command with -`'ling-aspec[fa]==0.1.0b0'`, choosing the version that matches the audited runtime requirement -files. Installing `ling-aspec` alone does not install or validate the complete pinned GPU runtime; -keep the requirements and explicit overrides above. +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 -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: +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 runtime, applies the patch and runs doctor. Contributors can use +`bash setup.sh --dev`. Run `bash setup.sh --help` for environment and interpreter options. -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: +## Install a release wheel -```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 -``` +Replace `ling-aspec` in the PyPI install command with the path to the downloaded wheel. +Then prepare Target nodes and optionally check the environment as above. The installed commands, +recipes and Quickstart data work outside the source checkout. -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. +See [advanced installation](installation-advanced.md) for native-library paths and manual +dependency management, or [troubleshooting](troubleshooting.md) if a check fails. diff --git a/pyproject.toml b/pyproject.toml index 5c0ec29..b016e7a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,12 +12,19 @@ license-files = ["LICENSE", "licenses/*"] requires-python = ">=3.11" dependencies = [ "accelerate", + "apache-tvm-ffi==0.1.11", "cmake", + "cuda-tile==1.6.0rc5", "datasets>=5.0.1,<6", - "mooncake-transfer-engine>=0.3.10.post1", + "flash-attn-4==4.0.0b34", + "flashinfer-python[cu13]==0.6.18", + "mooncake-transfer-engine==0.3.13.post1", "ninja", "numba", "numpy<2.4", + "nvidia-cuda-runtime-cu12==12.9.79", + "nvidia-cudnn-frontend==1.31.0", + "nvidia-cutlass-dsl[cu13]==4.7.1", "omegaconf", "openai", "openai-harmony", @@ -26,10 +33,15 @@ dependencies = [ "pydantic", "pyzmq", "qwen-vl-utils", - "ray", + "quack-kernels[cu13]==0.6.5", + "ray==2.59.0", "setuptools", + "sglang-kernel==0.4.7", + "torch==2.13.0", + "torchvision==0.28.0", "tqdm", - "transformers", + "transformers==5.12.1", + "uv==0.13.5", "wandb>=0.19.10", ] diff --git a/tests/test_dependency_audit.py b/tests/test_dependency_audit.py index 6080cbe..e84812f 100644 --- a/tests/test_dependency_audit.py +++ b/tests/test_dependency_audit.py @@ -55,6 +55,25 @@ 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: {"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 diff --git a/tests/test_dependency_extras.py b/tests/test_dependency_extras.py index a22928b..464fbe6 100644 --- a/tests/test_dependency_extras.py +++ b/tests/test_dependency_extras.py @@ -60,3 +60,17 @@ def test_sglang_requirements_match_the_shared_cute_runtime() -> None: } assert expected <= requirements + + +def test_default_install_declares_the_training_runtime() -> None: + project = tomllib.loads((ROOT / "pyproject.toml").read_text())["project"] + runtime = { + line.strip() + for line in (ROOT / "requirements/sglang-v0.5.21.txt").read_text().splitlines() + if line.strip() and not line.lstrip().startswith("#") + } + # SGLang's upstream pins conflict with the validated Draft kernels; the + # explicit Target setup command prepares it with the declared overrides. + assert runtime - {"sglang==0.5.21"} <= set(project["dependencies"]) + assert not any(requirement.startswith("sglang==") for requirement in project["dependencies"]) + assert "uv==0.13.5" in project["dependencies"] diff --git a/tests/test_open_source_entrypoints.py b/tests/test_open_source_entrypoints.py index f8d5843..deded37 100644 --- a/tests/test_open_source_entrypoints.py +++ b/tests/test_open_source_entrypoints.py @@ -67,6 +67,54 @@ 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_is_prepared_only_for_explicit_apply(monkeypatch, tmp_path, check): + events = [] + + def installed_site(**kwargs): + events.append(("locate", kwargs["writable"])) + if len(events) == 1: + raise importlib.metadata.PackageNotFoundError("sglang") + return tmp_path + + monkeypatch.setattr(patcher, "installed_site", installed_site) + monkeypatch.setattr(patcher, "install_target_runtime", lambda: events.append(("install",))) + monkeypatch.setattr(patcher, "patch_state", lambda _: "patched") + monkeypatch.setattr(sys, "argv", ["patch", *(["--check"] if check else [])]) + if check: + with pytest.raises(SystemExit) as error: + patcher.main() + assert error.value.code == 1 + assert events == [("locate", False)] + else: + patcher.main() + assert events == [("locate", True), ("install",), ("locate", True)] + + +def test_target_install_uses_bundled_overrides_and_current_interpreter(monkeypatch): + calls = [] + monkeypatch.setattr(patcher, "command_environment", lambda: {"PATH": "/verified/bin"}) + monkeypatch.setattr(patcher.shutil, "which", lambda *_, **__: "/verified/bin/git") + monkeypatch.setattr(patcher.subprocess, "run", lambda args, **kwargs: calls.append((args, kwargs))) + patcher.install_target_runtime() + args, kwargs = calls[0] + assert args[:5] == [sys.executable, "-m", "uv", "pip", "install"] + assert args[args.index("--python") + 1] == sys.executable + assert Path(args[args.index("--override") + 1]).name == "cute-overrides.txt" + assert Path(args[args.index("-r") + 1]).name == "sglang-v0.5.21.txt" + assert kwargs == {"env": {"PATH": "/verified/bin"}, "check": True} + + +def test_target_install_rejects_an_invalid_environment_before_installing(monkeypatch): + def invalid(): + raise ValueError("Activate the intended venv") + + monkeypatch.setattr(patcher, "command_environment", invalid) + monkeypatch.setattr(patcher.subprocess, "run", lambda *_, **__: pytest.fail("must not install")) + with pytest.raises(ValueError, match="intended venv"): + patcher.install_target_runtime() + + 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) diff --git a/tools/check_wheel_install.py b/tools/check_wheel_install.py index 2f46761..14ca29e 100644 --- a/tools/check_wheel_install.py +++ b/tools/check_wheel_install.py @@ -86,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(): From 9bcc5d2744f174bfb7bdb8aba405704ce787cd81 Mon Sep 17 00:00:00 2001 From: zhouaihui Date: Thu, 12 Nov 2026 12:13:22 +0800 Subject: [PATCH 03/13] feat(runtime): separate Target and Draft dependency environments --- .github/workflows/build-release.yml | 2 +- aspec/_target_worker.py | 20 +++++++ aspec/commands/convert_to_hf.py | 2 +- aspec/commands/doctor.py | 37 ++++++------ aspec/commands/quickstart.py | 3 + aspec/compat/sglang_patch.py | 28 +-------- aspec/config/inference_config.py | 2 + aspec/evidence.py | 4 +- aspec/inference/engine/target_actor.py | 82 ++++++++++++++++++++++++++ aspec/inference/factory.py | 22 ++++++- aspec/inference/target_runtime.py | 21 +++++++ aspec/training/trainer_actor.py | 11 ++++ aspec/utils/env.py | 1 + pyproject.toml | 38 ++---------- requirements/cute-overrides.txt | 4 -- requirements/draft.txt | 12 ++++ requirements/sglang-v0.5.21.txt | 21 +------ setup.py | 2 +- setup.sh | 13 +++- tests/test_convert_to_hf.py | 2 +- tests/test_dependency_audit.py | 9 ++- tests/test_dependency_extras.py | 35 ++++------- tests/test_open_source_entrypoints.py | 53 +++-------------- tests/test_release_audit.py | 20 ++++++- tests/test_sgl_engine_factory.py | 1 + tests/test_wheel_assets.py | 4 +- tools/audit_release.py | 40 +++++++------ tools/build_release.py | 28 ++++++++- tools/check_wheel_install.py | 2 +- tools/install.sh | 27 ++++++--- 30 files changed, 328 insertions(+), 218 deletions(-) create mode 100644 aspec/_target_worker.py create mode 100644 aspec/inference/engine/target_actor.py create mode 100644 aspec/inference/target_runtime.py delete mode 100644 requirements/cute-overrides.txt create mode 100644 requirements/draft.txt diff --git a/.github/workflows/build-release.yml b/.github/workflows/build-release.yml index 5568df8..52f8c3d 100644 --- a/.github/workflows/build-release.yml +++ b/.github/workflows/build-release.yml @@ -77,7 +77,7 @@ 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: "" diff --git a/aspec/_target_worker.py b/aspec/_target_worker.py new file mode 100644 index 0000000..3a2f586 --- /dev/null +++ b/aspec/_target_worker.py @@ -0,0 +1,20 @@ +# Copyright 2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 +"""Bootstrap a Ray worker with its own venv libraries before native imports.""" +import os +import sys + +from aspec._environment import execute + + +def main(): + os.environ["VIRTUAL_ENV"] = sys.prefix + # The Ray head belongs to Draft. Never load its CUDA libraries in Target. + os.environ.pop("LD_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/commands/convert_to_hf.py b/aspec/commands/convert_to_hf.py index bc42451..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("ling-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 679a9da..0f041ba 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,12 +27,6 @@ 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): @@ -53,19 +48,14 @@ 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") +def expected_versions(target=False): + source = asset_path("requirements/sglang-v0.5.21.txt" if target else "requirements/draft.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)]} @@ -82,8 +72,14 @@ def inspect_environment(*, gpu=False, target=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") + 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().items(): + for name, expected in expected_versions(target).items(): try: actual = importlib.metadata.version(name) except importlib.metadata.PackageNotFoundError: @@ -96,6 +92,13 @@ def inspect_environment(*, gpu=False, target=False): 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") + 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: diff --git a/aspec/commands/quickstart.py b/aspec/commands/quickstart.py index a86b883..253696d 100644 --- a/aspec/commands/quickstart.py +++ b/aspec/commands/quickstart.py @@ -54,10 +54,13 @@ 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-python", help="Python interpreter of the isolated Target venv") 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_python: + config.inference.target_python = args.target_python output = args.output.resolve() if args.plan: with tempfile.TemporaryDirectory(prefix="aspec-plan-") as folder: diff --git a/aspec/compat/sglang_patch.py b/aspec/compat/sglang_patch.py index 2ddd081..9603fc0 100644 --- a/aspec/compat/sglang_patch.py +++ b/aspec/compat/sglang_patch.py @@ -12,19 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Prepare the pinned Target runtime and apply/check its training patch.""" +"""Apply/check the training patch in the active Target environment.""" import argparse import importlib.metadata import importlib.util import os -import shutil import subprocess import sys from pathlib import Path -from aspec._environment import command_environment -from aspec.assets import asset_path VERSION = "0.5.21" PATCH = Path(__file__).with_name("patches") / "sglang-v0.5.21.patch" @@ -92,33 +89,12 @@ def installed_site(*, writable=False): return package.parent -def install_target_runtime(): - """Reuse tools/install.sh's explicit overrides in this isolated environment.""" - env = command_environment() - if shutil.which("git", path=env.get("PATH")) is None: - raise RuntimeError("Install Git before preparing the SGLang training patch") - print(f"Installing the pinned SGLang Target runtime in {sys.prefix}", flush=True) - subprocess.run( - [sys.executable, "-m", "uv", "pip", "install", "--link-mode", "copy", - "--python", sys.executable, - "--override", str(asset_path("requirements/cute-overrides.txt")), - "-r", str(asset_path("requirements/sglang-v0.5.21.txt"))], - env=env, check=True, - ) - - def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--check", action="store_true", help="Require the complete patch; do not change files") args = parser.parse_args() try: - try: - site = installed_site(writable=not args.check) - except importlib.metadata.PackageNotFoundError: - if args.check: - raise - install_target_runtime() - site = installed_site(writable=True) + site = installed_site(writable=not args.check) state = patch_state(site) if state == "unpatched": if args.check: diff --git a/aspec/config/inference_config.py b/aspec/config/inference_config.py index ffacc0d..b651939 100644 --- a/aspec/config/inference_config.py +++ b/aspec/config/inference_config.py @@ -164,6 +164,8 @@ class OfflineTrainingConfig: @dataclass class InferenceConfig: + # Absolute interpreter path on each Target node; independent of Draft's venv. + target_python: Optional[str] = None aux_hidden_states_layers: Optional[list] = None inference_batch_size: int = 1 inference_buffer_threshold: int = 32 diff --git a/aspec/evidence.py b/aspec/evidence.py index 3e2ad1a..e1401d3 100644 --- a/aspec/evidence.py +++ b/aspec/evidence.py @@ -41,6 +41,8 @@ EVIDENCE_SCHEMA_VERSION = "aspec.run-evidence/v0alpha8" _PACKAGE_NAMES = ( "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("ling-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 diff --git a/aspec/inference/engine/target_actor.py b/aspec/inference/engine/target_actor.py new file mode 100644 index 0000000..74194cf --- /dev/null +++ b/aspec/inference/engine/target_actor.py @@ -0,0 +1,82 @@ +# 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 aspec + import ray + import sglang + import torch + + result = self.engine.get_status() + 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()), + "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..b98e91c 100644 --- a/aspec/inference/factory.py +++ b/aspec/inference/factory.py @@ -21,6 +21,9 @@ """Inference engine creation and initialization with Ray placement groups.""" import secrets +import os +import shlex +from pathlib import Path import ray from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy @@ -242,12 +245,24 @@ 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.engine.target_actor import TargetActor + from aspec.inference.target_runtime import code_identity + + target_python = getattr(args, "target_python", None) or os.environ.get("ASPEC_TARGET_PYTHON") + if not target_python or not Path(target_python).is_absolute(): + raise ValueError("Set inference.target_python or ASPEC_TARGET_PYTHON to the absolute " + "Python path of the isolated Target environment") pg_obj, reordered_bundle_indices, reordered_gpu_ids = pg - SglRayActor = ray.remote(SglEngine) + SglRayActor = ray.remote(TargetActor) env_vars = get_aspec_env_vars(args) + # The Target bootstrap reconstructs loader paths from its own installation. + env_vars.pop("LD_LIBRARY_PATH", None) + env_vars.pop("PYTHONPATH", None) + runtime_env = {"env_vars": env_vars, "py_executable": shlex.join( + [target_python, "-m", "aspec._target_worker"])} + expected_runtime = code_identity() # Step 1: Create all engine actors (without calling init yet) engines = [] @@ -267,8 +282,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( + expected_runtime=expected_runtime, args=args, rank=i, base_gpu_id=base_gpu_id, diff --git a/aspec/inference/target_runtime.py b/aspec/inference/target_runtime.py new file mode 100644 index 0000000..7e131d8 --- /dev/null +++ b/aspec/inference/target_runtime.py @@ -0,0 +1,21 @@ +# 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 +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()} diff --git a/aspec/training/trainer_actor.py b/aspec/training/trainer_actor.py index 0460822..52b363f 100644 --- a/aspec/training/trainer_actor.py +++ b/aspec/training/trainer_actor.py @@ -178,6 +178,12 @@ def get_global_step(self) -> int: def get_status(self) -> dict: """Return the rank placement actually selected by Ray and CUDA.""" + import importlib.metadata + import sys + + import aspec + + from aspec.inference.target_runtime import code_identity return { "rank": self._rank, @@ -188,6 +194,11 @@ def get_status(self) -> dict: "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, + "cutlass": importlib.metadata.version("nvidia-cutlass-dsl"), + "quack": importlib.metadata.version("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..55b628b 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", diff --git a/pyproject.toml b/pyproject.toml index b016e7a..1b9f913 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,7 +3,7 @@ requires = ["setuptools>=77.0", "wheel"] build-backend = "setuptools.build_meta" [project] -name = "ling-aspec" +name = "ling-aspec-runtime" version = "0.1.0b0" description = "A Torch-native framework for speculative decoding training" readme = "README.md" @@ -11,38 +11,10 @@ license = "Apache-2.0" license-files = ["LICENSE", "licenses/*"] requires-python = ">=3.11" dependencies = [ - "accelerate", - "apache-tvm-ffi==0.1.11", - "cmake", - "cuda-tile==1.6.0rc5", - "datasets>=5.0.1,<6", - "flash-attn-4==4.0.0b34", - "flashinfer-python[cu13]==0.6.18", - "mooncake-transfer-engine==0.3.13.post1", - "ninja", - "numba", - "numpy<2.4", - "nvidia-cuda-runtime-cu12==12.9.79", - "nvidia-cudnn-frontend==1.31.0", - "nvidia-cutlass-dsl[cu13]==4.7.1", - "omegaconf", - "openai", - "openai-harmony", - "packaging", - "psutil", - "pydantic", - "pyzmq", - "qwen-vl-utils", - "quack-kernels[cu13]==0.6.5", - "ray==2.59.0", - "setuptools", - "sglang-kernel==0.4.7", - "torch==2.13.0", - "torchvision==0.28.0", - "tqdm", - "transformers==5.12.1", - "uv==0.13.5", - "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] 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..40c2624 100644 --- a/requirements/sglang-v0.5.21.txt +++ b/requirements/sglang-v0.5.21.txt @@ -1,23 +1,6 @@ -# 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 +# 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_convert_to_hf.py b/tests/test_convert_to_hf.py index a3a5a67..0823a5d 100644 --- a/tests/test_convert_to_hf.py +++ b/tests/test_convert_to_hf.py @@ -54,7 +54,7 @@ def _model(): def test_convert_checkpoint_writes_strict_loadable_serving_artifact(tmp_path, monkeypatch): - versions = {"ling-aspec": "0.1.0b0", "aspec": "0.0.16"} + 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)) diff --git a/tests/test_dependency_audit.py b/tests/test_dependency_audit.py index e84812f..96e91c1 100644 --- a/tests/test_dependency_audit.py +++ b/tests/test_dependency_audit.py @@ -28,9 +28,8 @@ def test_declared_overrides_remain_visible(): 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", [ @@ -61,7 +60,7 @@ def test_doctor_requires_sglang_only_for_target_nodes(monkeypatch, target): monkeypatch.setattr(doctor, "check_dependency_metadata", lambda: {"conflicts": [], "overrides": []}) monkeypatch.setattr(doctor, "venv_error", lambda: None) - monkeypatch.setattr(doctor, "expected_versions", lambda: {"sglang": "0.5.21"}) + monkeypatch.setattr(doctor, "expected_versions", lambda target=False: {"sglang": "0.5.21"}) def missing(_=None): raise doctor.importlib.metadata.PackageNotFoundError("sglang") @@ -80,7 +79,7 @@ def test_doctor_enforces_only_the_python_minimum(monkeypatch, tmp_path, version) 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 464fbe6..7b9e86c 100644 --- a/tests/test_dependency_extras.py +++ b/tests/test_dependency_extras.py @@ -48,29 +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_default_install_declares_the_training_runtime() -> None: +def test_draft_and_target_dependencies_are_separate(): project = tomllib.loads((ROOT / "pyproject.toml").read_text())["project"] - runtime = { - line.strip() - for line in (ROOT / "requirements/sglang-v0.5.21.txt").read_text().splitlines() - if line.strip() and not line.lstrip().startswith("#") - } - # SGLang's upstream pins conflict with the validated Draft kernels; the - # explicit Target setup command prepares it with the declared overrides. - assert runtime - {"sglang==0.5.21"} <= set(project["dependencies"]) - assert not any(requirement.startswith("sglang==") for requirement in project["dependencies"]) - assert "uv==0.13.5" in project["dependencies"] + 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_open_source_entrypoints.py b/tests/test_open_source_entrypoints.py index deded37..b5d3ba8 100644 --- a/tests/test_open_source_entrypoints.py +++ b/tests/test_open_source_entrypoints.py @@ -68,51 +68,16 @@ def test_unknown_wheel_changes_are_not_overwritten(wheel): @pytest.mark.parametrize("check", [False, True]) -def test_missing_target_runtime_is_prepared_only_for_explicit_apply(monkeypatch, tmp_path, check): - events = [] - - def installed_site(**kwargs): - events.append(("locate", kwargs["writable"])) - if len(events) == 1: - raise importlib.metadata.PackageNotFoundError("sglang") - return tmp_path - - monkeypatch.setattr(patcher, "installed_site", installed_site) - monkeypatch.setattr(patcher, "install_target_runtime", lambda: events.append(("install",))) - monkeypatch.setattr(patcher, "patch_state", lambda _: "patched") +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 [])]) - if check: - with pytest.raises(SystemExit) as error: - patcher.main() - assert error.value.code == 1 - assert events == [("locate", False)] - else: + with pytest.raises(SystemExit) as error: patcher.main() - assert events == [("locate", True), ("install",), ("locate", True)] - - -def test_target_install_uses_bundled_overrides_and_current_interpreter(monkeypatch): - calls = [] - monkeypatch.setattr(patcher, "command_environment", lambda: {"PATH": "/verified/bin"}) - monkeypatch.setattr(patcher.shutil, "which", lambda *_, **__: "/verified/bin/git") - monkeypatch.setattr(patcher.subprocess, "run", lambda args, **kwargs: calls.append((args, kwargs))) - patcher.install_target_runtime() - args, kwargs = calls[0] - assert args[:5] == [sys.executable, "-m", "uv", "pip", "install"] - assert args[args.index("--python") + 1] == sys.executable - assert Path(args[args.index("--override") + 1]).name == "cute-overrides.txt" - assert Path(args[args.index("-r") + 1]).name == "sglang-v0.5.21.txt" - assert kwargs == {"env": {"PATH": "/verified/bin"}, "check": True} - - -def test_target_install_rejects_an_invalid_environment_before_installing(monkeypatch): - def invalid(): - raise ValueError("Activate the intended venv") - - monkeypatch.setattr(patcher, "command_environment", invalid) - monkeypatch.setattr(patcher.subprocess, "run", lambda *_, **__: pytest.fail("must not install")) - with pytest.raises(ValueError, match="intended venv"): - patcher.install_target_runtime() + assert error.value.code == 1 def test_patch_check_works_inside_an_unrelated_git_checkout(wheel): @@ -210,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 ca2cfc6..b7608f3 100644 --- a/tests/test_release_audit.py +++ b/tests/test_release_audit.py @@ -31,7 +31,7 @@ def _artifacts( wheel_version="0.1.0b0", sdist_version="0.1.0b0", sdist_link=False, - distribution="ling-aspec", + distribution="ling-aspec-runtime", ): archive_name = distribution.replace("-", "_") wheel = directory / f"{archive_name}-0.1.0b0-py3-none-any.whl" @@ -61,13 +61,27 @@ def _artifacts( 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", @@ -99,7 +113,7 @@ 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) diff --git a/tests/test_sgl_engine_factory.py b/tests/test_sgl_engine_factory.py index ea0d145..44654a0 100644 --- a/tests/test_sgl_engine_factory.py +++ b/tests/test_sgl_engine_factory.py @@ -48,6 +48,7 @@ def _make_args(**overrides): """Build a minimal args Namespace for _init_sgl_engines.""" defaults = dict( target_model_path="/fake/model", + target_python="/isolated-target/bin/python", inference_num_gpus=8, inference_num_gpus_per_engine=8, inference_num_gpus_per_node=8, diff --git a/tests/test_wheel_assets.py b/tests/test_wheel_assets.py index 260c8b3..1bf17da 100644 --- a/tests/test_wheel_assets.py +++ b/tests/test_wheel_assets.py @@ -58,7 +58,7 @@ def test_installed_command_runner_preserves_cwd_and_arguments(tmp_path, monkeypa def test_package_identity_hashes_bundled_patch_and_does_not_borrow_parent_git(tmp_path, monkeypatch): - versions = {"ling-aspec": "0.1.0b0", "aspec": "0.0.16"} + versions = {"ling-aspec-runtime": "0.1.0b0", "aspec": "0.0.16"} def version(name): if name not in versions: @@ -67,7 +67,7 @@ def version(name): monkeypatch.setattr(importlib.metadata, "version", version) observed = _package_versions() - assert observed["ling-aspec"] == "0.1.0b0" + 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" diff --git a/tools/audit_release.py b/tools/audit_release.py index c2b3696..97cbed4 100644 --- a/tools/audit_release.py +++ b/tools/audit_release.py @@ -101,7 +101,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,8 +138,8 @@ 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 != "ling-aspec": - raise ValueError(f"{artifact.name}: unexpected distribution name {name!r}; expected 'ling-aspec'") + if name not in {"ling-aspec", "ling-aspec-target", "ling-aspec-runtime"}: + raise ValueError(f"{artifact.name}: unexpected distribution name {name!r}") return name, version @@ -155,7 +155,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 @@ -181,10 +183,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("ling_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") @@ -198,18 +200,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..4d334fd 100644 --- a/tools/build_release.py +++ b/tools/build_release.py @@ -12,6 +12,9 @@ """Build and audit wheel/sdist, then collect runtime pins and SHA256 checksums.""" import argparse +import json +import tempfile +import tomllib import hashlib import shutil import subprocess @@ -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 14ca29e..da3767a 100644 --- a/tools/check_wheel_install.py +++ b/tools/check_wheel_install.py @@ -59,7 +59,7 @@ def run(arguments, expected=0, contains=None): identity=_source_identity() assert identity['revision'] is None and identity['dirty'] is None assert identity['source_file_count'] > 0 -assert identity['version'] == importlib.metadata.version('ling-aspec') +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') diff --git a/tools/install.sh b/tools/install.sh index beb6809..daadbea 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 ;; @@ -46,11 +48,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[@]}" From 0520c70a7fd424cc0481a738c8fb356b5d9495ca Mon Sep 17 00:00:00 2001 From: zhouaihui Date: Thu, 12 Nov 2026 12:31:22 +0800 Subject: [PATCH 04/13] fix(runtime): isolate backend imports and record role interpreter identity --- .github/workflows/ci.yml | 1 + Dockerfile | 2 + README.md | 8 +- README_zh-CN.md | 6 +- aspec/_environment.py | 4 + aspec/_target_worker.py | 5 + aspec/cli.py | 23 +++++ aspec/commands/doctor.py | 13 +-- aspec/compat/sglang_patch.py | 5 +- aspec/config/train_config.py | 4 + aspec/inference/engine/__init__.py | 18 ++-- aspec/inference/engine/target_actor.py | 3 +- aspec/inference/factory.py | 7 +- aspec/inference/target_runtime.py | 32 ++++++- aspec/planning.py | 1 + aspec/ray/cluster_preflight.py | 25 ++++- aspec/training/trainer_actor.py | 9 +- docs/framework/DECISIONS.md | 12 ++- docs/releasing.md | 5 +- docs/user-guide/installation-advanced.md | 54 ++++------- docs/user-guide/installation.md | 31 +++++-- tests/test_dependency_audit.py | 4 +- tests/test_target_environment.py | 111 +++++++++++++++++++++++ tools/audit_release.py | 4 + tools/build_release.py | 6 +- 25 files changed, 308 insertions(+), 85 deletions(-) create mode 100644 tests/test_target_environment.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index e1a5480..081e697 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -62,6 +62,7 @@ jobs: tests/test_environment_scripts.py \ tests/test_onboarding.py \ tests/test_dependency_extras.py \ + tests/test_target_environment.py \ tests/test_cudnn_gqa_mode.py \ tests/test_dependency_audit.py \ tests/test_draft_model_config.py \ diff --git a/Dockerfile b/Dockerfile index e676bf0..c977731 100644 --- a/Dockerfile +++ b/Dockerfile @@ -6,6 +6,8 @@ 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 ASPEC_TARGET_PYTHON=/opt/aspec-target/bin/python ENV VIRTUAL_ENV=/opt/aspec-venv ENV PATH="/opt/aspec-venv/bin:${PATH}" CMD ["bash"] diff --git a/README.md b/README.md index 943e925..bb681de 100644 --- a/README.md +++ b/README.md @@ -60,13 +60,17 @@ In an activated Linux virtual environment: pip install ling-aspec ``` -This installs the training dependencies. On nodes that run the Target model, prepare SGLang -and apply the training patch: +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 ``` +In the Draft environment, set `ASPEC_TARGET_PYTHON=/absolute/path/to/target-venv/bin/python` +before training. Target and Draft keep independent dependencies. + 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. diff --git a/README_zh-CN.md b/README_zh-CN.md index 5e15ad0..27a2907 100644 --- a/README_zh-CN.md +++ b/README_zh-CN.md @@ -55,12 +55,16 @@ Qwen、Ling 上的 DSpark/DFlash 是主要优化方向和 recipe 维护重点； pip install ling-aspec ``` -该命令会安装训练依赖。需要运行 Target 模型的节点，再准备 SGLang 并应用训练补丁： +该命令安装 Draft 训练依赖。在独立的 Target 虚拟环境中安装 SGLang 集成并应用补丁： ```bash +pip install ling-aspec-target aspec patch-sglang ``` +Target 使用独立虚拟环境。在 Draft 环境中设置 +`ASPEC_TARGET_PYTHON=/absolute/path/to/target-venv/bin/python` 后启动训练，两边分别维护依赖。 + 环境要求、源码安装和本地 wheel 安装见[安装指南](docs/user-guide/installation.md)， 也可使用[容器](docs/user-guide/containers.md)。可选执行 `aspec doctor` 检查环境。 diff --git a/aspec/_environment.py b/aspec/_environment.py index 00d3206..e33030f 100644 --- a/aspec/_environment.py +++ b/aspec/_environment.py @@ -40,6 +40,10 @@ 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(): + env["CUDA_HOME"] = str(toolkit) + env["PATH"] = str(toolkit / "bin") + os.pathsep + env.get("PATH", "") paths.extend(str(p) for p in (purelib / "nvidia/cuda_runtime/lib", purelib / "torch/lib") if p.is_dir()) if paths: diff --git a/aspec/_target_worker.py b/aspec/_target_worker.py index 3a2f586..77656de 100644 --- a/aspec/_target_worker.py +++ b/aspec/_target_worker.py @@ -3,12 +3,17 @@ """Bootstrap a Ray worker with its own venv libraries before native imports.""" import os import sys +from pathlib import Path from aspec._environment import execute def main(): os.environ["VIRTUAL_ENV"] = sys.prefix + 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 Ray head belongs to Draft. Never load its CUDA libraries in Target. os.environ.pop("LD_LIBRARY_PATH", None) os.environ.pop("PYTHONPATH", None) diff --git a/aspec/cli.py b/aspec/cli.py index dc7515c..f75d28f 100644 --- a/aspec/cli.py +++ b/aspec/cli.py @@ -13,7 +13,10 @@ # limitations under the License. """Workflow shortcuts and native-library bootstrap for the public CLI.""" +import importlib.metadata +import os import sys +from pathlib import Path from aspec._environment import execute @@ -45,6 +48,26 @@ 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") + target_command = module in {"aspec.compat.sglang_patch", "aspec.commands.evaluate_draft", + "aspec.commands.prepare_target_dataset"} + target_command = target_command or (module == "aspec.commands.doctor" and "--target" in arguments) + target_python = os.environ.get("ASPEC_TARGET_PYTHON") + if target_command and not target_role and not help_requested and target_python: + if not Path(target_python).is_absolute(): + raise SystemExit("ASPEC_TARGET_PYTHON must be an absolute Target interpreter path") + execute([target_python, "-m", "aspec._target_worker", "-m", module, *arguments]) + return try: execute([sys.executable, "-m", module, *arguments]) except (ValueError, OSError) as error: diff --git a/aspec/commands/doctor.py b/aspec/commands/doctor.py index 0f041ba..410eeaa 100644 --- a/aspec/commands/doctor.py +++ b/aspec/commands/doctor.py @@ -28,9 +28,8 @@ from aspec.compat.sglang_patch import installed_site, patch_state, venv_error - 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} @@ -56,9 +55,11 @@ def check_dependency_metadata(distributions=None): def expected_versions(target=False): source = asset_path("requirements/sglang-v0.5.21.txt" if target else "requirements/draft.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)]} + 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, target=False): @@ -88,7 +89,7 @@ def inspect_environment(*, gpu=False, target=False): has_sglang = actual is not None if not target and not has_sglang: continue - ok = actual is not None and actual.split("+", 1)[0] == expected + 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") diff --git a/aspec/compat/sglang_patch.py b/aspec/compat/sglang_patch.py index 9603fc0..292f3b9 100644 --- a/aspec/compat/sglang_patch.py +++ b/aspec/compat/sglang_patch.py @@ -22,7 +22,6 @@ import sys from pathlib import Path - VERSION = "0.5.21" PATCH = Path(__file__).with_name("patches") / "sglang-v0.5.21.patch" @@ -105,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 (ValueError, 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/train_config.py b/aspec/config/train_config.py index 59108f4..c7c00a7 100644 --- a/aspec/config/train_config.py +++ b/aspec/config/train_config.py @@ -1006,6 +1006,10 @@ def load_config( configs_to_merge.append(cli_config) config = OmegaConf.merge(*configs_to_merge) + if config.inference.target_python is None: + config.inference.target_python = os.environ.get("ASPEC_TARGET_PYTHON") or None + if config.inference.target_python is not None and not Path(config.inference.target_python).is_absolute(): + raise ValueError("inference.target_python must be an absolute Target interpreter path") _resolve_relative_paths(config, os.getcwd()) _validate_current_scope(config) 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 index 74194cf..ec5e6fd 100644 --- a/aspec/inference/engine/target_actor.py +++ b/aspec/inference/engine/target_actor.py @@ -48,11 +48,12 @@ def find_free_port(self, *args, **kwargs): return self.engine.find_free_port(*args, **kwargs) def get_status(self): - import aspec import ray import sglang import torch + import aspec + result = self.engine.get_status() result["runtime"] = { **self.runtime_identity, diff --git a/aspec/inference/factory.py b/aspec/inference/factory.py index b98e91c..f33ab21 100644 --- a/aspec/inference/factory.py +++ b/aspec/inference/factory.py @@ -20,9 +20,10 @@ """Inference engine creation and initialization with Ray placement groups.""" -import secrets import os +import secrets import shlex +import sys from pathlib import Path import ray @@ -252,6 +253,8 @@ def _prepare_sgl_engines( if not target_python or not Path(target_python).is_absolute(): raise ValueError("Set inference.target_python or ASPEC_TARGET_PYTHON to the absolute " "Python path of the isolated Target environment") + if Path(target_python).parent == Path(sys.executable).parent: + raise ValueError("Target and Draft must use separate virtual environments") pg_obj, reordered_bundle_indices, reordered_gpu_ids = pg @@ -260,6 +263,8 @@ def _prepare_sgl_engines( # The Target bootstrap reconstructs loader paths from its own installation. env_vars.pop("LD_LIBRARY_PATH", None) env_vars.pop("PYTHONPATH", None) + env_vars.pop("CUTE_DSL_CACHE_DIR", None) + env_vars.pop("FLASH_ATTENTION_CUTE_DSL_CACHE_DIR", None) runtime_env = {"env_vars": env_vars, "py_executable": shlex.join( [target_python, "-m", "aspec._target_worker"])} expected_runtime = code_identity() diff --git a/aspec/inference/target_runtime.py b/aspec/inference/target_runtime.py index 7e131d8..b4ebd40 100644 --- a/aspec/inference/target_runtime.py +++ b/aspec/inference/target_runtime.py @@ -3,6 +3,8 @@ """Identity shared across the independently installed Draft and Target roles.""" import hashlib import importlib.metadata +import json +import sys from pathlib import Path @@ -18,4 +20,32 @@ def code_identity(): version = importlib.metadata.version("ling-aspec-runtime") except importlib.metadata.PackageNotFoundError: version = None - return {"version": version, "code_sha256": digest.hexdigest()} + 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/planning.py b/aspec/planning.py index 7515a0e..0faa722 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", + "python_executable": _select(config, "inference.target_python"), }, "locality_domains": locality_domains, }, diff --git a/aspec/ray/cluster_preflight.py b/aspec/ray/cluster_preflight.py index a2570d6..67300bd 100644 --- a/aspec/ray/cluster_preflight.py +++ b/aspec/ray/cluster_preflight.py @@ -283,7 +283,8 @@ 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_python: str | None = None, ) -> dict[str, Any]: marker = Path(marker_path) if storage_mode == "node_local": @@ -310,6 +311,17 @@ 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_python: + from aspec.inference.target_runtime import code_identity + + result = subprocess.run([target_python, "-m", "aspec._target_worker", "-m", + "aspec.inference.target_runtime"], capture_output=True, text=True, + timeout=60, check=True) + target_runtime = json.loads(result.stdout) + expected = code_identity() + if any(target_runtime.get(key) != value for key, value in expected.items()): + 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,9 @@ def _collect_local_inventory( "python": sys.version.split()[0], "source_tree_sha256": _source_tree_sha256(), "packages": _package_versions(), + "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"), @@ -384,6 +397,11 @@ def validate_cluster_inventories( } if len(package_sets) != 1: raise ClusterPreflightError("required package versions differ across nodes") + 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 = set() for node in inventories: artifacts = node.get("runtime_artifacts") or {} @@ -1159,7 +1177,8 @@ def run_cluster_preflight(args) -> dict[str, Any] | None: 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), + getattr(args, "target_python", None) or os.environ.get("ASPEC_TARGET_PYTHON")) ) inventories = ray.get( refs, diff --git a/aspec/training/trainer_actor.py b/aspec/training/trainer_actor.py index 52b363f..2d0ac71 100644 --- a/aspec/training/trainer_actor.py +++ b/aspec/training/trainer_actor.py @@ -178,12 +178,12 @@ def get_global_step(self) -> int: def get_status(self) -> dict: """Return the rank placement actually selected by Ray and CUDA.""" - import importlib.metadata import sys - import aspec + import torch - from aspec.inference.target_runtime import code_identity + import aspec + from aspec.inference.target_runtime import code_identity, dependency_versions return { "rank": self._rank, @@ -197,8 +197,7 @@ def get_status(self) -> dict: "runtime": {**code_identity(), "python": sys.executable, "prefix": sys.prefix, "aspec": aspec.__file__, "ray": ray.__file__, "torch": torch.__version__, "cuda": torch.version.cuda, - "cutlass": importlib.metadata.version("nvidia-cutlass-dsl"), - "quack": importlib.metadata.version("quack-kernels")}, + "dependencies": dependency_versions("nvidia-cutlass-dsl", "quack-kernels")}, } def save_model(self, step: int, force_sync: bool = False) -> None: diff --git a/docs/framework/DECISIONS.md b/docs/framework/DECISIONS.md index c59bbf2..4313813 100644 --- a/docs/framework/DECISIONS.md +++ b/docs/framework/DECISIONS.md @@ -19,12 +19,14 @@ These decisions describe the maintained implementation and its public 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.** The default distribution installs the Draft runtime. Target nodes explicitly - run `aspec patch-sglang`, which prepares a missing SGLang runtime with the documented overrides - and applies the patch. Read-only patch checks never install dependencies. Doctor checks the - Target runtime when present or explicitly requested with `--target`. +- **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. The launcher selects + Target's explicit interpreter and verifies code identity before loading SGLang. This boundary is + inspired by SpecForge's independently provisioned capture-server interpreter, adapted to the + retained Ray actor/transport path rather than introducing another training pipeline. - **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 77ebbf5..cfd861b 100644 --- a/docs/releasing.md +++ b/docs/releasing.md @@ -1,6 +1,7 @@ # Releasing -ASpec is distributed on PyPI as `ling-aspec`. Python imports and the CLI use `aspec`. +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`. ## Versioning @@ -37,7 +38,7 @@ maintenance tools. `setup.py` collects resources from their canonical source loc In the release repository, configure a GitHub environment named `pypi`: -- Add the environment secret `PYPI_API_TOKEN` with a PyPI token authorized for `ling-aspec`. +- 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. diff --git a/docs/user-guide/installation-advanced.md b/docs/user-guide/installation-advanced.md index 6c604fb..4b6e30b 100644 --- a/docs/user-guide/installation-advanced.md +++ b/docs/user-guide/installation-advanced.md @@ -1,16 +1,16 @@ # Advanced installation and native prerequisites Start with the [installation guide](installation.md). This page explains the supported runtime, -native prerequisites, compatibility exceptions and manual recovery procedures. +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 runtime. On Target nodes, -`aspec patch-sglang` also prepares a missing SGLang installation using the bundled dependency -manifest and compatibility overrides before applying the patch. `aspec doctor --target` checks -that setup; `aspec doctor` can check a training environment without SGLang. +`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. ## Requirements @@ -20,9 +20,12 @@ that setup; `aspec doctor` can check a training environment without SGLang. - 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, @@ -60,30 +63,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 together with upstream SGLang. Target setup handles this explicit override separately -from the default pip installation. 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 @@ -114,8 +97,7 @@ reduction, and the same BF16/FP16 operand precision as the eager reference. It r 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 ``` @@ -187,7 +169,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 cf13afa..560469b 100644 --- a/docs/user-guide/installation.md +++ b/docs/user-guide/installation.md @@ -13,21 +13,32 @@ source .venv/bin/activate pip install ling-aspec ``` -This installs ASpec and its training dependencies. The Python package and command are named +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. ## Target nodes -On nodes that run the Target model: +In a separate virtual environment on nodes that run the Target model: ```bash +python -m venv .venv-target +source .venv-target/bin/activate +pip install ling-aspec-target aspec patch-sglang ``` -The command installs the compatible SGLang runtime if it is missing, then applies the training -patch. Runtime versions and compatibility overrides are shipped with ASpec; no manual dependency -selection is needed. See [advanced installation](installation-advanced.md#dedicated-environment) -for the override details. +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. + +In the Draft environment, point training at the Target interpreter: + +```bash +export ASPEC_TARGET_PYTHON=/absolute/path/to/.venv-target/bin/python +``` + +Alternatively set `inference.target_python` in the training config, or pass `--target-python` +to Quickstart. Target and Draft use the same Python minor version and Ray version. On multiple Target nodes, +the configured absolute path must exist on each node. Recheck the patch after reinstalling SGLang: @@ -66,13 +77,15 @@ bash setup.sh source .venv/bin/activate ``` -Setup installs the runtime, applies the patch and runs doctor. Contributors can use +Setup installs the Draft environment. Run `bash setup.sh --target` to prepare `.venv-target`, +then set `ASPEC_TARGET_PYTHON` as above. 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 -Replace `ling-aspec` in the PyPI install command with the path to the downloaded wheel. -Then prepare Target nodes and optionally check the environment as above. The installed commands, +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. See [advanced installation](installation-advanced.md) for native-library paths and manual diff --git a/tests/test_dependency_audit.py b/tests/test_dependency_audit.py index 96e91c1..5736735 100644 --- a/tests/test_dependency_audit.py +++ b/tests/test_dependency_audit.py @@ -23,7 +23,7 @@ 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"), @@ -38,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), ]) diff --git a/tests/test_target_environment.py b/tests/test_target_environment.py new file mode 100644 index 0000000..1b61da8 --- /dev/null +++ b/tests/test_target_environment.py @@ -0,0 +1,111 @@ +# 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 os +import subprocess +import sys +from argparse import Namespace +from unittest.mock import MagicMock, patch + +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 + + +@pytest.mark.parametrize("interpreter", [None, "relative/bin/python"]) +def test_target_interpreter_must_be_explicit_and_absolute(monkeypatch, interpreter): + from aspec.inference.factory import _prepare_sgl_engines + + monkeypatch.delenv("ASPEC_TARGET_PYTHON", raising=False) + args = Namespace(target_python=interpreter, inference_num_gpus=1, + inference_num_gpus_per_engine=1) + with pytest.raises(ValueError, match="absolute"): + _prepare_sgl_engines(args, None) + + +def test_target_actor_selects_its_interpreter_without_draft_loader_paths(monkeypatch): + from aspec.inference.factory import _prepare_sgl_engines + + args = Namespace(target_python="/target env/bin/python", inference_num_gpus=1, + inference_num_gpus_per_engine=1, sglang_nnodes=1) + remote_class = MagicMock() + with (patch("aspec.inference.factory.ray") as ray, + patch("aspec.inference.factory.get_aspec_env_vars", return_value={ + "LD_LIBRARY_PATH": "/draft/lib", "PYTHONPATH": "/draft/source", + "ASPEC_CUDA_COMPAT_DIR": "/driver/compat"})): + ray.remote.return_value = remote_class + ray.get.return_value = 25000 + _prepare_sgl_engines(args, (MagicMock(), [0], [0])) + environment = remote_class.options.call_args.kwargs["runtime_env"] + assert environment["py_executable"] == "'/target env/bin/python' -m aspec._target_worker" + assert environment["env_vars"] == {"ASPEC_CUDA_COMPAT_DIR": "/driver/compat"} + assert "code_sha256" in remote_class.options.return_value.remote.call_args.kwargs["expected_runtime"] + + +def test_target_bootstrap_rebuilds_environment_before_starting_worker(monkeypatch): + from aspec import _target_worker + + for name in ("PYTHONPATH", "PYTHONHOME", "LD_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")) + + +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_use_the_selected_target_interpreter(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][:4] == ["/target/bin/python", "-m", "aspec._target_worker", "-m"] + + +@pytest.mark.parametrize("explicit", [None, "/configured-target/bin/python"]) +def test_target_interpreter_selection_is_recorded_in_plan(monkeypatch, explicit): + 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_python = explicit + plan = compile_execution_plan(load_config(base_config=config)) + expected = explicit or "/environment-target/bin/python" + assert plan.to_runtime_args().target_python == expected + assert plan.to_dict()["topology"]["inference"]["python_executable"] == expected diff --git a/tools/audit_release.py b/tools/audit_release.py index 97cbed4..d5b27e3 100644 --- a/tools/audit_release.py +++ b/tools/audit_release.py @@ -140,6 +140,10 @@ def _identity(raw: bytes, artifact: Path) -> tuple[str, str]: 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 diff --git a/tools/build_release.py b/tools/build_release.py index 4d334fd..ee00aef 100644 --- a/tools/build_release.py +++ b/tools/build_release.py @@ -12,13 +12,13 @@ """Build and audit wheel/sdist, then collect runtime pins and SHA256 checksums.""" import argparse -import json -import tempfile -import tomllib import hashlib +import json import shutil import subprocess import sys +import tempfile +import tomllib from pathlib import Path from tools.audit_release import audit From 39dbe9edcbe8704a8522495e345a3fb56df347f2 Mon Sep 17 00:00:00 2001 From: zhouaihui Date: Thu, 12 Nov 2026 12:43:03 +0800 Subject: [PATCH 05/13] fix(install): expose pip CUDA libraries to Target JIT linker --- aspec/_environment.py | 45 ++++++++++++++++++++++-- aspec/_target_worker.py | 4 +++ docs/user-guide/installation-advanced.md | 5 +-- requirements/sglang-v0.5.21.txt | 1 + tests/test_environment_scripts.py | 26 ++++++++++++++ tests/test_target_environment.py | 6 ++-- 6 files changed, 80 insertions(+), 7 deletions(-) diff --git a/aspec/_environment.py b/aspec/_environment.py index e33030f..46ac32d 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", 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.") @@ -42,10 +76,17 @@ def command_environment(): 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(): - env["CUDA_HOME"] = str(toolkit) - env["PATH"] = str(toolkit / "bin") + os.pathsep + env.get("PATH", "") + 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 index 77656de..a67ba7c 100644 --- a/aspec/_target_worker.py +++ b/aspec/_target_worker.py @@ -10,12 +10,16 @@ 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 Ray head 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:]]) diff --git a/docs/user-guide/installation-advanced.md b/docs/user-guide/installation-advanced.md index 4b6e30b..fa2b32f 100644 --- a/docs/user-guide/installation-advanced.md +++ b/docs/user-guide/installation-advanced.md @@ -48,8 +48,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. diff --git a/requirements/sglang-v0.5.21.txt b/requirements/sglang-v0.5.21.txt index 40c2624..eb37893 100644 --- a/requirements/sglang-v0.5.21.txt +++ b/requirements/sglang-v0.5.21.txt @@ -1,6 +1,7 @@ # Target-only environment: retain SGLang's declared dependency versions. sglang==0.5.21 flash-attn-4==4.0.0b34 +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 nvidia-cuda-runtime-cu12==12.9.79 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_target_environment.py b/tests/test_target_environment.py index 1b61da8..ef3815c 100644 --- a/tests/test_target_environment.py +++ b/tests/test_target_environment.py @@ -57,7 +57,7 @@ def test_target_actor_selects_its_interpreter_without_draft_loader_paths(monkeyp def test_target_bootstrap_rebuilds_environment_before_starting_worker(monkeypatch): from aspec import _target_worker - for name in ("PYTHONPATH", "PYTHONHOME", "LD_LIBRARY_PATH"): + 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") @@ -70,7 +70,7 @@ def test_target_bootstrap_rebuilds_environment_before_starting_worker(monkeypatc 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")) + 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): @@ -108,4 +108,4 @@ def test_target_interpreter_selection_is_recorded_in_plan(monkeypatch, explicit) plan = compile_execution_plan(load_config(base_config=config)) expected = explicit or "/environment-target/bin/python" assert plan.to_runtime_args().target_python == expected - assert plan.to_dict()["topology"]["inference"]["python_executable"] == expected + assert plan.to_dict()["policy"]["topology"]["inference"]["python_executable"] == expected From aa4b8777287436ab77274c875e06a65449873c11 Mon Sep 17 00:00:00 2001 From: zhouaihui Date: Thu, 12 Nov 2026 12:49:33 +0800 Subject: [PATCH 06/13] fix(evidence): collect SGLang provenance from isolated Target workers --- CONTRIBUTING.md | 6 ++- aspec/_environment.py | 2 +- aspec/evidence.py | 12 ++++++ aspec/inference/engine/target_actor.py | 9 ++++ docs/user-guide/installation-advanced.md | 2 +- docs/user-guide/installation.md | 3 +- requirements/ci-cpu.txt | 2 +- tests/test_evidence.py | 53 +++++++++++++++++++++++- tools/install.sh | 8 ---- 9 files changed, 83 insertions(+), 14 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 2b081a3..9f4876b 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -21,11 +21,15 @@ 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 diff --git a/aspec/_environment.py b/aspec/_environment.py index 46ac32d..c580d42 100644 --- a/aspec/_environment.py +++ b/aspec/_environment.py @@ -28,7 +28,7 @@ def _toolkit_view(toolkit, env): the installed wheels and host toolkit remain untouched. """ cache = Path(env.get("ASPEC_CACHE_HOME") or - Path(env.get("XDG_CACHE_HOME", Path.home() / ".cache")) / "aspec") + 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" diff --git a/aspec/evidence.py b/aspec/evidence.py index e1401d3..ccfc32e 100644 --- a/aspec/evidence.py +++ b/aspec/evidence.py @@ -305,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(), @@ -429,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/target_actor.py b/aspec/inference/engine/target_actor.py index ec5e6fd..f904827 100644 --- a/aspec/inference/engine/target_actor.py +++ b/aspec/inference/engine/target_actor.py @@ -53,7 +53,14 @@ def get_status(self): 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["runtime"] = { **self.runtime_identity, @@ -61,6 +68,8 @@ def get_status(self): "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"), diff --git a/docs/user-guide/installation-advanced.md b/docs/user-guide/installation-advanced.md index fa2b32f..74c5948 100644 --- a/docs/user-guide/installation-advanced.md +++ b/docs/user-guide/installation-advanced.md @@ -30,7 +30,7 @@ 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 +matching wheel is available; pip 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 compatibility; frontend 1.32 requires CUTLASS >=4.8. diff --git a/docs/user-guide/installation.md b/docs/user-guide/installation.md index 560469b..81c8232 100644 --- a/docs/user-guide/installation.md +++ b/docs/user-guide/installation.md @@ -1,7 +1,8 @@ # Installation Use Linux x86-64, Python >=3.11 and an isolated virtual environment. Training requires NVIDIA -GPUs, a CUDA-compatible driver and a CUDA development toolkit with `nvcc`. +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). 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/tests/test_evidence.py b/tests/test_evidence.py index f213a15..59bc18b 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() @@ -731,6 +777,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 +934,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/tools/install.sh b/tools/install.sh index daadbea..0a088ab 100755 --- a/tools/install.sh +++ b/tools/install.sh @@ -21,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) From 587bf0526ec41b803f9548d25f5427d68d067d05 Mon Sep 17 00:00:00 2001 From: zhouaihui Date: Thu, 12 Nov 2026 13:04:55 +0800 Subject: [PATCH 07/13] test(release): require split runtime assets and complete execution evidence --- CHANGELOG.md | 8 ++++++-- docs/user-guide/data.md | 3 +++ docs/user-guide/evaluation.md | 3 +++ docs/user-guide/installation.md | 2 +- tests/test_evidence.py | 1 - tests/test_release_audit.py | 4 +++- tests/test_target_environment.py | 2 +- tools/audit_release.py | 4 ++++ 8 files changed, 21 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f81ed0d..7e0b833 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,8 +2,12 @@ ## Unreleased -- Install training dependencies with the default distribution. Target nodes can prepare SGLang - and its training patch with `aspec patch-sglang`; `aspec doctor --target` checks that runtime. +- 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. Select Target's interpreter with + `ASPEC_TARGET_PYTHON` or `inference.target_python`. +- 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 diff --git a/docs/user-guide/data.md b/docs/user-guide/data.md index 9ba8b00..8897c73 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. +`aspec prepare-target-data` uses the Target environment selected by `ASPEC_TARGET_PYTHON`. +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..d24cf1d 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. +Serving-based `aspec evaluate` runs in the Target environment selected by `ASPEC_TARGET_PYTHON`. +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.md b/docs/user-guide/installation.md index 81c8232..0b8abc0 100644 --- a/docs/user-guide/installation.md +++ b/docs/user-guide/installation.md @@ -48,7 +48,7 @@ aspec patch-sglang --check ``` Applying the patch again verifies an already patched installation. Conflicting changes or an -unsupported SGLang version are rejected. `--check` never installs dependencies or modifies files. +unsupported SGLang version are rejected. `--check` never installs dependencies or modifies SGLang files. ## Check the environment diff --git a/tests/test_evidence.py b/tests/test_evidence.py index 59bc18b..997f9ab 100644 --- a/tests/test_evidence.py +++ b/tests/test_evidence.py @@ -736,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", diff --git a/tests/test_release_audit.py b/tests/test_release_audit.py index b7608f3..e517d25 100644 --- a/tests/test_release_audit.py +++ b/tests/test_release_audit.py @@ -85,7 +85,9 @@ def test_release_audit_accepts_complete_artifact_pair(tmp_path): @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"]) 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"): diff --git a/tests/test_target_environment.py b/tests/test_target_environment.py index ef3815c..6141b3b 100644 --- a/tests/test_target_environment.py +++ b/tests/test_target_environment.py @@ -108,4 +108,4 @@ def test_target_interpreter_selection_is_recorded_in_plan(monkeypatch, explicit) plan = compile_execution_plan(load_config(base_config=config)) expected = explicit or "/environment-target/bin/python" assert plan.to_runtime_args().target_python == expected - assert plan.to_dict()["policy"]["topology"]["inference"]["python_executable"] == expected + assert plan.to_dict()["policy"]["resources"]["inference"]["python_executable"] == expected diff --git a/tools/audit_release.py b/tools/audit_release.py index d5b27e3..a477750 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,9 @@ "setup.py", } WHEEL_REQUIRED = { + "aspec/_target_worker.py", + "aspec/inference/target_runtime.py", + "aspec/inference/engine/target_actor.py", "aspec/compat/fa4_cutlass.py", "aspec/compat/sglang_patch.py", "aspec/compat/patches/sglang-v0.5.21.patch", From 06c0018fe8a662dfc0b103e8f1dc8cf4f7727f3a Mon Sep 17 00:00:00 2001 From: zhouaihui Date: Thu, 12 Nov 2026 13:12:08 +0800 Subject: [PATCH 08/13] fix(runtime): validate cluster dependencies within Target and Draft roles --- aspec/inference/factory.py | 3 +- aspec/ray/cluster_preflight.py | 49 ++++++++++++++++++--------------- aspec/utils/env.py | 2 +- tests/test_cluster_preflight.py | 17 ++++++++++++ 4 files changed, 47 insertions(+), 24 deletions(-) diff --git a/aspec/inference/factory.py b/aspec/inference/factory.py index f33ab21..2d13b49 100644 --- a/aspec/inference/factory.py +++ b/aspec/inference/factory.py @@ -253,7 +253,8 @@ def _prepare_sgl_engines( if not target_python or not Path(target_python).is_absolute(): raise ValueError("Set inference.target_python or ASPEC_TARGET_PYTHON to the absolute " "Python path of the isolated Target environment") - if Path(target_python).parent == Path(sys.executable).parent: + if (getattr(args, "cluster_mode", "single_node") == "single_node" + and Path(target_python).parent == Path(sys.executable).parent): raise ValueError("Target and Draft must use separate virtual environments") pg_obj, reordered_bundle_indices, reordered_gpu_ids = pg diff --git a/aspec/ray/cluster_preflight.py b/aspec/ray/cluster_preflight.py index 67300bd..94577aa 100644 --- a/aspec/ray/cluster_preflight.py +++ b/aspec/ray/cluster_preflight.py @@ -285,6 +285,7 @@ def _collect_local_inventory( required_paths: Sequence[str], marker_path: str, marker_sha256: str, storage_mode: str = "shared", output_dir: str | None = None, target_python: str | None = None, + roles: Sequence[str] | None = None, ) -> dict[str, Any]: marker = Path(marker_path) if storage_mode == "node_local": @@ -329,6 +330,7 @@ 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": target_runtime["sglang"] if target_runtime else module_runtime_identity( @@ -392,22 +394,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") + 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 = set() + 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") @@ -417,21 +428,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" ) @@ -1171,6 +1171,11 @@ 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") + target_python = ((getattr(args, "target_python", None) or os.environ.get("ASPEC_TARGET_PYTHON")) + if "target" in roles else None) refs.append( probe.options( scheduling_strategy=NodeAffinitySchedulingStrategy( @@ -1178,7 +1183,7 @@ def run_cluster_preflight(args) -> dict[str, Any] | None: ) ).remote(_required_paths(args), str(marker), marker_sha256, getattr(args, "cluster_storage_mode", "shared"), str(output_dir), - getattr(args, "target_python", None) or os.environ.get("ASPEC_TARGET_PYTHON")) + target_python, roles) ) inventories = ray.get( refs, diff --git a/aspec/utils/env.py b/aspec/utils/env.py index 55b628b..9e00c0c 100644 --- a/aspec/utils/env.py +++ b/aspec/utils/env.py @@ -84,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/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"], From e1e6e7e1a6326ef4604b48b47ecc9898c028f4b6 Mon Sep 17 00:00:00 2001 From: zhouaihui Date: Thu, 12 Nov 2026 15:08:40 +0800 Subject: [PATCH 09/13] feat(runtime): launch Target independently with a shared service interface --- .github/workflows/ci.yml | 1 + CHANGELOG.md | 5 +- Dockerfile | 1 - README.md | 6 +- README_zh-CN.md | 5 +- aspec/_target_worker.py | 4 +- aspec/cli.py | 14 +- aspec/commands/quickstart.py | 6 +- aspec/commands/recipes.py | 20 +- aspec/config/inference_config.py | 12 +- aspec/config/train_config.py | 19 +- aspec/controller/loop.py | 5 + aspec/inference/factory.py | 40 ++- aspec/inference/managed.py | 81 ++++++ aspec/inference/service.py | 323 +++++++++++++++++++++++ aspec/inference/service_client.py | 91 +++++++ aspec/inference/service_protocol.py | 113 ++++++++ aspec/multi_draft.py | 3 + aspec/offline/generate.py | 7 +- aspec/offline/saving_actor.py | 3 + aspec/planning.py | 2 +- aspec/ray/cluster_preflight.py | 20 +- aspec/ray/placement_group.py | 9 +- aspec/ray/ray_actor.py | 6 + aspec/train_entry.py | 10 + aspec/training/trainer_actor.py | 2 + docs/framework/ARCHITECTURE.md | 7 + docs/framework/DECISIONS.md | 11 +- docs/user-guide/containers.md | 5 +- docs/user-guide/data.md | 2 +- docs/user-guide/evaluation.md | 2 +- docs/user-guide/installation-advanced.md | 10 +- docs/user-guide/installation.md | 31 ++- tests/test_release_audit.py | 6 +- tests/test_sgl_engine_factory.py | 1 - tests/test_target_environment.py | 51 +--- tests/test_target_service.py | 180 +++++++++++++ tools/audit_release.py | 4 + 38 files changed, 985 insertions(+), 133 deletions(-) create mode 100644 aspec/inference/managed.py create mode 100644 aspec/inference/service.py create mode 100644 aspec/inference/service_client.py create mode 100644 aspec/inference/service_protocol.py create mode 100644 tests/test_target_service.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 081e697..70d31c6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -63,6 +63,7 @@ jobs: 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 7e0b833..5c0fd71 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,8 +4,9 @@ - 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. Select Target's interpreter with - `ASPEC_TARGET_PYTHON` or `inference.target_python`. + 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` diff --git a/Dockerfile b/Dockerfile index c977731..7b9232c 100644 --- a/Dockerfile +++ b/Dockerfile @@ -7,7 +7,6 @@ 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 ASPEC_TARGET_PYTHON=/opt/aspec-target/bin/python ENV VIRTUAL_ENV=/opt/aspec-venv ENV PATH="/opt/aspec-venv/bin:${PATH}" CMD ["bash"] diff --git a/README.md b/README.md index bb681de..8bccd20 100644 --- a/README.md +++ b/README.md @@ -66,10 +66,12 @@ the SGLang integration and apply its patch: ```bash pip install ling-aspec-target aspec patch-sglang +aspec target ``` -In the Draft environment, set `ASPEC_TARGET_PYTHON=/absolute/path/to/target-venv/bin/python` -before training. Target and Draft keep independent dependencies. +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). diff --git a/README_zh-CN.md b/README_zh-CN.md index 27a2907..d282f49 100644 --- a/README_zh-CN.md +++ b/README_zh-CN.md @@ -60,10 +60,11 @@ pip install ling-aspec ```bash pip install ling-aspec-target aspec patch-sglang +aspec target ``` -Target 使用独立虚拟环境。在 Draft 环境中设置 -`ASPEC_TARGET_PYTHON=/absolute/path/to/target-venv/bin/python` 后启动训练，两边分别维护依赖。 +随后在 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` 检查环境。 diff --git a/aspec/_target_worker.py b/aspec/_target_worker.py index a67ba7c..5655ad4 100644 --- a/aspec/_target_worker.py +++ b/aspec/_target_worker.py @@ -1,6 +1,6 @@ # Copyright 2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 -"""Bootstrap a Ray worker with its own venv libraries before native imports.""" +"""Bootstrap the optional managed-local Target with its own native libraries.""" import os import sys from pathlib import Path @@ -17,7 +17,7 @@ def main(): if cache := os.environ.get("ASPEC_CACHE_HOME"): if Path(cache).name != "target": os.environ["ASPEC_CACHE_HOME"] = str(Path(cache) / "target") - # The Ray head belongs to Draft. Never load its CUDA libraries in 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) diff --git a/aspec/cli.py b/aspec/cli.py index f75d28f..d6c2c59 100644 --- a/aspec/cli.py +++ b/aspec/cli.py @@ -14,7 +14,6 @@ """Workflow shortcuts and native-library bootstrap for the public CLI.""" import importlib.metadata -import os import sys from pathlib import Path @@ -24,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", @@ -59,15 +60,8 @@ def main(): 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") - target_command = module in {"aspec.compat.sglang_patch", "aspec.commands.evaluate_draft", - "aspec.commands.prepare_target_dataset"} - target_command = target_command or (module == "aspec.commands.doctor" and "--target" in arguments) - target_python = os.environ.get("ASPEC_TARGET_PYTHON") - if target_command and not target_role and not help_requested and target_python: - if not Path(target_python).is_absolute(): - raise SystemExit("ASPEC_TARGET_PYTHON must be an absolute Target interpreter path") - execute([target_python, "-m", "aspec._target_worker", "-m", module, *arguments]) - return + 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/quickstart.py b/aspec/commands/quickstart.py index 253696d..d622551 100644 --- a/aspec/commands/quickstart.py +++ b/aspec/commands/quickstart.py @@ -54,13 +54,13 @@ 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-python", help="Python interpreter of the isolated Target venv") + parser.add_argument("--target-address", help="Address of an already running Target service") 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_python: - config.inference.target_python = args.target_python + if args.target_address: + config.inference.target_service = {"address": args.target_address} 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/config/inference_config.py b/aspec/config/inference_config.py index b651939..1f8a7b0 100644 --- a/aspec/config/inference_config.py +++ b/aspec/config/inference_config.py @@ -162,10 +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: - # Absolute interpreter path on each Target node; independent of Draft's venv. - target_python: Optional[str] = None + 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 c7c00a7..501558b 100644 --- a/aspec/config/train_config.py +++ b/aspec/config/train_config.py @@ -1006,10 +1006,21 @@ def load_config( configs_to_merge.append(cli_config) config = OmegaConf.merge(*configs_to_merge) - if config.inference.target_python is None: - config.inference.target_python = os.environ.get("ASPEC_TARGET_PYTHON") or None - if config.inference.target_python is not None and not Path(config.inference.target_python).is_absolute(): - raise ValueError("inference.target_python must be an absolute Target interpreter path") + 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/inference/factory.py b/aspec/inference/factory.py index 2d13b49..9203b91 100644 --- a/aspec/inference/factory.py +++ b/aspec/inference/factory.py @@ -20,11 +20,7 @@ """Inference engine creation and initialization with Ray placement groups.""" -import os import secrets -import shlex -import sys -from pathlib import Path import ray from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy @@ -246,29 +242,25 @@ def _prepare_sgl_engines( f"({gpus_per_engine} GPU(s) each, nnodes={nnodes}, replicas={num_replicas})" ) - from aspec.inference.engine.target_actor import TargetActor - from aspec.inference.target_runtime import code_identity - - target_python = getattr(args, "target_python", None) or os.environ.get("ASPEC_TARGET_PYTHON") - if not target_python or not Path(target_python).is_absolute(): - raise ValueError("Set inference.target_python or ASPEC_TARGET_PYTHON to the absolute " - "Python path of the isolated Target environment") - if (getattr(args, "cluster_mode", "single_node") == "single_node" - and Path(target_python).parent == Path(sys.executable).parent): - raise ValueError("Target and Draft must use separate virtual environments") + from aspec.inference.service_client import TargetServiceClient pg_obj, reordered_bundle_indices, reordered_gpu_ids = pg - SglRayActor = ray.remote(TargetActor) + SglRayActor = ray.remote(TargetServiceClient) env_vars = get_aspec_env_vars(args) - # The Target bootstrap reconstructs loader paths from its own installation. - env_vars.pop("LD_LIBRARY_PATH", None) - env_vars.pop("PYTHONPATH", None) - env_vars.pop("CUTE_DSL_CACHE_DIR", None) - env_vars.pop("FLASH_ATTENTION_CUTE_DSL_CACHE_DIR", None) - runtime_env = {"env_vars": env_vars, "py_executable": shlex.join( - [target_python, "-m", "aspec._target_worker"])} - expected_runtime = code_identity() + 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 = [] @@ -290,7 +282,7 @@ def _prepare_sgl_engines( scheduling_strategy=scheduling_strategy, runtime_env=runtime_env, ).remote( - expected_runtime=expected_runtime, + 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..dd2f97f --- /dev/null +++ b/aspec/inference/service.py @@ -0,0 +1,323 @@ +# 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) + if not isinstance(session, str) or not session: + raise ValueError("Target requires an explicit attempt identity") + 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. + actor = ray.remote(TargetActor).options( + num_cpus=0.2, num_gpus=0, runtime_env={"env_vars": get_aspec_env_vars(args)}, + ).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": 0} + 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 + 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 + 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: + 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) + self.call(key, "shutdown", [], {}) + + def idle(self): + with self.lock: + return not self.claims + + +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 retained while reader exit is unverified") + 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/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 0faa722..a4b61d4 100644 --- a/aspec/planning.py +++ b/aspec/planning.py @@ -1589,7 +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", - "python_executable": _select(config, "inference.target_python"), + "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 94577aa..4860c92 100644 --- a/aspec/ray/cluster_preflight.py +++ b/aspec/ray/cluster_preflight.py @@ -284,7 +284,7 @@ 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, - target_python: str | None = None, + target_address: str | None = None, roles: Sequence[str] | None = None, ) -> dict[str, Any]: marker = Path(marker_path) @@ -313,15 +313,14 @@ def _collect_local_inventory( 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_python: + if target_address: + from aspec.inference.service_protocol import request from aspec.inference.target_runtime import code_identity - result = subprocess.run([target_python, "-m", "aspec._target_worker", "-m", - "aspec.inference.target_runtime"], capture_output=True, text=True, - timeout=60, check=True) - target_runtime = json.loads(result.stdout) + info = request(target_address, "/v1/info", timeout=30) + target_runtime = info["runtime"] expected = code_identity() - if any(target_runtime.get(key) != value for key, value in expected.items()): + if info["identity"] != expected: raise ClusterPreflightError("Target and Draft ASpec code identity differs") return { "node_id": ray.get_runtime_context().get_node_id(), @@ -1174,8 +1173,9 @@ def run_cluster_preflight(args) -> dict[str, Any] | None: 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") - target_python = ((getattr(args, "target_python", None) or os.environ.get("ASPEC_TARGET_PYTHON")) - if "target" in roles else None) + 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( @@ -1183,7 +1183,7 @@ def run_cluster_preflight(args) -> dict[str, Any] | None: ) ).remote(_required_paths(args), str(marker), marker_sha256, getattr(args, "cluster_storage_mode", "shared"), str(output_dir), - target_python, roles) + 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 2d0ac71..2587568 100644 --- a/aspec/training/trainer_actor.py +++ b/aspec/training/trainer_actor.py @@ -180,6 +180,7 @@ def get_status(self) -> dict: """Return the rank placement actually selected by Ray and CUDA.""" import sys + import psutil import torch import aspec @@ -191,6 +192,7 @@ 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"), 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 4313813..62bc057 100644 --- a/docs/framework/DECISIONS.md +++ b/docs/framework/DECISIONS.md @@ -23,10 +23,13 @@ These decisions describe the maintained implementation and its public behavior. 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. The launcher selects - Target's explicit interpreter and verifies code identity before loading SGLang. This boundary is - inspired by SpecForge's independently provisioned capture-server interpreter, adapted to the - retained Ray actor/transport path rather than introducing another training pipeline. + 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/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 8897c73..dc6f96a 100644 --- a/docs/user-guide/data.md +++ b/docs/user-guide/data.md @@ -3,7 +3,7 @@ Target generation, training rows and decode prefixes have different boundaries. Bundled synthetic data is a smoke fixture, not a quality recipe. -`aspec prepare-target-data` uses the Target environment selected by `ASPEC_TARGET_PYTHON`. +Run `aspec prepare-target-data` in the Target environment. Other local preprocessing commands run in Draft. ## Conversations diff --git a/docs/user-guide/evaluation.md b/docs/user-guide/evaluation.md index d24cf1d..8053853 100644 --- a/docs/user-guide/evaluation.md +++ b/docs/user-guide/evaluation.md @@ -2,7 +2,7 @@ Training health, held-out learning, asset compatibility and decode usefulness are separate questions. -Serving-based `aspec evaluate` runs in the Target environment selected by `ASPEC_TARGET_PYTHON`. +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 diff --git a/docs/user-guide/installation-advanced.md b/docs/user-guide/installation-advanced.md index 74c5948..7c877ea 100644 --- a/docs/user-guide/installation-advanced.md +++ b/docs/user-guide/installation-advanced.md @@ -11,6 +11,12 @@ run from the repository root. Both installation forms use the same role dependen `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 @@ -31,7 +37,7 @@ 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; pip 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 +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 @@ -92,7 +98,7 @@ 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: diff --git a/docs/user-guide/installation.md b/docs/user-guide/installation.md index 0b8abc0..c466903 100644 --- a/docs/user-guide/installation.md +++ b/docs/user-guide/installation.md @@ -26,20 +26,23 @@ python -m venv .venv-target source .venv-target/bin/activate pip install ling-aspec-target aspec patch-sglang +aspec target ``` 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. -In the Draft environment, point training at the Target interpreter: +Keep this service running. In the Draft environment, start training: ```bash -export ASPEC_TARGET_PYTHON=/absolute/path/to/.venv-target/bin/python +aspec train --config /path/to/run.yaml ``` -Alternatively set `inference.target_python` in the training config, or pass `--target-python` -to Quickstart. Target and Draft use the same Python minor version and Ray version. On multiple Target nodes, -the configured absolute path must exist on each node. +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. Recheck the patch after reinstalling SGLang: @@ -79,7 +82,7 @@ source .venv/bin/activate ``` Setup installs the Draft environment. Run `bash setup.sh --target` to prepare `.venv-target`, -then set `ASPEC_TARGET_PYTHON` as above. Each setup runs the appropriate doctor. Contributors can use +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 @@ -91,3 +94,19 @@ recipes and Quickstart data work outside the source checkout. 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 +``` + +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. diff --git a/tests/test_release_audit.py b/tests/test_release_audit.py index e517d25..0d79272 100644 --- a/tests/test_release_audit.py +++ b/tests/test_release_audit.py @@ -87,7 +87,11 @@ def test_release_audit_accepts_complete_artifact_pair(tmp_path): @pytest.mark.parametrize("member", ["aspec/compat/sglang_patch.py", "aspec/compat/patches/sglang-v0.5.21.patch", "aspec/_target_worker.py", - "aspec/inference/engine/target_actor.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"): diff --git a/tests/test_sgl_engine_factory.py b/tests/test_sgl_engine_factory.py index 44654a0..ea0d145 100644 --- a/tests/test_sgl_engine_factory.py +++ b/tests/test_sgl_engine_factory.py @@ -48,7 +48,6 @@ def _make_args(**overrides): """Build a minimal args Namespace for _init_sgl_engines.""" defaults = dict( target_model_path="/fake/model", - target_python="/isolated-target/bin/python", inference_num_gpus=8, inference_num_gpus_per_engine=8, inference_num_gpus_per_node=8, diff --git a/tests/test_target_environment.py b/tests/test_target_environment.py index 6141b3b..c8c3aea 100644 --- a/tests/test_target_environment.py +++ b/tests/test_target_environment.py @@ -1,11 +1,8 @@ # 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 os import subprocess import sys -from argparse import Namespace -from unittest.mock import MagicMock, patch import pytest @@ -24,37 +21,9 @@ def find_spec(self, fullname, *args): assert result.returncode == 0, result.stderr -@pytest.mark.parametrize("interpreter", [None, "relative/bin/python"]) -def test_target_interpreter_must_be_explicit_and_absolute(monkeypatch, interpreter): - from aspec.inference.factory import _prepare_sgl_engines - - monkeypatch.delenv("ASPEC_TARGET_PYTHON", raising=False) - args = Namespace(target_python=interpreter, inference_num_gpus=1, - inference_num_gpus_per_engine=1) - with pytest.raises(ValueError, match="absolute"): - _prepare_sgl_engines(args, None) - - -def test_target_actor_selects_its_interpreter_without_draft_loader_paths(monkeypatch): - from aspec.inference.factory import _prepare_sgl_engines - - args = Namespace(target_python="/target env/bin/python", inference_num_gpus=1, - inference_num_gpus_per_engine=1, sglang_nnodes=1) - remote_class = MagicMock() - with (patch("aspec.inference.factory.ray") as ray, - patch("aspec.inference.factory.get_aspec_env_vars", return_value={ - "LD_LIBRARY_PATH": "/draft/lib", "PYTHONPATH": "/draft/source", - "ASPEC_CUDA_COMPAT_DIR": "/driver/compat"})): - ray.remote.return_value = remote_class - ray.get.return_value = 25000 - _prepare_sgl_engines(args, (MagicMock(), [0], [0])) - environment = remote_class.options.call_args.kwargs["runtime_env"] - assert environment["py_executable"] == "'/target env/bin/python' -m aspec._target_worker" - assert environment["env_vars"] == {"ASPEC_CUDA_COMPAT_DIR": "/driver/compat"} - assert "code_sha256" in remote_class.options.return_value.remote.call_args.kwargs["expected_runtime"] - - 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"): @@ -81,7 +50,7 @@ def test_target_rejects_different_code_before_loading_engine(monkeypatch): @pytest.mark.parametrize("command", ["evaluate", "prepare-target-data", "patch-sglang"]) -def test_target_cli_commands_use_the_selected_target_interpreter(monkeypatch, command): +def test_target_cli_commands_run_in_the_active_environment(monkeypatch, command): from aspec import cli def missing(_): @@ -93,19 +62,19 @@ def missing(_): calls = [] monkeypatch.setattr(cli, "execute", calls.append) cli.main() - assert calls[0][:4] == ["/target/bin/python", "-m", "aspec._target_worker", "-m"] + assert calls[0][:2] == [sys.executable, "-m"] -@pytest.mark.parametrize("explicit", [None, "/configured-target/bin/python"]) -def test_target_interpreter_selection_is_recorded_in_plan(monkeypatch, explicit): +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_python = explicit + config.inference.target_service = {"address": "http://target:30000"} plan = compile_execution_plan(load_config(base_config=config)) - expected = explicit or "/environment-target/bin/python" - assert plan.to_runtime_args().target_python == expected - assert plan.to_dict()["policy"]["resources"]["inference"]["python_executable"] == expected + 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..a850e62 --- /dev/null +++ b/tests/test_target_service.py @@ -0,0 +1,180 @@ +# 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 + monkeypatch.setattr(module, "ray", SimpleNamespace(remote=lambda _: actor, + get=lambda x, **_: x, kill=MagicMock())) + 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"}, + "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_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_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/tools/audit_release.py b/tools/audit_release.py index a477750..663a75d 100644 --- a/tools/audit_release.py +++ b/tools/audit_release.py @@ -85,6 +85,10 @@ "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", From bf2f9d52bce21e116ac23017e6bc98cd4adae080 Mon Sep 17 00:00:00 2001 From: zhouaihui Date: Thu, 12 Nov 2026 15:13:52 +0800 Subject: [PATCH 10/13] feat(cli): expose optional local supervision in Quickstart --- aspec/commands/quickstart.py | 9 +++++++++ docs/user-guide/installation.md | 1 + 2 files changed, 10 insertions(+) diff --git a/aspec/commands/quickstart.py b/aspec/commands/quickstart.py index d622551..536856f 100644 --- a/aspec/commands/quickstart.py +++ b/aspec/commands/quickstart.py @@ -55,12 +55,21 @@ def main(): 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/docs/user-guide/installation.md b/docs/user-guide/installation.md index c466903..7ca9e59 100644 --- a/docs/user-guide/installation.md +++ b/docs/user-guide/installation.md @@ -110,3 +110,4 @@ Then run the same `aspec train --config ...` command. This optional supervisor n 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`. From 19f9c1a048b270f8923fdaa5bf1dd52b99586cc1 Mon Sep 17 00:00:00 2001 From: zhouaihui Date: Thu, 12 Nov 2026 15:24:33 +0800 Subject: [PATCH 11/13] fix(target): retire failed source processes after verified reader fencing --- aspec/inference/engine/target_actor.py | 2 + aspec/inference/service.py | 66 +++++++++++++++++++++++++- tests/test_target_service.py | 25 +++++++++- 3 files changed, 90 insertions(+), 3 deletions(-) diff --git a/aspec/inference/engine/target_actor.py b/aspec/inference/engine/target_actor.py index f904827..1fdd10d 100644 --- a/aspec/inference/engine/target_actor.py +++ b/aspec/inference/engine/target_actor.py @@ -48,6 +48,7 @@ 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 @@ -62,6 +63,7 @@ def get_status(self): "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, diff --git a/aspec/inference/service.py b/aspec/inference/service.py index dd2f97f..fa0f9fa 100644 --- a/aspec/inference/service.py +++ b/aspec/inference/service.py @@ -119,12 +119,14 @@ def create(self, value): self.engines[key] = {"actor": actor, "state": "created", "session": session, "gpus": gpus, "heartbeat": time.monotonic(), "owners": getattr(args, "aspec_read_owners", None), - "status": {}, "features": {}, "active": 0} + "status": {}, "features": {}, "active": 0, "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. @@ -199,6 +201,12 @@ def reap_disconnected(self, *, expiry=30): 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 @@ -208,13 +216,67 @@ def reap_disconnected(self, *, expiry=30): 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) - self.call(key, "shutdown", [], {}) + 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 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, *_): diff --git a/tests/test_target_service.py b/tests/test_target_service.py index a850e62..50cdc38 100644 --- a/tests/test_target_service.py +++ b/tests/test_target_service.py @@ -71,8 +71,11 @@ def service(monkeypatch): 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())) + get=lambda x, **_: x, kill=MagicMock(), + exceptions=ray.exceptions)) instance = TargetService(model="qwen", devices=[0, 1], runtime={}) instance.test_handle = handle return instance @@ -121,6 +124,26 @@ def test_disconnect_retains_unverified_reader_storage(service, monkeypatch): 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 From f9533318ff6720a72caad262799a0efe8961085d Mon Sep 17 00:00:00 2001 From: zhouaihui Date: Thu, 12 Nov 2026 15:31:55 +0800 Subject: [PATCH 12/13] fix(target): propagate session identity to native scheduler workers --- aspec/inference/service.py | 29 ++++++++++++++++++++++++----- tests/test_target_service.py | 29 ++++++++++++++++++++++++++++- 2 files changed, 52 insertions(+), 6 deletions(-) diff --git a/aspec/inference/service.py b/aspec/inference/service.py index fa0f9fa..583d413 100644 --- a/aspec/inference/service.py +++ b/aspec/inference/service.py @@ -87,8 +87,9 @@ def create(self, value): 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) - if not isinstance(session, str) or not session: - raise ValueError("Target requires an explicit attempt identity") + 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") @@ -111,15 +112,17 @@ def create(self, value): # 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": get_aspec_env_vars(args)}, + 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": 0, "processes": {}} + "status": {}, "features": {}, "active": 1, "processes": {}} self.claims.update({gpu: key for gpu in gpus}) try: status = ray.get(actor.get_status.remote(), timeout=120) @@ -136,6 +139,9 @@ def create(self, value): 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): @@ -252,6 +258,19 @@ def observe_process_tree(identities): 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({ + value["contract_id"] for value in contracts.values() + }) > 1)), + } + + def stop_owned_processes(identities): """Called only after the read fence; never discover or stop unrelated PIDs.""" import psutil @@ -374,7 +393,7 @@ def stopping(*_): service.reap_disconnected() time.sleep(1) except BaseException: - logger.exception("Target retained while reader exit is unverified") + logger.exception("Target cleanup is incomplete; retaining the service") server.server_close() ray.shutdown() for sig, handler in previous.items(): diff --git a/tests/test_target_service.py b/tests/test_target_service.py index 50cdc38..a3c4b02 100644 --- a/tests/test_target_service.py +++ b/tests/test_target_service.py @@ -83,7 +83,8 @@ def service(monkeypatch): def creation(service, key="engine", gpu=0): return {"key": key, "identity": service.identity, - "arguments": {"target_model_path": "qwen", "aspec_attempt_id": "attempt-1"}, + "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} @@ -113,6 +114,32 @@ def test_model_identity_and_allowed_gpu_fail_before_creation(service): 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" + + def test_disconnect_retains_unverified_reader_storage(service, monkeypatch): service.create(creation(service)) service.reap_disconnected(expiry=-1) From b151362e7e2ede1e0f468bc787d25708c477541d Mon Sep 17 00:00:00 2001 From: zhouaihui Date: Thu, 12 Nov 2026 15:33:29 +0800 Subject: [PATCH 13/13] test(target): preserve shared capture view identity across the service --- aspec/inference/service.py | 4 +--- tests/test_target_service.py | 12 ++++++++++++ 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/aspec/inference/service.py b/aspec/inference/service.py index 583d413..89c920b 100644 --- a/aspec/inference/service.py +++ b/aspec/inference/service.py @@ -265,9 +265,7 @@ def session_environment(args): "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({ - value["contract_id"] for value in contracts.values() - }) > 1)), + "ASPEC_SHARED_FEATURE_VIEWS": str(int(len(set(contracts.values())) > 1)), } diff --git a/tests/test_target_service.py b/tests/test_target_service.py index a3c4b02..5fbdf25 100644 --- a/tests/test_target_service.py +++ b/tests/test_target_service.py @@ -140,6 +140,18 @@ def test_native_scheduler_uses_the_callers_run_identity(service, monkeypatch): 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)