diff --git a/examples/disaggregated/slurm/stress_ab/README.md b/examples/disaggregated/slurm/stress_ab/README.md new file mode 100644 index 000000000000..931b25b6485f --- /dev/null +++ b/examples/disaggregated/slurm/stress_ab/README.md @@ -0,0 +1,256 @@ + + +# NVBUG 6649384: full disaggregated stress A/B + +This investigation harness compares the executor before and after the revert +of PR #16687 in PR #18327. It supports the historical comparison by default +and explicitly pinned current-head profiles. It does not establish that #16687 caused the bug, +and the revert PR must not be merged as a production fix on this evidence. + +| Historical default identity | Pinned revision | +| --- | --- | +| Control A: historical source with #16687 | `0f2c3a95f9415045bdf06a7230759475692483b6` | +| Treatment B: exact executor revert | `3245fc3ecd76e2fb610f42f2422102e2430c28fe` | +| Shared test harness | The new PR head, recorded separately in each manifest | + +Build the control once. Both arms use that wheel's identical compiled artifacts +and dependencies; only `tensorrt_llm/_torch/pyexecutor/py_executor.py` changes. +The driver extracts the two executor files directly from the immutable Git +revisions, verifies the original wheel contains the control file, and checks +compiled artifact hashes. Do not rebase the PR, change these runtime revisions, +or replace the historical wheel with an rc25 or current release overlay. + +## Current-head comparison + +To investigate whether the behavior still matters today, pin a current control +commit, prepare a reviewed semantic revert on that same commit, and build the +control once. Later refactors can move the original behavior across files: for +example, the September 18 comparison changes both `py_executor.py` and +`disagg_adapter.py`. A historical single-file overlay on a newer wheel is not +the same experiment. + +Pass an explicit JSON profile to both the builder and runner with `--profile`; +the Slurm launcher forwards `AB_PROFILE`. Its schema is: + +```json +{ + "schema_version": 1, + "name": "current-head-comparison", + "control": "", + "treatment": "", + "harness_sha": "", + "runtime_files": [ + "tensorrt_llm/_torch/pyexecutor/disagg_adapter.py", + "tensorrt_llm/_torch/pyexecutor/py_executor.py" + ], + "non_runtime_files": [], + "expected_requests": 60000, + "dependency_versions": { + "aiperf": "0.8.0", + "lm_eval": "0.4.10", + "nixl-cu13": "1.4.0" + } +} +``` + +Use actual full commit IDs and source-verified dependency pins. The profile +must match the wheel provenance, the entire control/treatment diff and the +clean common test checkout. Runtime changes are limited to explicitly listed +Python files; compiled changes require a different comparison strategy. +All overlaid modules are checked for exact import paths and file hashes. + +Keep the common test checkout at `AB_HARNESS` and the runner checkout at +`AB_RUNNER_ROOT` when they differ. Both must be inside the mounted project +root. Fetch the exact Git objects for both variants before running, including +any local treatment and harness commits transferred with a Git bundle. +Record the current source's test/configuration changes from the historical +case. Both arms use the same current test and its original gates. + +Resolve the current source's build environment to an immutable digest. A +verified public base used by its Dockerfile can support a source build when +the CI development image is inaccessible; record that environment difference +and validate its prerequisites. Do not present such a run as historical or +complete CI reproduction. If both current arms pass, the finding is simply +no reproduction under the current conditions. + +## Experiment + +The launcher runs **A, B, B, A**, sequentially in one exclusive allocation on +one physical node with eight B200 GPUs and approximately 2 TB host RAM. Each +trial gets a fresh Slurm step and container so the scheduler can clean up MPI +descendants before the next trial. It invokes the original test directly: + +```text +disaggregated/test_disaggregated.py::test_disaggregated_stress_test[input8k-output1k-conc512-gpt_oss_120b_eagle_trtllm_stress] +``` + +The historical test retains its GPT-OSS-120B plus Eagle3 setup: context TP4/EP4 +on GPUs 0–3, generation TP4/EP4 on GPUs 4–7, attention DP, PP1, disabled overlap, +Python NIXL, maximum batch 128, maximum draft length 3, concurrency 512, +60,000 requests, input length 8192, output length 1024 ± 102, seed 100 and 10% +intentional cancellations after 0.5 seconds. The test's non-cancellation error +gate remains 5%, and its subsequent GSM8K accuracy gate remains 0.42. No waive +file is supplied. A skipped test is invalid, never a pass. + +The test helper's opt-in `TLLM_DISAGG_STRESS_KEEP_LOGS=1` preserves the worker +logs that were previously deleted even with `--keep-workspace`. Both arms use +the same modified test helper. Default cleanup for other callers is unchanged. + +## Prepare on an approved compute node + +Check the target cluster's current policy before source preparation or builds. +Use a current allocation owned by the submitting user. Keep account, QoS, +partition, registry credentials and private paths in a local run record. + +Use a persistent project root containing the complete Git clone, both source +checkouts, build directory, wheel and results. Mount that root at the same +absolute path in every container; detached worktree Git pointers and the build +virtual environment must remain valid. Initialize the control's pinned +submodules and keep that checkout completely clean. Use a new directory for +each build and experiment. Do not modify shared model files during the matrix. + +The historical CI build environment is: + +```text +artifactory.nvidia.com/sw-tensorrt-llm-docker-local/tensorrt-llm:pytorch-26.05-py3-x86_64-ubuntu24.04-skip-tritondevel-202607311529-16970 +``` + +Resolve an accessible registry manifest and record its immutable `sha256:...` +digest. A tag alone is insufficient. An installed wheel inside that image is +not proof of control-source identity. If the historical image is unavailable, +stop and resolve an exact mirror or explicitly agree on a common replacement +and its comparability limitation before testing. No replacement is automatic. + +In a bounded build allocation, launch that image **by digest**, mount the +persistent project root, set `TLLM_AB_IMAGE` and `TLLM_AB_IMAGE_DIGEST` to the +recorded reference and digest, and run with the image's system Python: + +```bash +python3 "$AB_HARNESS/examples/disaggregated/slurm/stress_ab/build_runtime.py" \ + --source "$AB_PROJECT_ROOT/control" \ + --output "$AB_PROJECT_ROOT/control-build" \ + --build-root "$AB_PROJECT_ROOT/control-build-state" \ + --image "$AB_IMAGE" --image-digest "$AB_IMAGE_DIGEST" \ + --jobs 16 --timeout 21600 +``` + +The builder uses the historical `scripts/build_wheel.py`, an out-of-tree clean +build and B200 `100-real` kernels. It pre-stages the unchanged +`requirements-grpc-smg.txt` omitted by that revision's packaging helper and +records the file hash. It records the command, source and submodules, +image digest, build log, wheel hash and prepared build virtual environment. +The build's six-hour limit is separate from the test allocation. Select a +permitted build allocation with time for preparation and artifact flush too. + +With an explicit profile, use that control revision's build script and pass +the same `--profile /persistent/path/profile.json` to the builder. Optional +`--cpp-build-dir /node-local/path/cpp` puts large CMake intermediates on local +scratch while preserving the build venv, wheel and logs persistently. +`--skip-stubs` skips Python type-stub generation for a CPU-only build; it does +not skip native compilation. These choices are recorded in the build command. + +**Retain the build virtual environment.** `build_wheel.py` installs the pinned +source's development requirements there, not into the base image's Python. +The experiment runs with the recorded interpreter and verifies its dependency +snapshot. Installing the wheel with `--no-deps` into an unprepared base image +is not sufficient. Do not install or upgrade dependencies between trials. + +Before allocating eight GPUs, verify available storage for the image, build, +wheel, installed runtime copies and four complete log sets. Do not use a small +home quota or assume the build fits because the model weights are shared. +The driver retains runtime trees and refuses a trial if free space is less +than four wheel sizes plus 20 GiB; that check is not a build-size estimate. + +Required readable shared assets below `AB_MODELS_ROOT`: + +- `gpt_oss/gpt-oss-120b`: configuration, local tokenizer and all indexed shards. +- `gpt_oss/gpt-oss-120b-Eagle3`: configuration and all weight shards. +- `datasets/openai/gsm8k/main/test-00000-of-00001.parquet` with `question` and + `answer` columns. + +The runner preflights local-only tokenizer loading, dataset readability, +`aiperf==0.8.0`, `lm_eval==0.4.10`, CLI entry points, NIXL, runtime import paths +and GPU identity. Model checkpoints being present alone is not a runtime +validation. The driver records input hashes and versions; keep the shared +mount unchanged throughout the matrix. + +## Submit one bounded matrix + +Set these variables in a private run configuration. `AB_PYTHON` must be the +literal `runtime_python` recorded in `provenance.json`; do not resolve its +symlink to the system Python. All project paths must be beneath the mounted +`AB_PROJECT_ROOT`. + +```bash +export AB_PROJECT_ROOT=/persistent/path/to/experiment +export AB_HARNESS="$AB_PROJECT_ROOT/harness" +export AB_WHEEL="$AB_PROJECT_ROOT/control-build/wheels/.whl" +export AB_PROVENANCE="$AB_PROJECT_ROOT/control-build/provenance.json" +export AB_PYTHON="$AB_PROJECT_ROOT/control-build-state/venv-3.12/bin/python3" +export AB_MODELS_ROOT=/shared/path/to/llm-models +export AB_RUN_ROOT="$AB_PROJECT_ROOT/results/abba-" +export AB_IMAGE='' +export AB_IMAGE_DIGEST='sha256:' +export AB_TRIAL_TIMEOUT=12600 + +# For an explicit current-head comparison, also set: +# export AB_PROFILE="$AB_PROJECT_ROOT/profile.json" +# export AB_RUNNER_ROOT="$AB_PROJECT_ROOT/runner" + +sbatch --parsable --account='' \ + --partition='' --qos='' \ + --output="$AB_PROJECT_ROOT/slurm-%j.out" \ + "$AB_HARNESS/examples/disaggregated/slurm/stress_ab/launch.slurm" +``` + +The script requests one node, eight GPUs, 224 CPUs, all host memory and 16 hours. +Four 3.5-hour test limits plus setup and flush require at least 15 hours. It +checks the actual granted wall time before starting. A four-hour allocation +cannot fit this matrix. Queue wait and the separate source build are not part +of these 16 hours. Use a permitted longer QoS or a real reservation; a longer +partition maximum alone does not establish QoS eligibility. + +On a preemptible route, `--no-requeue` prevents automatic repetitions. An +interruption leaves the matrix invalid/incomplete. Record the job ID once, +verify scheduler resources and arrange status monitoring. Do not resubmit +after an ambiguous submission response without reconciling the unique run ID. + +## Results and interpretation + +`plan.json` records the intended matrix. `slurm-steps.tsv` preserves every +step exit code. Each step writes its own `summary.json`, `manifest.json`, pip +and preflight logs, JUnit XML, pytest output and retained worker/request logs. +Progress is saved to persistent storage before and after each trial. + +The final `summary.json` reports per-arm pass/fail/invalid counts. Exit codes +are 0 for all four passing trials, 1 for a complete matrix containing a test +failure, and 2 for an invalid matrix. Scheduler interruption can return a +different nonzero status. Missing or extra tests, skipped/setup failures, +missing worker logs, incomplete 60,000-request accounting, runtime mismatch, +timeout and residual processes cannot become a pass. A valid test failure can +continue to the next arm; an invalid setup stops the matrix without retries. + +Compare the **first initiating context/generation worker exception**, request +accounting and accuracy result across both repetitions. The proxy's later +`LLM is shutting down` error alone does not identify the cause. A failing +control and passing revert supports an effect in this environment; two runs +per arm do not establish flake causality. Both passing means no reproduction +under these conditions. Both failing with different initiating exceptions +does not establish the same bug or exonerate the candidate change. + +The earlier TP2 rc25 eight-request smoke test and PR CI build/release checks +are not this experiment. The earlier CI success ran zero mapped tests and +must not be reported as B200 stress coverage. + +## Local validation without GPUs + +```bash +python3 -m unittest discover -s examples/disaggregated/slurm/stress_ab -v +bash -n examples/disaggregated/slurm/stress_ab/launch.slurm +python3 examples/disaggregated/slurm/stress_ab/run_ab.py \ + --harness . --wheel /unused/control.whl --provenance /unused/provenance.json \ + --models-root /unused/models --output /unused/results --dry-plan +``` + +These checks validate the harness and plan only. They do not build the control, +validate model/runtime availability or execute the GPU test. diff --git a/examples/disaggregated/slurm/stress_ab/build_runtime.py b/examples/disaggregated/slurm/stress_ab/build_runtime.py new file mode 100644 index 000000000000..81e5989c9b89 --- /dev/null +++ b/examples/disaggregated/slurm/stress_ab/build_runtime.py @@ -0,0 +1,413 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Build the pinned control once and record provenance for both A/B arms.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import re +import shutil +import signal +import subprocess +import sys +from datetime import datetime, timezone +from pathlib import Path + +CONTROL = "0f2c3a95f9415045bdf06a7230759475692483b6" +TREATMENT = "3245fc3ecd76e2fb610f42f2422102e2430c28fe" +EXECUTOR = "tensorrt_llm/_torch/pyexecutor/py_executor.py" +HISTORICAL_PROFILE = { + "schema_version": 1, + "name": "historical-6649384", + "control": CONTROL, + "treatment": TREATMENT, + "harness_sha": None, + "runtime_files": [EXECUTOR], + "non_runtime_files": ["tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py"], + "expected_requests": 60000, + "dependency_versions": {"aiperf": "0.8.0", "lm_eval": "0.4.10", "nixl-cu13": "1.3.1"}, +} + + +def load_profile(path: Path | None = None) -> dict: + """Read an immutable comparison; only listed Python runtime files may differ.""" + profile = json.loads(path.read_text()) if path else dict(HISTORICAL_PROFILE) + if ( + not isinstance(profile, dict) + or set(profile) != set(HISTORICAL_PROFILE) + or profile["schema_version"] != 1 + ): + raise ValueError("unknown or incomplete A/B profile schema") + if not isinstance(profile["name"], str) or not profile["name"]: + raise ValueError("profile name must be nonempty") + for field in ("control", "treatment", "harness_sha"): + if field == "harness_sha" and path is None: + continue + if not isinstance(profile[field], str) or not re.fullmatch(r"[0-9a-f]{40}", profile[field]): + raise ValueError(f"profile {field} must be a full immutable commit SHA") + if profile["control"] == profile["treatment"]: + raise ValueError("comparison commits must differ") + for field, prefix in ( + ("runtime_files", "tensorrt_llm/"), + ("non_runtime_files", "tests/unittest/"), + ): + paths = profile[field] + if not isinstance(paths, list) or not all(isinstance(item, str) for item in paths): + raise ValueError(f"profile {field} must contain file paths") + if len(set(paths)) != len(paths) or (field == "runtime_files" and not paths): + raise ValueError(f"profile {field} must contain unique file paths") + for item in paths: + if ( + not item.startswith(prefix) + or not item.endswith(".py") + or any(not part.isidentifier() for part in item[:-3].split("/")) + ): + raise ValueError(f"unsafe or non-Python comparison path: {item}") + if type(profile["expected_requests"]) is not int or profile["expected_requests"] < 1: + raise ValueError("expected_requests must be a positive integer") + dependencies = profile["dependency_versions"] + if not isinstance(dependencies, dict) or set(dependencies) != { + "aiperf", + "lm_eval", + "nixl-cu13", + }: + raise ValueError("profile must pin aiperf, lm_eval and nixl-cu13") + if dependencies["aiperf"] != "0.8.0" or dependencies["lm_eval"] != "0.4.10": + raise ValueError("this stress accounting requires AIPerf 0.8.0 and lm_eval 0.4.10") + if not all( + isinstance(version, str) and re.fullmatch(r"[0-9]+(?:\.[0-9]+)+", version) + for version in dependencies.values() + ): + raise ValueError("dependency versions must be exact numeric versions") + return profile + + +def validate_source_profile(source: Path, profile: dict) -> None: + """Reject compiled, configuration, or unlisted changes before reusing one wheel.""" + changed = set( + _git( + source, "diff", "--name-only", "--no-renames", profile["control"], profile["treatment"] + ).splitlines() + ) + if changed != set(profile["runtime_files"] + profile["non_runtime_files"]): + raise ValueError("comparison has unexpected source changes") + for arm in ("control", "treatment"): + for relative in profile["runtime_files"]: + entry = _git(source, "ls-tree", profile[arm], "--", relative) + if not entry.startswith("100644 blob "): + raise ValueError( + f"comparison runtime file must be a regular Python file: {relative}" + ) + requirements = _git(source, "show", f"{profile['control']}:requirements-dev.txt") + for name, version in profile["dependency_versions"].items(): + pattern = rf"^{re.escape(name)}(?:\[[^\]]+\])?=={re.escape(version)}\s*(?:#.*)?$" + if not re.search(pattern, requirements, re.MULTILINE): + raise ValueError(f"profile dependency does not match control requirements: {name}") + + +def _git(source: Path, *args: str) -> str: + return subprocess.check_output(["git", "-C", str(source), *args], text=True).strip() + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(8 * 1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def runtime_identity(python: Path) -> dict: + """Inspect one interpreter without user-site, PYTHONPATH, or working-directory overlays.""" + code = """ +import importlib.metadata, json, os, re, sys +packages = [] +for distribution in importlib.metadata.distributions(): + name = distribution.metadata.get('Name') + if not name or not distribution.version: + raise ValueError('installed distribution has incomplete metadata') + packages.append({'name': re.sub(r'[-_.]+', '-', name).lower(), + 'version': distribution.version, + 'location': os.path.abspath(distribution.locate_file(''))}) +packages.sort(key=lambda item: (item['name'], item['version'], item['location'])) +print(json.dumps({'runtime_python': os.path.abspath(sys.executable), + 'runtime_prefix': os.path.abspath(sys.prefix), + 'runtime_python_version': sys.version, + 'runtime_sys_path': sys.path, + 'runtime_distributions': packages})) +""" + return json.loads( + subprocess.check_output([str(python), "-I", "-c", code], text=True, timeout=120) + ) + + +def _prepare_cli_wrappers(python: Path) -> dict: + code = """ +import importlib.metadata, json +result = {} +for name, package, version in (('aiperf', 'aiperf', '0.8.0'), ('lm_eval', 'lm_eval', '0.4.10')): + distribution = importlib.metadata.distribution(package) + if distribution.version != version: + raise ValueError(f'{package} must be {version}, got {distribution.version}') + entries = [entry for entry in distribution.entry_points + if entry.group == 'console_scripts' and entry.name == name] + if len(entries) != 1: + raise ValueError(f'{package} has no unique {name} console entry point') + result[name] = {'distribution': package, 'version': version, 'entrypoint': entries[0].value} +print(json.dumps(result)) +""" + entries = json.loads( + subprocess.check_output([str(python), "-I", "-c", code], text=True, timeout=120) + ) + for name, metadata in entries.items(): + script = python.parent / name + created = not script.exists() + if created: + script.write_text( + f"#!{python}\n" + "from importlib.metadata import distribution\n" + f"entry = next(item for item in distribution({metadata['distribution']!r}).entry_points\n" + f" if item.group == 'console_scripts' and item.name == {name!r}\n" + f" and item.value == {metadata['entrypoint']!r})\n" + "raise SystemExit(entry.load()())\n" + ) + script.chmod(0o755) + first_line = script.read_text().splitlines()[0] + if not first_line.startswith("#!") or Path(first_line[2:]).parent != python.parent: + raise ValueError(f"{script} is not bound to the baseline venv interpreter") + if not os.access(script, os.X_OK): + raise ValueError(f"baseline CLI is not executable: {script}") + metadata.update(path=str(script), sha256=_sha256(script), generated=created) + return entries + + +def main() -> int: + """Build in an approved compute allocation, never on a login frontend.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--source", type=Path, required=True) + parser.add_argument("--profile", type=Path) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--build-root", type=Path, required=True) + parser.add_argument("--cpp-build-dir", type=Path) + parser.add_argument("--skip-stubs", action="store_true") + parser.add_argument("--image", required=True) + parser.add_argument("--image-digest", required=True) + parser.add_argument("--jobs", type=int, default=16) + parser.add_argument("--timeout", type=int, default=21600) + parser.add_argument("--dry-run", action="store_true") + args = parser.parse_args() + if sys.prefix != sys.base_prefix: + parser.error("invoke the builder with container system Python, outside an existing venv") + profile = load_profile(args.profile) + source = args.source.resolve() + output = args.output.resolve() + build_root = args.build_root.resolve() + if args.jobs < 1 or args.timeout < 1: + parser.error("jobs and timeout must be positive") + if _git(source, "rev-parse", "HEAD") != profile["control"]: + parser.error(f"source must be the exact control commit {profile['control']}") + validate_source_profile(source, profile) + if _git(source, "status", "--porcelain", "--untracked-files=all"): + parser.error("control source must be clean, including untracked files") + submodules = _git(source, "submodule", "status", "--recursive") + if any(line.startswith(("-", "+", "U")) for line in submodules.splitlines()): + parser.error("initialize all pinned submodules before building") + if not re.fullmatch(r"sha256:[0-9a-f]{64}", args.image_digest): + parser.error("resolve and record the actual registry image digest first") + if output == build_root or source in output.parents or source in build_root.parents: + parser.error("output and build-root must be distinct and outside the source checkout") + if output.exists() or build_root.exists(): + parser.error("output and build-root must be new paths; existing runs are never overwritten") + # Keep every dependency installation in the persistent runtime, including build setup. + runtime_prefix = build_root / f"venv-{sys.version_info.major}.{sys.version_info.minor}" + runtime_python = runtime_prefix / "bin/python3" + venv_command = [ + sys.executable, + "-I", + "-m", + "venv", + "--system-site-packages", + str(runtime_prefix), + ] + command = [ + str(runtime_python), + str(source / "scripts/build_wheel.py"), + "--no-venv", + "--clean", + "--out-of-tree", + "--build_root", + str(build_root), + "--dist_dir", + str(output / "wheels"), + "--job_count", + str(args.jobs), + "--cuda_architectures", + "100-real", + "--yes", + ] + if args.cpp_build_dir: + cpp_build_dir = args.cpp_build_dir.resolve() + if cpp_build_dir.exists() or cpp_build_dir == source or source in cpp_build_dir.parents: + parser.error("cpp-build-dir must be a new directory outside the source checkout") + command.extend(["--build_dir", str(cpp_build_dir)]) + if args.skip_stubs: + command.append("--skip-stubs") + if args.dry_run: + print( + json.dumps( + { + "source_sha": profile["control"], + "profile": profile, + "venv_creation_command": venv_command, + "build_command": command, + }, + indent=2, + ) + ) + return 0 + if not os.environ.get("SLURM_JOB_ID"): + parser.error("build in an approved Slurm compute allocation, not a frontend") + if ( + os.environ.get("TLLM_AB_IMAGE") != args.image + or os.environ.get("TLLM_AB_IMAGE_DIGEST") != args.image_digest + ): + parser.error("container launcher must set matching TLLM_AB_IMAGE and TLLM_AB_IMAGE_DIGEST") + output.mkdir(parents=True) + build_root.mkdir(parents=True) + staging_additions = [] + if profile["control"] == CONTROL: + # Historical out-of-tree staging omits this setup.py dependency. + extra_requirement = source / "requirements-grpc-smg.txt" + staged_requirement = build_root / "package/requirements-grpc-smg.txt" + staged_requirement.parent.mkdir() + shutil.copy2(extra_requirement, staged_requirement) + staging_additions.append( + { + "source": str(extra_requirement), + "destination": str(staged_requirement), + "sha256": _sha256(extra_requirement), + } + ) + build_log = output / "build.log" + manifest = { + "schema_version": 1, + "source_sha": profile["control"], + "profile": profile, + "clean_source": True, + "submodules": submodules.splitlines(), + "image": args.image, + "image_digest": args.image_digest, + "venv_creation_command": venv_command, + "build_command": command, + "build_log": str(build_log), + "staging_additions": staging_additions, + "slurm_job_id": os.environ["SLURM_JOB_ID"], + "started_at": datetime.now(timezone.utc).isoformat(), + "status": "building", + } + manifest_path = output / "build-status.json" + manifest_path.write_text(json.dumps(manifest, indent=2) + "\n") + build_env = dict(os.environ) + for key in ( + "PYTHONPATH", + "PYTHONHOME", + "PYTHONUSERBASE", + "PIP_TARGET", + "PIP_PREFIX", + "PIP_USER", + ): + build_env.pop(key, None) + build_env.update( + PATH=f"{runtime_python.parent}:{os.environ['PATH']}", + VIRTUAL_ENV=str(runtime_prefix), + PYTHONNOUSERSITE="1", + ) + with build_log.open("w") as stream: + for step, timeout in ((venv_command, 300), (command, args.timeout)): + process = subprocess.Popen( + step, + cwd=source, + env=build_env, + stdout=stream, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + try: + returncode = process.wait(timeout=timeout) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGTERM) + try: + process.wait(timeout=30) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGKILL) + process.wait() + returncode = 124 + if step is venv_command: + manifest["venv_creation_returncode"] = returncode + manifest_path.write_text(json.dumps(manifest, indent=2) + "\n") + if returncode != 0: + break + manifest["returncode"] = returncode + manifest["finished_at"] = datetime.now(timezone.utc).isoformat() + wheels = list((output / "wheels").glob("tensorrt_llm-*.whl")) + source_unchanged = not _git(source, "status", "--porcelain", "--untracked-files=all") + manifest["source_unchanged_after_build"] = source_unchanged + if returncode != 0 or len(wheels) != 1 or not source_unchanged: + manifest["status"] = "invalid" + manifest_path.write_text(json.dumps(manifest, indent=2) + "\n") + print(f"Build invalid; inspect {manifest_path} and {build_log}", file=sys.stderr) + return returncode if 0 < returncode < 256 else 2 + # Preserve the venv path rather than resolving its symlink to the system interpreter. + try: + entrypoints = _prepare_cli_wrappers(runtime_python) + identity = runtime_identity(runtime_python) + for name, version in profile["dependency_versions"].items(): + normalized = name.replace("_", "-") + installed = [ + item["version"] + for item in identity["runtime_distributions"] + if item["name"] == normalized + ] + if not installed or any(value != version for value in installed): + raise ValueError(f"baseline dependency must match profile: {name}=={version}") + if identity["runtime_python"] != str(runtime_python) or identity["runtime_prefix"] != str( + runtime_prefix + ): + raise ValueError("built runtime interpreter does not belong to the expected venv") + except (OSError, ValueError, subprocess.SubprocessError) as error: + manifest.update( + status="invalid", reason=f"runtime environment verification failed: {error}" + ) + manifest_path.write_text(json.dumps(manifest, indent=2) + "\n") + print(f"Build invalid; inspect {manifest_path}", file=sys.stderr) + return 2 + manifest.update( + status="built", + wheel=wheels[0].name, + wheel_sha256=_sha256(wheels[0]), + runtime_entrypoints=entrypoints, + **identity, + ) + manifest_path.write_text(json.dumps(manifest, indent=2) + "\n") + (output / "provenance.json").write_text(json.dumps(manifest, indent=2) + "\n") + print(f"Wheel: {wheels[0]}\nProvenance: {output / 'provenance.json'}") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/examples/disaggregated/slurm/stress_ab/launch.slurm b/examples/disaggregated/slurm/stress_ab/launch.slurm new file mode 100644 index 000000000000..3f56d451d779 --- /dev/null +++ b/examples/disaggregated/slurm/stress_ab/launch.slurm @@ -0,0 +1,147 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +#SBATCH --job-name=disagg-stress-ab +#SBATCH --nodes=1 +#SBATCH --ntasks=1 +#SBATCH --gpus-per-node=8 +#SBATCH --cpus-per-task=224 +#SBATCH --mem=0 +#SBATCH --exclusive +#SBATCH --time=16:00:00 +#SBATCH --no-requeue +#SBATCH --signal=B:TERM@120 + +# Account, partition and QoS must be explicitly selected at submission according +# to the site's current policy. This script never submits or retries a job. +set -Eeuo pipefail + +: "${AB_PROJECT_ROOT:?Persistent project root containing the clone, worktree, wheel and results}" +: "${AB_HARNESS:?Pinned PR checkout under AB_PROJECT_ROOT}" +: "${AB_WHEEL:?Control wheel from build_runtime.py}" +: "${AB_PROVENANCE:?Matching provenance.json from build_runtime.py}" +: "${AB_PYTHON:?Absolute runtime_python from provenance.json; retain its build venv}" +: "${AB_MODELS_ROOT:?Readable GPT-OSS and Eagle3 model root}" +: "${AB_RUN_ROOT:?New persistent run directory under AB_PROJECT_ROOT}" +: "${AB_IMAGE:?Historical container registry/repository:tag}" +: "${AB_IMAGE_DIGEST:?Resolved immutable sha256 digest}" +: "${SLURM_JOB_ID:?Launch through sbatch in an approved compute allocation}" +[[ ${SLURM_JOB_NUM_NODES:-0} == 1 ]] || { echo 'Exactly one node is required' >&2; exit 2; } +[[ $AB_IMAGE_DIGEST =~ ^sha256:[0-9a-f]{64}$ ]] || { echo 'Invalid image digest' >&2; exit 2; } +[[ $AB_IMAGE =~ ^[a-zA-Z0-9._/-]+:[a-zA-Z0-9._-]+$ ]] || { echo 'Use a registry/repository:tag image' >&2; exit 2; } +[[ ! -e $AB_RUN_ROOT ]] || { echo 'Run directory already exists; refusing to overwrite' >&2; exit 2; } +runner_root=${AB_RUNNER_ROOT:-$AB_HARNESS} +profile_args=() +if [[ -n ${AB_PROFILE:-} ]]; then + [[ -f $AB_PROFILE && $AB_PROFILE == "$AB_PROJECT_ROOT/"* && $AB_PROFILE != *','* && $AB_PROFILE != *':'* && $AB_PROFILE != *$'\n'* ]] || { + echo 'Profile must be an existing file under AB_PROJECT_ROOT' >&2; exit 2; + } + profile_args=(--profile "$AB_PROFILE") +fi +for path in "$AB_PROJECT_ROOT" "$AB_HARNESS" "$runner_root" "$AB_WHEEL" "$AB_PROVENANCE" "$AB_PYTHON" "$AB_MODELS_ROOT" "$AB_RUN_ROOT"; do + [[ $path == /* && $path != *','* && $path != *':'* && $path != *$'\n'* ]] || { + echo 'Use absolute mount paths without comma, colon or newline' >&2; exit 2; + } +done +for path in "$AB_HARNESS" "$runner_root" "$AB_WHEEL" "$AB_PROVENANCE" "$AB_PYTHON" "$AB_RUN_ROOT"; do + [[ $path == "$AB_PROJECT_ROOT/"* ]] || { echo 'Project artifacts must be under AB_PROJECT_ROOT' >&2; exit 2; } +done +[[ -f $AB_WHEEL && -f $AB_PROVENANCE && -d $AB_HARNESS && -d $AB_MODELS_ROOT ]] || { + echo 'Missing prepared source, wheel, provenance or models' >&2; exit 2; +} + +trial_timeout=${AB_TRIAL_TIMEOUT:-12600} +[[ $trial_timeout =~ ^[1-9][0-9]*$ ]] || { echo 'Invalid trial timeout' >&2; exit 2; } +# Four full test timeouts plus 10 minutes of preflight per step and 20 minutes +# for container/setup/flush. A shorter allocation is rejected before testing. +required_seconds=$((4 * (trial_timeout + 600) + 1200)) +job_time_limit=$(scontrol show job "$SLURM_JOB_ID" -o | python3 -c ' +import re, sys +value = re.search(r"\bTimeLimit=(\S+)", sys.stdin.read()).group(1) +days, clock = value.split("-", 1) if "-" in value else ("0", value) +parts = [int(v) for v in clock.split(":")] +if len(parts) == 2: + parts.insert(0, 0) +print(int(days)*86400 + parts[0]*3600 + parts[1]*60 + parts[2])') +(( job_time_limit >= required_seconds )) || { echo 'Allocation cannot fit the bounded ABBA plan' >&2; exit 2; } + +mkdir -p "$AB_RUN_ROOT" +runner_dir="$runner_root/examples/disaggregated/slurm/stress_ab" +export TLLM_AB_IMAGE="$AB_IMAGE" TLLM_AB_IMAGE_DIGEST="$AB_IMAGE_DIGEST" +image_repo=${AB_IMAGE%:*} +container_image="${image_repo/\//#}@${AB_IMAGE_DIGEST}" +# Preserve absolute paths so detached worktree .git pointers remain valid. +mounts="$AB_PROJECT_ROOT:$AB_PROJECT_ROOT,$AB_MODELS_ROOT:$AB_MODELS_ROOT:ro" +printf 'index\tarm\texit_code\n' > "$AB_RUN_ROOT/slurm-steps.tsv" +python3 - "$AB_RUN_ROOT/plan.json" "$AB_HARNESS" "$AB_IMAGE" "$AB_IMAGE_DIGEST" "$trial_timeout" "$runner_root" "${AB_PROFILE:-}" <<'PY' +import json, os, socket, subprocess, sys +from datetime import datetime, timezone +from pathlib import Path +path, harness, image, digest, timeout, runner, profile = sys.argv[1:] +plan = { + 'schema_version': 1, 'order': ['control', 'treatment', 'treatment', 'control'], + 'slurm_job_id': os.environ['SLURM_JOB_ID'], 'node': socket.gethostname(), + 'started_at': datetime.now(timezone.utc).isoformat(), 'trial_timeout': int(timeout), + 'harness_sha': subprocess.check_output(['git', '-C', harness, 'rev-parse', 'HEAD'], text=True).strip(), + 'runner_sha': subprocess.check_output(['git', '-C', runner, 'rev-parse', 'HEAD'], text=True).strip(), + 'profile': json.loads(Path(profile).read_text()) if profile else None, + 'image': image, 'image_digest': digest, 'same_node': True, + 'cleanup_boundary': 'fresh Slurm step and container for each trial', +} +Path(path).write_text(json.dumps(plan, indent=2) + '\n') +PY + +finish() { + status=$? + trap - EXIT + set +e + python3 "$runner_dir/summarize.py" "$AB_RUN_ROOT" + summary_status=$? + if (( status != 0 )); then exit "$status"; fi + exit "$summary_status" +} +trap finish EXIT +trap 'exit 143' TERM +trap 'exit 130' INT + +index=0 +for arm in control treatment treatment control; do + index=$((index + 1)) + trial_dir=$(printf '%s/%02d-%s' "$AB_RUN_ROOT" "$index" "$arm") + if srun --nodes=1 --ntasks=1 --gpus-per-node=8 --kill-on-bad-exit=1 \ + --container-image="$container_image" --container-mounts="$mounts" \ + --container-workdir="$AB_PROJECT_ROOT" \ + --output="$AB_RUN_ROOT/step-$index.out" --error="$AB_RUN_ROOT/step-$index.err" \ + "$AB_PYTHON" "$runner_dir/run_ab.py" --harness "$AB_HARNESS" \ + --wheel "$AB_WHEEL" --provenance "$AB_PROVENANCE" \ + --models-root "$AB_MODELS_ROOT" --output "$trial_dir" \ + --order "$arm" --trial-timeout "$trial_timeout" ${profile_args[@]+"${profile_args[@]}"}; then + step_status=0 + else + step_status=$? + fi + printf '%s\t%s\t%s\n' "$index" "$arm" "$step_status" >> "$AB_RUN_ROOT/slurm-steps.tsv" + # An expected pytest failure is an observation. Infrastructure, timeout or + # invalid setup stops this bounded run; no automatic retries are made. + if (( step_status != 0 && step_status != 1 )); then exit "$step_status"; fi + # srun itself can also return 1; only a complete classified pytest failure + # may continue to the next arm. + python3 - "$trial_dir/summary.json" "$arm" "$step_status" <<'PY' +import json, sys +from pathlib import Path +path, arm, code = sys.argv[1:] +data = json.loads(Path(path).read_text()) +trials = data.get('trials', []) +expected = 'pass' if code == '0' else 'fail' +if (data.get('status') != expected or len(trials) != 1 + or trials[0].get('arm') != arm or trials[0].get('status') != expected): + sys.exit(2) +PY +done diff --git a/examples/disaggregated/slurm/stress_ab/run_ab.py b/examples/disaggregated/slurm/stress_ab/run_ab.py new file mode 100644 index 000000000000..839e68afb1b8 --- /dev/null +++ b/examples/disaggregated/slurm/stress_ab/run_ab.py @@ -0,0 +1,701 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Run the pinned NVBUG 6649384 comparison inside an exclusive eight-B200 Slurm step.""" + +from __future__ import annotations + +import argparse +import ast +import hashlib +import json +import os +import re +import shutil +import signal +import socket +import subprocess +import sys +import time +import uuid +import xml.etree.ElementTree as ET +from pathlib import Path + +from build_runtime import load_profile, runtime_identity, validate_source_profile + +CASE = "test_disaggregated_stress_test[input8k-output1k-conc512-gpt_oss_120b_eagle_trtllm_stress]" +SELECTOR = f"disaggregated/test_disaggregated.py::{CASE}" + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise ValueError(message) + + +def _sha(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(8 * 1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _save(path: Path, data: dict) -> None: + temporary = path.with_suffix(".tmp") + temporary.write_text(json.dumps(data, indent=2) + "\n") + temporary.replace(path) + + +def _command(command: list[str], timeout: int = 300, **kwargs) -> str: + return subprocess.check_output(command, text=True, timeout=timeout, **kwargs).strip() + + +def _environment(runtime: Path, harness: Path, models: Path, run_id: str) -> dict[str, str]: + env = dict(os.environ) + for key in ( + "PYTHONPATH", + "PYTHONHOME", + "PYTHONSTARTUP", + "PYTHONUSERBASE", + "PYTHONOPTIMIZE", + "PYTEST_ADDOPTS", + ): + env.pop(key, None) + env.update( + PYTHONPATH=str(runtime), + PYTHONNOUSERSITE="1", + PYTHONDONTWRITEBYTECODE="1", + PATH=f"{runtime / 'bin'}:{Path(sys.executable).parent}:{os.environ['PATH']}", + VIRTUAL_ENV=sys.prefix, + LLM_ROOT=str(harness), + LLM_MODELS_ROOT=str(models), + TLLM_DISAGG_STRESS_KEEP_LOGS="1", + TLLM_AB_RUN_ID=run_id, + HF_HUB_OFFLINE="1", + TRANSFORMERS_OFFLINE="1", + ) + return env + + +def _validate_provenance(wheel: Path, provenance: dict, profile: dict | None = None) -> None: + profile = profile or load_profile() + _require(provenance.get("status") == "built", "baseline build did not finish successfully") + _require(provenance.get("source_unchanged_after_build") is True, "source changed during build") + _require( + provenance.get("source_sha") == profile["control"], + "wheel source SHA is not the pinned control", + ) + _require( + provenance.get("profile", load_profile()) == profile, "build and runtime profiles differ" + ) + _require(provenance.get("clean_source") is True, "baseline build source was not clean") + _require(provenance.get("wheel") == wheel.name, "provenance wheel filename mismatch") + _require(provenance.get("wheel_sha256") == _sha(wheel), "wheel SHA256 mismatch") + _require( + bool(provenance.get("build_log") and provenance.get("build_command")), + "missing build command/log provenance", + ) + _require( + bool(re.fullmatch(r"sha256:[0-9a-f]{64}", provenance.get("image_digest", ""))), + "missing immutable container digest", + ) + for field in ("image", "image_digest"): + _require( + os.environ.get(f"TLLM_AB_{field.upper()}") == provenance.get(field), + f"running container {field} does not match build provenance", + ) + + +def _validate_dependencies(provenance: dict) -> dict: + expected_python = provenance.get("runtime_python") + expected_prefix = provenance.get("runtime_prefix") + _require(sys.prefix != sys.base_prefix, "run the driver with the recorded baseline venv Python") + _require( + os.path.abspath(sys.executable) == expected_python, + "driver interpreter differs from recorded runtime_python; preserve the venv symlink path", + ) + _require( + os.path.abspath(sys.prefix) == expected_prefix, "driver uses a different runtime_prefix" + ) + actual = runtime_identity(Path(sys.executable)) + _require(bool(actual.get("runtime_distributions")), "no installed dependency inventory") + for field, value in actual.items(): + _require( + provenance.get(field) == value, f"baseline dependency environment changed: {field}" + ) + entrypoints = provenance.get("runtime_entrypoints", {}) + _require(set(entrypoints) == {"aiperf", "lm_eval"}, "missing baseline CLI provenance") + for name, metadata in entrypoints.items(): + script = Path(expected_prefix) / "bin" / name + _require( + metadata.get("path") == str(script) and metadata.get("sha256") == _sha(script), + f"baseline CLI changed: {name}", + ) + return actual + + +def _prepare_runtime(args: argparse.Namespace, provenance: dict) -> dict: + _validate_provenance(args.wheel, provenance, args.profile) + _require( + shutil.disk_usage(args.output).free > args.wheel.stat().st_size * 4 + 20 * 1024**3, + "insufficient free storage for the installed wheel and retained request/worker logs", + ) + git = ["git", "-C", str(args.harness)] + validate_source_profile(args.harness, args.profile) + harness = _validate_harness(args.harness, args.profile) + blobs = { + arm: { + relative: subprocess.check_output(git + ["show", f"{args.profile[arm]}:{relative}"]) + for relative in args.profile["runtime_files"] + } + for arm in ("control", "treatment") + } + base = args.output / "runtime" / "base" + base.mkdir(parents=True) + install = [ + sys.executable, + "-m", + "pip", + "install", + "--no-index", + "--no-deps", + "--no-compile", + "--target", + str(base), + str(args.wheel), + ] + with (args.output / "pip.log").open("w") as log: + subprocess.run(install, check=True, stdout=log, stderr=subprocess.STDOUT, timeout=600) + for relative, blob in blobs["control"].items(): + _require( + (base / relative).read_bytes() == blob, + f"wheel Python is not control source: {relative}", + ) + _require(not list(base.rglob("*.pyc")), "wheel contains stale bytecode") + binaries = { + str(path.relative_to(base)): _sha(path) + for path in base.rglob("*") + if path.is_file() and (".so" in path.name or path.suffix in (".a", ".cubin", ".fatbin")) + } + _require(bool(binaries), "wheel has no compiled artifacts") + for arm, files in blobs.items(): + _overlay_runtime(base, base.parent / arm, files, binaries) + return { + "provenance": provenance, + "install_command": install, + "binaries": binaries, + "runtime_file_sha256": { + arm: {relative: hashlib.sha256(blob).hexdigest() for relative, blob in files.items()} + for arm, files in blobs.items() + }, + **harness, + } + + +def _overlay_runtime(base: Path, target: Path, files: dict, binaries: dict) -> None: + shutil.copytree(base, target, copy_function=os.link) + for relative, blob in files.items(): + destination = target / relative + destination.unlink() # Never write through a hardlink to the baseline. + destination.write_bytes(blob) + for relative, digest in binaries.items(): + _require(_sha(target / relative) == digest, f"compiled artifact differs: {relative}") + _require((target / "bin/trtllm-serve").is_file(), "wheel lacks trtllm-serve entry point") + + +def _validate_harness(harness: Path, profile: dict) -> dict: + git = ["git", "-C", str(harness)] + head = _command(git + ["rev-parse", "HEAD"]) + if profile["harness_sha"]: + _require(head == profile["harness_sha"], "test harness SHA differs from profile") + _require( + not _command(git + ["status", "--porcelain", "--untracked-files=all"]), + "pinned test harness must be clean", + ) + test = harness / "tests/integration/defs/disaggregated/test_disaggregated.py" + tree = ast.parse(test.read_text()) + stress = next( + item + for item in tree.body + if isinstance(item, ast.FunctionDef) and item.name == "test_disaggregated_stress_test" + ) + configs = [] + for node in ast.walk(stress): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "TestConfig" + ): + values = {item.arg: ast.literal_eval(item.value) for item in node.keywords} + if values.get("test_desc") == "gpt_oss_120b_eagle_trtllm_stress": + configs.append(values) + _require( + len(configs) == 1 and configs[0].get("request_count") == profile["expected_requests"], + "test request count differs from profile", + ) + return { + "harness_sha": head, + "harness_diff": _command(git + ["diff", "HEAD"]), + "harness_files": { + relative: _sha(harness / relative) + for relative in ( + "tests/integration/defs/disaggregated/test_disaggregated.py", + "tests/integration/defs/conftest.py", + "tests/integration/defs/disaggregated/test_configs/disagg_config_ctxtp2_gentp2_gptoss_eagle_trtllm.yaml", + "tests/integration/lm_eval_configs/gsm8k_local.yaml", + "requirements-dev.txt", + ) + }, + } + + +def _gpu_inventory() -> list[list[str]]: + _require(bool(os.environ.get("SLURM_JOB_ID")), "an existing Slurm allocation is required") + _require( + os.environ.get("SLURM_JOB_NUM_NODES", os.environ.get("SLURM_NNODES")) == "1", + "exactly one Slurm node is required", + ) + _require( + os.environ.get("SLURM_GPUS_ON_NODE") == "8", + "allocation must provide eight GPUs on this node", + ) + rows = [ + line.split(", ") + for line in _command( + [ + "nvidia-smi", + "--query-gpu=uuid,name,memory.total,driver_version", + "--format=csv,noheader,nounits", + ] + ).splitlines() + ] + _require( + len(rows) == 8 and all(len(row) == 4 and "B200" in row[1] for row in rows), + "exactly eight B200 GPUs are required", + ) + return rows + + +def _models(models_root: Path) -> dict: + result = {} + for name in ("gpt-oss-120b", "gpt-oss-120b-Eagle3"): + directory = models_root / "gpt_oss" / name + config = directory / "config.json" + json.loads(config.read_text()) + indices = list(directory.glob("*.index.json")) + shards = { + directory / shard + for index in indices + for shard in json.loads(index.read_text()).get("weight_map", {}).values() + } + if not shards: + shards = set(directory.glob("*.safetensors")) | set( + directory.glob("pytorch_model*.bin") + ) + _require(bool(shards), f"no model shards found: {directory}") + for shard in shards: + with shard.open("rb") as stream: + _require(bool(stream.read(1)), f"empty model shard: {shard}") + result[name] = {"config_sha256": _sha(config), "shards": sorted(map(str, shards))} + return result + + +def _accuracy_inputs(args: argparse.Namespace) -> dict: + env = _environment(args.output / "runtime/base", args.harness, args.models_root, "preflight") + executables = { + name: shutil.which(name, path=env["PATH"]) for name in ("aiperf", "lm_eval", "python3") + } + _require(all(executables.values()), "aiperf, lm_eval and python3 must already be installed") + venv_bin = Path(sys.executable).parent + _require( + all(Path(path).parent == venv_bin for path in executables.values()), + "aiperf, lm_eval and python3 must resolve in the recorded venv/bin", + ) + template = args.harness / "tests/integration/lm_eval_configs/gsm8k_local.yaml" + dataset = args.models_root / "datasets/openai/gsm8k/main/test-00000-of-00001.parquet" + target = args.models_root / "gpt_oss/gpt-oss-120b" + code = """ +import importlib.metadata, json, sys +import pyarrow.parquet as pq +from transformers import AutoTokenizer +versions = {name: importlib.metadata.version(name) for name in ('lm_eval', 'aiperf')} +if versions != {'lm_eval': '0.4.10', 'aiperf': '0.8.0'}: + raise ValueError('lm_eval==0.4.10 and aiperf==0.8.0 are required: ' + str(versions)) +columns = pq.ParquetFile(sys.argv[2]).schema.names +if not {'question', 'answer'}.issubset(columns): + raise ValueError('GSM8K parquet lacks question/answer columns') +tokenizer = AutoTokenizer.from_pretrained(sys.argv[1], local_files_only=True, trust_remote_code=True) +if not tokenizer.encode('A/B tokenizer preflight'): + raise ValueError('tokenizer returned no tokens') +print('AB_INPUTS=' + json.dumps({'versions': versions, 'dataset_columns': columns})) +""" + log = args.output / "inputs-preflight.log" + try: + text = _command( + [sys.executable, "-c", code, str(target), str(dataset)], + cwd=args.output, + env=env, + stderr=subprocess.STDOUT, + ) + except subprocess.CalledProcessError as error: + log.write_text(error.output) + raise + log.write_text(text + "\n") + return { + "executables": executables, + "dataset": str(dataset), + "dataset_sha256": _sha(dataset), + "accuracy_template_sha256": _sha(template), + "tokenizer_files": { + str(path): _sha(path) for path in target.glob("*token*") if path.is_file() + }, + } + + +def _cleanliness(run_id: str) -> dict: + gpu = _command( + ["nvidia-smi", "--query-compute-apps=gpu_uuid,pid", "--format=csv,noheader"], timeout=10 + ) + owned = [] + marker = f"TLLM_AB_RUN_ID={run_id}".encode() + for entry in Path("/proc").iterdir(): + if not entry.name.isdigit(): + continue + try: + if entry.stat().st_uid == os.getuid() and marker in ( + entry / "environ" + ).read_bytes().split(b"\0"): + state = (entry / "stat").read_text().rsplit(")", 1)[1].split()[0] + if state != "Z": + owned.append(int(entry.name)) + except (FileNotFoundError, PermissionError, ProcessLookupError): + continue + return {"clean": not gpu and not owned, "gpu_processes": gpu.splitlines(), "owned_pids": owned} + + +def _runtime_preflight( + runtime: Path, expected_hashes: dict, env: dict, log: Path, dependency_versions: dict +) -> dict: + code = """ +import hashlib, importlib, importlib.metadata, json, pathlib, sys +import tensorrt_llm +import tensorrt_llm.bindings as bindings +files = {} +for relative in json.loads(sys.argv[1]): + name = relative.removesuffix('.py').replace('/', '.').removesuffix('.__init__') + module = importlib.import_module(name) + path = pathlib.Path(module.__file__).resolve() + files[relative] = {'path': str(path), 'sha256': hashlib.sha256(path.read_bytes()).hexdigest()} +print('AB_RUNTIME=' + json.dumps({'package': str(pathlib.Path(tensorrt_llm.__file__).resolve()), + 'runtime_files': files, 'bindings': str(pathlib.Path(bindings.__file__).resolve()), + 'versions': {name: importlib.metadata.version(name) for name in + ('tensorrt-llm', 'torch', 'aiperf', 'lm_eval', 'nixl-cu13', 'pytest')}})) +""" + try: + text = _command( + [sys.executable, "-c", code, json.dumps(list(expected_hashes))], + cwd=runtime, + env=env, + stderr=subprocess.STDOUT, + ) + except subprocess.CalledProcessError as error: + log.write_text(error.output) + raise + log.write_text(text + "\n") + data = json.loads( + next( + line.removeprefix("AB_RUNTIME=") + for line in text.splitlines() + if line.startswith("AB_RUNTIME=") + ) + ) + for field in ("package", "bindings"): + _require( + Path(data[field]).is_relative_to(runtime), f"{field} imported outside selected runtime" + ) + _require(set(data["runtime_files"]) == set(expected_hashes), "incomplete runtime file imports") + for relative, expected in expected_hashes.items(): + actual = data["runtime_files"][relative] + _require( + Path(actual["path"]) == runtime / relative, + f"{relative} imported outside selected runtime", + ) + _require( + actual["sha256"] == expected, f"imported Python differs from pinned source: {relative}" + ) + for name, version in dependency_versions.items(): + _require( + data["versions"][name] == version, f"runtime dependency differs from profile: {name}" + ) + return data + + +def _accounting(path: Path, expected: int = 60000) -> dict: + counts = {"records": 0, "valid": 0, "cancelled": 0, "errors": 0} + with path.open() as stream: + for line in stream: + if not line.strip(): + continue + record = json.loads(line) + _require( + isinstance(record, dict) and ("metrics" in record or "error" in record), + "unexpected per-request record schema", + ) + error = record.get("error") or {} + metadata = record.get("metadata") or {} + _require( + isinstance(error, dict) and isinstance(metadata, dict), "invalid record fields" + ) + cancelled = ( + metadata.get("was_cancelled") is True + or error.get("code") == 499 + or error.get("type") == "RequestCancellationError" + ) + counts["records"] += 1 + counts["cancelled" if cancelled else "errors" if error else "valid"] += 1 + _require( + counts["records"] == expected, + f"expected {expected} request records, got {counts['records']}", + ) + considered = counts["records"] - counts["cancelled"] + _require(considered > 0, "all request records are cancelled") + return {**counts, "error_rate": counts["errors"] / considered} + + +def _classify(trial: Path, rc: int, timed_out: bool, expected_requests: int = 60000) -> dict: + _require(not timed_out, "trial timed out") + _require(rc in (0, 1), f"pytest setup/interruption exit status {rc}") + cases = ET.parse(trial / "junit.xml").findall(".//testcase") + _require(len(cases) == 1 and cases[0].get("name") == CASE, "missing or extra selected tests") + case = cases[0] + _require( + case.find("skipped") is None and case.find("error") is None, "test skipped or setup failed" + ) + failed = case.find("failure") is not None + _require((rc == 1) == failed, "pytest exit status and JUnit disagree") + profiles = list((trial / "workspace").rglob("profile_export.jsonl")) + _require(len(profiles) == 1, "missing or ambiguous per-request accounting") + logs = { + name: list((trial / "workspace").rglob(name)) + for name in ("worker_ctx_0.log", "worker_gen_0.log", "disagg_server.log") + } + _require( + all(len(paths) == 1 and paths[0].stat().st_size for paths in logs.values()), + "full context/generation/proxy logs were not retained", + ) + accounting = _accounting(profiles[0], expected_requests) + _require( + failed or accounting["error_rate"] <= 0.05, "pytest passed despite excessive error rate" + ) + return { + "status": "fail" if failed else "pass", + "accounting": accounting, + "artifacts": { + "profile": str(profiles[0]), + **{name: str(paths[0]) for name, paths in logs.items()}, + }, + } + + +def _stop_group(process: subprocess.Popen) -> None: + for sig, grace in ((signal.SIGTERM, 15), (signal.SIGKILL, 5)): + try: + os.killpg(process.pid, sig) + except ProcessLookupError: + break + try: + process.wait(timeout=grace) + except subprocess.TimeoutExpired: + continue + process.wait(timeout=5) + + +def _run_trial(args: argparse.Namespace, arm: str, index: int, identity: dict, run_id: str) -> dict: + trial = args.output / f"{index:02d}-{arm}" + trial.mkdir() + runtime = args.output / "runtime" / arm + env = _environment(runtime, args.harness, args.models_root, run_id) + command = [ + sys.executable, + "-m", + "pytest", + "-vs", + SELECTOR, + "--workspace", + str(trial / "workspace"), + "--keep-workspace", + "--junitxml", + str(trial / "junit.xml"), + ] + result = { + "arm": arm, + "status": "invalid", + "command": command, + "started": time.time(), + "pytest_returncode": None, + "timed_out": False, + } + _save(trial / "result.json", result) + try: + result["dependencies"] = _validate_dependencies(identity["provenance"]) + before = _cleanliness(run_id) + result["before"] = before + _require(before["clean"], "live GPU workload or previous owned process before trial") + result["runtime"] = _runtime_preflight( + runtime, + identity["runtime_file_sha256"][arm], + env, + trial / "preflight.log", + args.profile["dependency_versions"], + ) + with (trial / "pytest.log").open("w") as log: + process = subprocess.Popen( + command, + cwd=args.harness / "tests/integration/defs", + env=env, + stdout=log, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + result["pid"] = process.pid + _save(trial / "result.json", result) + try: + result["pytest_returncode"] = process.wait(timeout=args.trial_timeout) + except subprocess.TimeoutExpired: + result["timed_out"] = True + finally: + _stop_group(process) + if result["pytest_returncode"] is None: + result["pytest_returncode"] = process.returncode + result.update( + _classify( + trial, + result["pytest_returncode"], + result["timed_out"], + args.profile["expected_requests"], + ) + ) + except (ValueError, OSError, ET.ParseError, subprocess.SubprocessError, StopIteration) as error: + result.update(status="invalid", reason=str(error)) + finally: + # Allow CUDA contexts to retire; never kill an unrelated GPU process. + try: + for _ in range(15): + result["cleanup"] = _cleanliness(run_id) + if result["cleanup"]["clean"]: + break + time.sleep(1) + except (OSError, subprocess.SubprocessError) as error: + result["cleanup"] = {"clean": False, "error": str(error)} + if not result["cleanup"]["clean"]: + result.update(status="invalid", reason="live owned or GPU processes after trial") + try: + _validate_dependencies(identity["provenance"]) + result["dependency_environment_unchanged"] = True + except (ValueError, OSError, subprocess.SubprocessError) as error: + result.update( + status="invalid", reason=str(error), dependency_environment_unchanged=False + ) + result["finished"] = time.time() + _save(trial / "result.json", result) + return result + + +def _main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + for flag in ("wheel", "provenance", "output", "models-root", "harness"): + parser.add_argument(f"--{flag}", required=True, type=Path) + parser.add_argument("--profile", type=Path) + parser.add_argument("--order", default="control,treatment,treatment,control") + parser.add_argument("--trial-timeout", type=int, default=12600) + parser.add_argument("--dry-plan", action="store_true") + args = parser.parse_args() + args.profile = load_profile(args.profile) + _require(sys.flags.optimize == 0, "driver must run without Python optimization") + order = args.order.split(",") + _require(order and all(arm in ("control", "treatment") for arm in order), "invalid arm order") + _require(0 < args.trial_timeout <= 12600, "timeout must be between 1 and 12600 seconds") + for field in ("wheel", "provenance", "output", "models_root", "harness"): + setattr(args, field, getattr(args, field).resolve()) + if args.dry_plan: + print( + json.dumps( + { + "selector": SELECTOR, + "order": order, + "control": args.profile["control"], + "treatment": args.profile["treatment"], + "profile": args.profile, + "timeout": args.trial_timeout, + "status": "plan_only_not_runtime_validation", + }, + indent=2, + ) + ) + return 0 + args.output.mkdir(parents=True, exist_ok=True) + _require(not any(args.output.iterdir()), "output directory must be new or empty") + run_id = str(uuid.uuid4()) + summary = { + "schema_version": 1, + "status": "invalid", + "trials": [], + "run_id": run_id, + "order": order, + "selector": SELECTOR, + "control": args.profile["control"], + "treatment": args.profile["treatment"], + "profile": args.profile, + "hostname": socket.gethostname(), + "driver_sha256": _sha(Path(__file__)), + "interpretation": "Observations only; two trials per arm cannot establish flake causality.", + } + _save(args.output / "summary.json", summary) + try: + provenance = json.loads(args.provenance.read_text()) + _validate_provenance(args.wheel, provenance, args.profile) + summary["dependencies"] = _validate_dependencies(provenance) + os.environ["PATH"] = f"{Path(sys.executable).parent}:{os.environ['PATH']}" + summary["gpus"] = _gpu_inventory() + summary["models"] = _models(args.models_root) + summary["accuracy_inputs"] = _accuracy_inputs(args) + summary["identity"] = _prepare_runtime(args, provenance) + summary["slurm"] = { + name: os.environ.get(name) + for name in ("SLURM_JOB_ID", "SLURM_STEP_ID", "SLURM_JOB_NODELIST") + } + _save(args.output / "manifest.json", summary) + for index, arm in enumerate(order, 1): + trial = _run_trial(args, arm, index, summary["identity"], run_id) + summary["trials"].append(trial) + _save(args.output / "summary.json", summary) + if not trial.get("cleanup", {}).get("clean"): + break + statuses = [trial["status"] for trial in summary["trials"]] + summary["status"] = ( + "invalid" + if "invalid" in statuses or len(statuses) != len(order) + else "fail" + if "fail" in statuses + else "pass" + ) + except (ValueError, OSError, subprocess.SubprocessError, StopIteration, SyntaxError) as error: + summary["reason"] = str(error) + finally: + summary["finished"] = time.time() + _save(args.output / "summary.json", summary) + return {"pass": 0, "fail": 1, "invalid": 2}[summary["status"]] + + +if __name__ == "__main__": + sys.exit(_main()) diff --git a/examples/disaggregated/slurm/stress_ab/summarize.py b/examples/disaggregated/slurm/stress_ab/summarize.py new file mode 100644 index 000000000000..0731cfb3db12 --- /dev/null +++ b/examples/disaggregated/slurm/stress_ab/summarize.py @@ -0,0 +1,88 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Reconcile the four fresh Slurm steps without turning missing data into a pass.""" + +import argparse +import csv +import json +from pathlib import Path + +ORDER = ["control", "treatment", "treatment", "control"] + + +def summarize(root: Path) -> dict: + """Return status and per-arm counts; no automatic causal claim is made.""" + with (root / "slurm-steps.tsv").open() as stream: + rows = list(csv.DictReader(stream, delimiter="\t")) + step_rows = {int(row["index"]): row for row in rows} + if len(step_rows) != len(rows) or set(step_rows) - {1, 2, 3, 4}: + raise ValueError("Duplicate or unexpected step records") + trials = [] + counts = {arm: dict.fromkeys(("pass", "fail", "invalid"), 0) for arm in set(ORDER)} + for index, arm in enumerate(ORDER, 1): + path = root / f"{index:02d}-{arm}" / "summary.json" + step = step_rows.get(index) + status = "invalid" + reason = "Missing Slurm step or trial summary" + if path.is_file() and step: + data = json.loads(path.read_text()) + details = data.get("trials", []) + if ( + len(details) == 1 + and details[0].get("arm") == arm + and step["arm"] == arm + and data.get("status") in ("pass", "fail", "invalid") + and details[0].get("status") == data["status"] + ): + status = data["status"] + expected_code = {"pass": 0, "fail": 1, "invalid": 2}[status] + if int(step["exit_code"]) != expected_code: + status = "invalid" + reason = "Slurm step exit disagrees with trial summary" + else: + reason = details[0].get("reason", "") + else: + reason = "Summary does not contain the expected single arm" + counts[arm][status] += 1 + trials.append( + {"index": index, "arm": arm, "status": status, "reason": reason, "summary": str(path)} + ) + status = "pass" + if any(item["status"] == "invalid" for item in trials): + status = "invalid" + elif any(item["status"] == "fail" for item in trials): + status = "fail" + return { + "schema_version": 1, + "status": status, + "counts": counts, + "trials": trials, + "interpretation": "Compare initiating worker errors and request accounting; no causal verdict is automatic.", + } + + +def main() -> int: + """Write a persistent matrix result and preserve non-passing exit status.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("run_root", type=Path) + args = parser.parse_args() + result = summarize(args.run_root) + (args.run_root / "summary.json").write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + return {"pass": 0, "fail": 1, "invalid": 2}[result["status"]] + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/disaggregated/slurm/stress_ab/test_launch.py b/examples/disaggregated/slurm/stress_ab/test_launch.py new file mode 100644 index 000000000000..636e75760ee0 --- /dev/null +++ b/examples/disaggregated/slurm/stress_ab/test_launch.py @@ -0,0 +1,151 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Exercise the real shell launcher using local scheduler stand-ins, without GPUs.""" + +import json +import os +import shutil +import subprocess +import sys +import tempfile +import unittest +from pathlib import Path + +SCRIPTS = Path(__file__).parent + + +class LaunchTests(unittest.TestCase): + def _launch(self, mode: str, time_limit: str = "16:00:00") -> tuple[int, list, dict]: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary).resolve() + binaries = root / "bin" + binaries.mkdir() + (binaries / "python3").symlink_to(sys.executable) + fake_commands = { + "git": "print('1' * 40)", + "scontrol": f"print('JobId=123 TimeLimit={time_limit}')", + "srun": """ +import json, os, sys +from pathlib import Path +args = sys.argv[1:] +output = Path(args[args.index('--output') + 1]) +arm = args[args.index('--order') + 1] +if os.environ['TEST_MODE'] == 'profile': + profile = Path(args[args.index('--profile') + 1]) + if json.loads(profile.read_text()).get('name') != 'current': + sys.exit(2) + script = next(value for value in args if value.endswith('/run_ab.py')) + if '/runner/' not in script or args[args.index('--harness') + 1].endswith('/runner'): + sys.exit(2) +with (output.parent / 'calls.txt').open('a') as stream: + stream.write(arm + '\\n') +if os.environ['TEST_MODE'] == 'launch_error': + sys.exit(1) +output.mkdir() +status = 'fail' if os.environ['TEST_MODE'] == 'control_failure' and arm == 'control' else 'pass' +(output / 'summary.json').write_text(json.dumps({ + 'status': status, 'trials': [{'arm': arm, 'status': status}]})) +sys.exit(1 if status == 'fail' else 0) +""", + } + for name, body in fake_commands.items(): + path = binaries / name + path.write_text(f"#!{sys.executable}\n{body}\n") + path.chmod(0o755) + harness = root / "harness" + scripts = harness / "examples/disaggregated/slurm/stress_ab" + scripts.mkdir(parents=True) + shutil.copy2(SCRIPTS / "summarize.py", scripts) + models = root / "models" + models.mkdir() + wheel, provenance = root / "wheel.whl", root / "provenance.json" + wheel.touch() + provenance.touch() + run = root / "run" + env = dict(os.environ) + for name in ("AB_PROFILE", "AB_RUNNER_ROOT", "AB_TRIAL_TIMEOUT"): + env.pop(name, None) + env.update( + AB_PROJECT_ROOT=str(root), + AB_HARNESS=str(harness), + AB_WHEEL=str(wheel), + AB_PROVENANCE=str(provenance), + AB_PYTHON=str(binaries / "python3"), + AB_MODELS_ROOT=str(models), + AB_RUN_ROOT=str(run), + AB_IMAGE="registry.example/repository:tag", + AB_IMAGE_DIGEST="sha256:" + "0" * 64, + SLURM_JOB_ID="123", + SLURM_JOB_NUM_NODES="1", + PATH=f"{binaries}:{os.environ['PATH']}", + TEST_MODE=mode, + ) + if mode == "profile": + profile = root / "profile.json" + profile.write_text(json.dumps({"name": "current"})) + runner = root / "runner" + runner_scripts = runner / "examples/disaggregated/slurm/stress_ab" + runner_scripts.mkdir(parents=True) + shutil.copy2(SCRIPTS / "summarize.py", runner_scripts) + env.update(AB_PROFILE=str(profile), AB_RUNNER_ROOT=str(runner)) + result = subprocess.run( + ["bash", str(SCRIPTS / "launch.slurm")], + env=env, + capture_output=True, + text=True, + timeout=15, + ) + calls = ( + (run / "calls.txt").read_text().splitlines() if (run / "calls.txt").exists() else [] + ) + summary = ( + json.loads((run / "summary.json").read_text()) + if (run / "summary.json").exists() + else {} + ) + return result.returncode, calls, summary + + def test_four_passing_steps(self): + code, calls, summary = self._launch("pass") + self.assertEqual(code, 0) + self.assertEqual(calls, ["control", "treatment", "treatment", "control"]) + self.assertEqual(summary["status"], "pass") + + def test_test_failure_continues_and_remains_nonzero(self): + code, calls, summary = self._launch("control_failure") + self.assertEqual(code, 1) + self.assertEqual(len(calls), 4) + self.assertEqual(summary["counts"]["control"]["fail"], 2) + + def test_profile_uses_separate_runner_and_common_test_checkout(self): + code, calls, summary = self._launch("profile") + self.assertEqual(code, 0) + self.assertEqual(len(calls), 4) + self.assertEqual(summary["status"], "pass") + + def test_srun_exit_one_without_evidence_stops_matrix(self): + code, calls, summary = self._launch("launch_error") + self.assertNotEqual(code, 0) + self.assertEqual(calls, ["control"]) + self.assertEqual(summary["status"], "invalid") + + def test_short_allocation_never_starts_a_trial(self): + code, calls, _ = self._launch("pass", "04:00:00") + self.assertEqual(code, 2) + self.assertEqual(calls, []) + + +if __name__ == "__main__": + unittest.main() diff --git a/examples/disaggregated/slurm/stress_ab/test_run_ab.py b/examples/disaggregated/slurm/stress_ab/test_run_ab.py new file mode 100644 index 000000000000..3fb318f36ef3 --- /dev/null +++ b/examples/disaggregated/slurm/stress_ab/test_run_ab.py @@ -0,0 +1,474 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""GPU-free contract tests: python -m unittest discover -s .""" + +import json +import os +import subprocess +import sys +import tempfile +import unittest +import venv +from pathlib import Path +from unittest.mock import patch + +import build_runtime +import run_ab + + +class TrialEvidenceTests(unittest.TestCase): + def setUp(self) -> None: + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.root = Path(self.temporary.name).resolve() + self.workspace = self.root / "workspace" + self.workspace.mkdir() + self.profile = self.workspace / "profile_export.jsonl" + for name in ("worker_ctx_0.log", "worker_gen_0.log", "disagg_server.log"): + (self.workspace / name).write_text("full worker output\n") + self.profile.write_text('{"metrics": {}, "error": null}\n' * 60000) + self._junit("") + + def _junit(self, children: str, name: str = run_ab.CASE, extra: str = "") -> None: + (self.root / "junit.xml").write_text( + f'{children}' + f"{extra}" + ) + + def test_success_requires_one_real_test_and_complete_records(self) -> None: + self.assertEqual(run_ab._classify(self.root, 0, False)["status"], "pass") + self.profile.write_text('{"metrics": {}}\n' * 59999) + with self.assertRaisesRegex(ValueError, "expected 60000"): + run_ab._classify(self.root, 0, False) + self.assertEqual( + run_ab._classify(self.root, 0, False, expected_requests=59999)["status"], "pass" + ) + + def test_known_shutdown_storm_is_failure_not_invalid_or_pass(self) -> None: + self.profile.write_text( + '{"error": {"code": 500, "type": "Internal Server Error"}}\n' * 54038 + + '{"error": {"code": 499, "type": "RequestCancellationError"}}\n' * 5962 + ) + self._junit('') + result = run_ab._classify(self.root, 1, False) + self.assertEqual(result["status"], "fail") + self.assertEqual(result["accounting"]["errors"], 54038) + self.assertEqual(result["accounting"]["cancelled"], 5962) + self.assertEqual(result["accounting"]["error_rate"], 1) + self._junit("") + with self.assertRaisesRegex(ValueError, "excessive error rate"): + run_ab._classify(self.root, 0, False) + + def test_skipped_missing_extra_setup_and_timeout_are_invalid(self) -> None: + for children, name, extra, rc, timeout in ( + ("", run_ab.CASE, "", 0, False), + ("", run_ab.CASE, "", 1, False), + ("", "wrong-test", "", 0, False), + ("", run_ab.CASE, '', 0, False), + ("", run_ab.CASE, "", 0, True), + ("", run_ab.CASE, "", 5, False), + ("", run_ab.CASE, "", 0, False), + ): + with self.subTest(children=children, name=name, rc=rc, timeout=timeout): + self._junit(children, name, extra) + with self.assertRaises(ValueError): + run_ab._classify(self.root, rc, timeout) + (self.root / "junit.xml").unlink() + with self.assertRaises(FileNotFoundError): + run_ab._classify(self.root, 0, False) + + def test_missing_logs_and_ambiguous_profiles_are_invalid(self) -> None: + (self.workspace / "worker_ctx_0.log").unlink() + with self.assertRaisesRegex(ValueError, "not retained"): + run_ab._classify(self.root, 0, False) + nested = self.workspace / "other" + nested.mkdir() + (nested / "profile_export.jsonl").write_text("{}") + with self.assertRaisesRegex(ValueError, "ambiguous"): + run_ab._classify(self.root, 0, False) + + def test_accounting_rejects_corruption_and_all_cancellation(self) -> None: + for record in ( + "truncated", + "{}", + "[]", + '{"error": "broken"}', + '{"metrics": {}, "metadata": {"was_cancelled": true}}', + ): + with self.subTest(record=record): + self.profile.write_text(record) + with self.assertRaises(ValueError): + run_ab._accounting(self.profile, expected=1) + + def test_provenance_rejects_wrong_wheel_source_and_digest(self) -> None: + wheel = self.root / "baseline.whl" + wheel.write_bytes(b"baseline") + provenance = { + "status": "built", + "source_unchanged_after_build": True, + "source_sha": build_runtime.CONTROL, + "clean_source": True, + "wheel": wheel.name, + "wheel_sha256": run_ab._sha(wheel), + "image": "image:pinned", + "image_digest": "sha256:" + "a" * 64, + "build_command": ["build"], + "build_log": "build.log", + } + with patch.dict( + os.environ, + { + "TLLM_AB_IMAGE": provenance["image"], + "TLLM_AB_IMAGE_DIGEST": provenance["image_digest"], + }, + ): + run_ab._validate_provenance(wheel, provenance) + for key, value in ( + ("status", "building"), + ("source_unchanged_after_build", False), + ("source_sha", "unknown"), + ("clean_source", False), + ("wheel_sha256", "wrong"), + ("image_digest", "latest"), + ): + with self.subTest(key=key), self.assertRaises(ValueError): + run_ab._validate_provenance(wheel, {**provenance, key: value}) + + def test_import_path_must_be_the_selected_runtime(self) -> None: + report = { + "package": "/wrong/tensorrt_llm/__init__.py", + "executor": "/wrong/executor.py", + "bindings": "/wrong/bindings.so", + "executor_sha256": "a", + "versions": {"aiperf": "0.8.0"}, + } + with patch.object(run_ab, "_command", return_value="AB_RUNTIME=" + json.dumps(report)): + with self.assertRaisesRegex(ValueError, "outside selected runtime"): + run_ab._runtime_preflight( + self.root, {build_runtime.EXECUTOR: "a"}, {}, self.root / "preflight.log", {} + ) + + def test_environment_does_not_import_checkout_or_user_overlay(self) -> None: + with patch.dict( + os.environ, {"PYTHONPATH": "/bad", "PYTHONHOME": "/bad", "PYTEST_ADDOPTS": "-k missing"} + ): + env = run_ab._environment(self.root, self.root / "harness", self.root / "models", "run") + self.assertEqual(env["PYTHONPATH"], str(self.root)) + self.assertNotIn("PYTHONHOME", env) + self.assertNotIn("PYTEST_ADDOPTS", env) + self.assertEqual(env["PYTHONNOUSERSITE"], "1") + self.assertEqual( + env["PATH"].split(os.pathsep)[:2], + [str(self.root / "bin"), str(Path(sys.executable).parent)], + ) + self.assertEqual(env["VIRTUAL_ENV"], sys.prefix) + + def test_contamination_blocks_next_arm_and_preserves_invalid_summary(self) -> None: + output = self.root / "run" + provenance = self.root / "provenance.json" + provenance.write_text("{}") + arguments = [ + "run_ab.py", + "--wheel", + str(self.root / "wheel.whl"), + "--provenance", + str(provenance), + "--output", + str(output), + "--models-root", + str(self.root), + "--harness", + str(self.root), + ] + with ( + patch.object(sys, "argv", arguments), + patch.object(run_ab, "_prepare_runtime", return_value={}), + patch.object(run_ab, "_validate_provenance", return_value=None), + patch.object(run_ab, "_gpu_inventory", return_value=[]), + patch.object(run_ab, "_models", return_value={}), + patch.object(run_ab, "_accuracy_inputs", return_value={}), + patch.object(run_ab, "_validate_dependencies", return_value={}), + patch.object( + run_ab, + "_run_trial", + return_value={ + "status": "invalid", + "pytest_returncode": 1, + "cleanup": {"clean": False, "owned_pids": [123]}, + }, + ) as trial, + ): + self.assertEqual(run_ab._main(), 2) + self.assertEqual(trial.call_count, 1) + summary = json.loads((output / "summary.json").read_text()) + self.assertEqual(summary["status"], "invalid") + self.assertEqual(summary["trials"][0]["pytest_returncode"], 1) + + @unittest.skipUnless(hasattr(os, "killpg"), "POSIX process groups required") + def test_timeout_cleanup_only_stops_owned_process_group(self) -> None: + owned = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(120)"], start_new_session=True + ) + unrelated = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(120)"], start_new_session=True + ) + try: + run_ab._stop_group(owned) + self.assertIsNotNone(owned.poll()) + self.assertIsNone(unrelated.poll()) + finally: + for child in (owned, unrelated): + if child.poll() is None: + child.kill() + child.wait(timeout=10) + + +class ProfileTests(unittest.TestCase): + def setUp(self) -> None: + temporary = tempfile.TemporaryDirectory() + self.addCleanup(temporary.cleanup) + self.root = Path(temporary.name).resolve() + self.profile = json.loads(json.dumps(build_runtime.HISTORICAL_PROFILE)) + self.profile.update( + name="current-test", + control="a" * 40, + treatment="b" * 40, + harness_sha="c" * 40, + non_runtime_files=[], + runtime_files=["tensorrt_llm/executor.py", "tensorrt_llm/adapter.py"], + ) + self.profile["dependency_versions"]["nixl-cu13"] = "1.4.0" + + def _load(self, profile: dict) -> dict: + path = self.root / "profile.json" + path.write_text(json.dumps(profile)) + return build_runtime.load_profile(path) + + def test_profile_requires_immutable_commits_and_python_allowlist(self) -> None: + self.assertEqual(self._load(self.profile), self.profile) + for key, value in ( + ("control", "main"), + ("harness_sha", None), + ("runtime_files", ["tensorrt_llm/../../escape.py"]), + ("runtime_files", ["tensorrt_llm/bindings.so"]), + ("runtime_files", []), + ("expected_requests", True), + ("non_runtime_files", ["cpp/engine.cpp"]), + ): + with self.subTest(key=key, value=value), self.assertRaises(ValueError): + self._load({**self.profile, key: value}) + + def test_source_validation_rejects_unlisted_compiled_changes_and_dependency_skew(self) -> None: + def git(*args): + return subprocess.check_output(["git", "-C", str(self.root), *args], text=True).strip() + + git("init", "--quiet") + git("config", "user.name", "Test") + git("config", "user.email", "test@example.invalid") + for relative in self.profile["runtime_files"]: + path = self.root / relative + path.parent.mkdir(exist_ok=True) + path.write_text("value = 'control'\n") + (self.root / "requirements-dev.txt").write_text( + "aiperf==0.8.0\nlm_eval[api]==0.4.10\nnixl-cu13==1.4.0\n" + ) + git("add", ".") + git("commit", "--quiet", "-m", "control") + self.profile["control"] = git("rev-parse", "HEAD") + for relative in self.profile["runtime_files"]: + (self.root / relative).write_text("value = 'treatment'\n") + git("commit", "--quiet", "-am", "treatment") + self.profile["treatment"] = git("rev-parse", "HEAD") + build_runtime.validate_source_profile(self.root, self.profile) + skew = { + **self.profile, + "dependency_versions": {**self.profile["dependency_versions"], "nixl-cu13": "1.3.1"}, + } + with self.assertRaisesRegex(ValueError, "control requirements"): + build_runtime.validate_source_profile(self.root, skew) + (self.root / "compiled.cpp").write_text("compiled change\n") + git("add", ".") + git("commit", "--quiet", "-m", "compiled") + self.profile["treatment"] = git("rev-parse", "HEAD") + with self.assertRaisesRegex(ValueError, "unexpected source changes"): + build_runtime.validate_source_profile(self.root, self.profile) + + def test_multifile_overlay_does_not_mutate_control_or_compiled_artifacts(self) -> None: + base = self.root / "base" + (base / "bin").mkdir(parents=True) + (base / "bin/trtllm-serve").write_text("console entry point") + files = {relative: b"treatment" for relative in self.profile["runtime_files"]} + for relative in files: + path = base / relative + path.parent.mkdir(exist_ok=True) + path.write_bytes(b"control") + binary = base / "bindings.so" + binary.write_bytes(b"compiled control") + target = self.root / "treatment" + run_ab._overlay_runtime(base, target, files, {"bindings.so": run_ab._sha(binary)}) + for relative in files: + self.assertEqual((base / relative).read_bytes(), b"control") + self.assertEqual((target / relative).read_bytes(), b"treatment") + self.assertNotEqual((base / relative).stat().st_ino, (target / relative).stat().st_ino) + self.assertEqual(binary.stat().st_ino, (target / "bindings.so").stat().st_ino) + report = { + "package": str(target / "tensorrt_llm/__init__.py"), + "bindings": str(target / "bindings.so"), + "versions": {}, + "runtime_files": { + relative: {"path": str(target / relative), "sha256": run_ab._sha(target / relative)} + for relative in files + }, + } + hashes = {relative: run_ab._sha(target / relative) for relative in files} + + def preflight(): + with patch.object(run_ab, "_command", return_value="AB_RUNTIME=" + json.dumps(report)): + return run_ab._runtime_preflight( + target, hashes, {}, self.root / "preflight.log", {} + ) + + preflight() + second = self.profile["runtime_files"][1] + report["runtime_files"][second]["path"] = str(base / second) + with self.assertRaisesRegex(ValueError, "outside selected runtime"): + preflight() + report["runtime_files"][second]["path"] = str(target / second) + report["runtime_files"][second]["sha256"] = "wrong" + with self.assertRaisesRegex(ValueError, "differs from pinned source"): + preflight() + + def test_current_profile_cannot_use_historical_provenance(self) -> None: + wheel = self.root / "wheel.whl" + wheel.write_bytes(b"baseline") + provenance = { + "status": "built", + "source_unchanged_after_build": True, + "source_sha": self.profile["control"], + "profile": self.profile, + "clean_source": True, + "wheel": wheel.name, + "wheel_sha256": run_ab._sha(wheel), + "build_log": "log", + "build_command": ["build"], + "image": "image:tag", + "image_digest": "sha256:" + "a" * 64, + } + with patch.dict( + os.environ, + { + "TLLM_AB_IMAGE": provenance["image"], + "TLLM_AB_IMAGE_DIGEST": provenance["image_digest"], + }, + ): + run_ab._validate_provenance(wheel, provenance, self.profile) + with self.assertRaisesRegex(ValueError, "profiles differ"): + run_ab._validate_provenance( + wheel, {**provenance, "profile": build_runtime.load_profile()}, self.profile + ) + with self.assertRaisesRegex(ValueError, "pinned control"): + run_ab._validate_provenance( + wheel, {**provenance, "source_sha": build_runtime.CONTROL}, self.profile + ) + + +class DependencyEnvironmentTests(unittest.TestCase): + def setUp(self) -> None: + self.temporary = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary.cleanup) + self.root = Path(self.temporary.name).resolve() + self.prefix = self.root / "baseline-venv" + venv.EnvBuilder(with_pip=False, symlinks=True).create(self.prefix) + self.python = self.prefix / "bin/python3" + self.site = ( + self.prefix + / f"lib/python{sys.version_info.major}.{sys.version_info.minor}/site-packages" + ) + for name, version in (("aiperf", "0.8.0"), ("lm_eval", "0.4.10")): + self._distribution(name, version) + entrypoints = self.site / f"{name}-{version}.dist-info/entry_points.txt" + entrypoints.write_text(f"[console_scripts]\n{name} = fake_cli:main\n") + (self.site / "fake_cli.py").write_text( + "def main():\n print('baseline CLI')\n return 0\n" + ) + self.entrypoints = build_runtime._prepare_cli_wrappers(self.python) + self.provenance = { + **build_runtime.runtime_identity(self.python), + "runtime_entrypoints": self.entrypoints, + } + + def _distribution(self, name: str, version: str) -> None: + metadata = self.site / f"{name}-{version}.dist-info/METADATA" + metadata.parent.mkdir() + metadata.write_text(f"Metadata-Version: 2.1\nName: {name}\nVersion: {version}\n") + + def _validate_in_venv(self) -> subprocess.CompletedProcess: + code = ( + "import json, sys; sys.path.insert(0, sys.argv[1]); " + "import run_ab; run_ab._validate_dependencies(json.loads(sys.argv[2]))" + ) + return subprocess.run( + [ + str(self.python), + "-I", + "-c", + code, + str(Path(run_ab.__file__).parent), + json.dumps(self.provenance), + ], + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + timeout=30, + ) + + def test_snapshot_preserves_venv_interpreter_and_duplicate_distributions(self) -> None: + self._distribution("duplicate_name", "1.0") + self._distribution("Duplicate-Name", "2.0") + snapshot = build_runtime.runtime_identity(self.python) + self.assertEqual(snapshot["runtime_python"], str(self.python)) + self.assertEqual(snapshot["runtime_prefix"], str(self.prefix)) + self.assertNotEqual(snapshot["runtime_python"], str(self.python.resolve())) + duplicates = [ + item for item in snapshot["runtime_distributions"] if item["name"] == "duplicate-name" + ] + self.assertEqual([item["version"] for item in duplicates], ["1.0", "2.0"]) + self.assertEqual(snapshot, build_runtime.runtime_identity(self.python)) + + def test_recorded_venv_passes_but_dependency_drift_is_rejected(self) -> None: + passed = self._validate_in_venv() + self.assertEqual(passed.returncode, 0, passed.stderr) + self._distribution("new-dependency", "1.0") + drift = self._validate_in_venv() + self.assertNotEqual(drift.returncode, 0) + self.assertIn("runtime_distributions", drift.stderr) + + def test_same_binary_outside_recorded_venv_is_rejected(self) -> None: + with self.assertRaisesRegex(ValueError, "recorded baseline venv|driver interpreter"): + run_ab._validate_dependencies(self.provenance) + + def test_cli_wrappers_use_recorded_python_and_detect_mutation(self) -> None: + for name in ("aiperf", "lm_eval"): + script = self.prefix / "bin" / name + self.assertEqual(script.read_text().splitlines()[0], f"#!{self.python}") + output = subprocess.check_output([str(script)], text=True, timeout=30) + self.assertEqual(output.strip(), "baseline CLI") + script.write_text(script.read_text() + "# changed\n") + drift = self._validate_in_venv() + self.assertNotEqual(drift.returncode, 0) + self.assertIn("baseline CLI changed", drift.stderr) + + +if __name__ == "__main__": + unittest.main() diff --git a/examples/disaggregated/slurm/stress_ab/test_summarize.py b/examples/disaggregated/slurm/stress_ab/test_summarize.py new file mode 100644 index 000000000000..f4ad6e948d12 --- /dev/null +++ b/examples/disaggregated/slurm/stress_ab/test_summarize.py @@ -0,0 +1,74 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Ensure the matrix cannot hide missing trials or failed Slurm steps.""" + +import json +import tempfile +import unittest +from pathlib import Path + +import summarize + + +class MatrixTests(unittest.TestCase): + def setUp(self) -> None: + temporary = tempfile.TemporaryDirectory() + self.addCleanup(temporary.cleanup) + self.root = Path(temporary.name) + self.steps = self.root / "slurm-steps.tsv" + self.steps.write_text("index\tarm\texit_code\n") + + def _trial(self, index: int, status: str, code: int) -> None: + arm = summarize.ORDER[index - 1] + directory = self.root / f"{index:02d}-{arm}" + directory.mkdir() + (directory / "summary.json").write_text( + json.dumps({"status": status, "trials": [{"arm": arm, "status": status}]}) + ) + with self.steps.open("a") as stream: + stream.write(f"{index}\t{arm}\t{code}\n") + + def test_control_failures_remain_visible_when_treatment_passes(self) -> None: + for index in range(1, 5): + failed = index in (1, 4) + self._trial(index, "fail" if failed else "pass", int(failed)) + result = summarize.summarize(self.root) + self.assertEqual(result["status"], "fail") + self.assertEqual(result["counts"]["control"]["fail"], 2) + self.assertEqual(result["counts"]["treatment"]["pass"], 2) + + def test_partial_matrix_cannot_pass(self) -> None: + self._trial(1, "pass", 0) + result = summarize.summarize(self.root) + self.assertEqual(result["status"], "invalid") + self.assertEqual(result["counts"]["treatment"]["invalid"], 2) + + def test_nonzero_step_cannot_be_overridden_by_passing_summary(self) -> None: + for index in range(1, 5): + self._trial(index, "pass", 143 if index == 4 else 0) + result = summarize.summarize(self.root) + self.assertEqual(result["status"], "invalid") + self.assertEqual(result["trials"][3]["status"], "invalid") + + def test_duplicate_step_is_rejected(self) -> None: + self._trial(1, "pass", 0) + with self.steps.open("a") as stream: + stream.write("1\tcontrol\t0\n") + with self.assertRaisesRegex(ValueError, "Duplicate"): + summarize.summarize(self.root) + + +if __name__ == "__main__": + unittest.main() diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 64122c1bfcde..b4a5317f2ac3 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -586,11 +586,6 @@ def __init__( # responses and flushing them at a synchronised point in the executor # loop avoids the mismatch. self._pending_transfer_responses: List[Tuple[int, LlmResponse]] = [] - # Requests with a buffered terminal response are terminated only after - # the synchronized flush has published that response. This preserves - # queue-backed client delivery while retaining normal termination as - # the single owner of resource and result-queue cleanup. - self._pending_response_terminations: List[LlmRequest] = [] # Same buffer-then-synced-flush pattern as _pending_transfer_responses # above: _handle_responses and _append_iter_stats are reached from # per-rank-divergent gates, so their tp_allgather collectives are @@ -1022,10 +1017,7 @@ def _end_transfer_and_maybe_terminate(self, request: LlmRequest): (request.py_request_id, response)) if self.async_transfer_manager.end_transfer(request): self.active_requests.remove(request) - if response: - self._pending_response_terminations.append(request) - else: - self._terminate_request(request) + self._terminate_request(request) return if self.async_transfer_manager.end_transfer(request): if transfer_failed: @@ -1043,15 +1035,11 @@ def _flush_pending_transfer_responses(self): """ responses = self._pending_transfer_responses self._pending_transfer_responses = [] - requests_to_terminate = self._pending_response_terminations - self._pending_response_terminations = [] if responses or self.enable_attention_dp: # Even when this rank has no responses we must participate in the # collective when ADP is enabled so that the other rank's gather # can complete. self._enqueue_responses(responses) - for request in requests_to_terminate: - self._terminate_request(request) def _handle_kv_transfer_timeouts_synced(self): """ADP-safe drain of the KV-transfer-timeout consensus collective. @@ -4003,8 +3991,7 @@ def _handle_disagg_cache_errors_synced(self): self.is_shutdown = True self._handle_errors(error_msg, requests=None, - charge_budget=False, - fatal_is_collective_aligned=True) + charge_budget=False) return if not (self.enable_attention_dp and self.dist.world_size != 1): @@ -4099,13 +4086,6 @@ def _executor_loop(self): scheduled_batch, iter_stats = self._prepare_and_schedule_batch() if scheduled_batch is None: - # _handle_disagg_cache_errors_synced() can buffer a - # non-fatal response before scheduling observes shutdown. - # Drain it before leaving the loop so the client does not - # wait for its own timeout. Scheduling shutdown is - # model-parallel synchronized, so every ADP rank reaches - # this collective together. - self._flush_pending_transfer_responses() self._event_loop_completed = True break @@ -4120,11 +4100,6 @@ def _executor_loop(self): scheduled_batch) self._pause_recompute_paused_requests(scheduled_batch) self._finalize_adp_dummy_allocation(False) - # _check_benchmark_disagg_gate() makes this retry decision - # with a model-parallel all-gather. Flush before retrying so - # a response buffered at the top of this pass is not held - # until the benchmark fill gate opens. - self._flush_pending_transfer_responses() continue if (self._mm_encoder_item_scheduling_enabled @@ -4281,6 +4256,7 @@ def _executor_loop(self): self.kv_cache_manager.update_context_resources( scheduled_batch) self._send_kv_async(scheduled_batch.all_requests()) + self._flush_pending_transfer_responses() self._handle_canceled_requests() finished_requests = self._handle_responses() @@ -4312,14 +4288,6 @@ def _executor_loop(self): self._kv_connector_terminate_requests() - # This is the one ADP-synchronized response flush for a - # completed executor-loop pass. It is deliberately outside - # ``if can_queue``: non-fatal errors can be buffered on an - # idle pass, and every ADP rank must enter the tp_gather in - # the same order. Keep it after all per-pass error handling - # so a response is not needlessly delayed to the next pass. - self._flush_pending_transfer_responses() - if self.enable_iter_perf_stats and sample_state is not None: self._process_iter_stats( finished_requests, self.active_requests, @@ -4933,7 +4901,6 @@ def _executor_loop_overlap(self): scheduled_batch, iter_stats = self._prepare_and_schedule_batch() if scheduled_batch is None: - self._flush_pending_transfer_responses() self._event_loop_completed = True break @@ -4948,7 +4915,6 @@ def _executor_loop_overlap(self): scheduled_batch) self._pause_recompute_paused_requests(scheduled_batch) self._finalize_adp_dummy_allocation(False) - self._flush_pending_transfer_responses() continue if (self._mm_encoder_item_scheduling_enabled @@ -7754,8 +7720,7 @@ def _handle_errors(self, error_msg: Optional[str] = None, *, requests: Optional[List[LlmRequest]] = None, - charge_budget: bool = True, - fatal_is_collective_aligned: bool = False) -> None: + charge_budget: bool = True) -> None: """Fail requests and optionally initiate shutdown on fatal errors. When ``charge_budget`` is True (the default), classifies the error @@ -7791,14 +7756,6 @@ def _handle_errors(self, """ error_responses: Dict[int, LlmResponse] = {} error_msg = error_msg or "error" - multi_rank_adp = (self.enable_attention_dp - and self.dist.world_size != 1) - # ``fatal_is_collective_aligned`` is set only by the synchronized caller - # (_handle_disagg_cache_errors_synced, after its world allreduce), which - # guarantees every ADP rank enters the fatal path in the same collective - # order. Do NOT infer it from ``self._fatal_error is not None``: a - # rank-local setter would then route into a tp_gather while peers are - # elsewhere, recreating the desync this path exists to prevent. budget_fatal = (self._error_budget.consume(error_msg) if charge_budget else False) @@ -7853,17 +7810,14 @@ def _handle_errors(self, client_id=getattr(item.request, 'client_id', None)))) - if not multi_rank_adp: + adp_collective_required = (self.enable_attention_dp + and self.dist.world_size != 1) + if waiting_responses or adp_collective_required: + self._enqueue_responses(waiting_responses) if waiting_responses: - self._enqueue_responses(waiting_responses) logger.info( f"Drained {len(waiting_responses)} queued requests " "on fatal error") - elif fatal_is_collective_aligned: - # Synchronized fatal: every ADP rank enters this drain gather - # together, so issue it even when this rank has no queued - # responses, to stay peer-aligned. - self._enqueue_responses(waiting_responses) failed_requests = (list(self.active_requests) if requests is None else requests) @@ -7881,65 +7835,12 @@ def _handle_errors(self, request for request in self.active_requests if request not in requests ] - defer_termination = False - publish_immediately = False - if is_fatal: - if multi_rank_adp: - if fatal_is_collective_aligned: - # Synchronized fatal: all ranks agree and march through the - # aligned publish/terminate collectives together, so publish - # here instead of diverging. - publish_immediately = True - else: - # A novel rank-local fatal can be observed by one ADP rank - # before its peers reach their next collective. Do not issue - # tp_gather here: it would desynchronize the group exactly - # like a rank-local non-fatal response. Skip the gather and - # raise below; the distributed supervisor tears down peers. - logger.error( - "Skipping rank-local fatal response gather under ADP") - else: - publish_immediately = True - elif multi_rank_adp: - # Under attention DP, _enqueue_responses performs a tp_gather - # that every rank must enter in the same order. Non-fatal errors - # (e.g. a failed disagg KV transfer) are observed by a single - # rank, so enqueueing here would pair this rank's gather against - # a different collective on its peers — typically the per-step - # tp_allgather(batch_size) — corrupting both sides. Buffer the - # responses instead; every rank flushes the buffer together at - # _flush_pending_transfer_responses. - self._pending_transfer_responses.extend(error_responses.items()) - self._pending_response_terminations.extend(failed_requests) - defer_termination = True - else: - # Without multi-rank ADP there is no rank-divergent collective, so - # publish the error immediately. - publish_immediately = True - - if publish_immediately: - # A fatal executor exits before the next normal flush; a non-ADP - # executor has no rank-divergent collective. Both can publish the - # current errors plus any previously buffered terminal responses. - pending = self._pending_transfer_responses - self._pending_transfer_responses = [] - pending_terminations = self._pending_response_terminations - self._pending_response_terminations = [] - self._enqueue_responses(pending + list(error_responses.items())) - for request in pending_terminations: - self._terminate_request(request) - - if not defer_termination: - for request in failed_requests: - self._terminate_request(request) + self._enqueue_responses(list(error_responses.items())) + for request in failed_requests: + self._terminate_request(request) if self._fatal_error is not None: self.executor_request_queue.enqueue_shutdown_request() - if multi_rank_adp and not fatal_is_collective_aligned: - # Only a novel rank-local fatal raises to force local teardown; - # a synchronized fatal has already published in lockstep and - # tears down every rank together via is_shutdown. - raise self._fatal_error def _terminate_request(self, request: LlmRequest) -> None: # Dummy requests don't participate in disagg KV cache transfers, @@ -8066,27 +7967,8 @@ def _enqueue_responses(self, responses: Iterable[Tuple[int, LlmResponse]]): gather_responses = [] if responses_list is not None: for resp in responses_list: - if resp is None: - continue - if not isinstance(resp, (list, tuple)): - # A non-list contribution means the TP collective - # was matched against a *different* collective on - # a peer rank (collectives pair by call order, not - # by type). A common mismatch partner is the - # per-step tp_allgather(batch_size), which makes - # the stray payload an int. The gathered data is - # corrupt on every rank, so fail fast instead of - # raising an opaque TypeError below or silently - # enqueueing garbage responses. - raise RuntimeError( - f"_enqueue_responses: TP collective desync — " - f"gathered a {type(resp).__name__} " - f"instead of a response list. A peer rank " - f"entered a different collective (e.g. " - f"tp_allgather(batch_size)); check for " - f"per-rank-divergent callers of " - f"_enqueue_responses.") - gather_responses.extend(resp) + if resp is not None: + gather_responses.extend(resp) responses = gather_responses logger.debug( f'after gather, rank = {self.dist.rank}, responses = {responses}') diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index f35b49961440..70d3bcee459b 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -2801,7 +2801,10 @@ def run_disaggregated_aiperf(config_file, raise finally: terminate(*ctx_workers, *gen_workers, disagg_server) - shutil.rmtree(work_dir, ignore_errors=True) + if os.environ.get("TLLM_DISAGG_STRESS_KEEP_LOGS") == "1": + logger.info(f"Preserving disaggregated stress logs: {work_dir}") + else: + shutil.rmtree(work_dir, ignore_errors=True) def run_accuracy_test(model_path: str, server_url: str, concurrency: int, diff --git a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py index 0294bcd0ecf4..08a6f4701761 100644 --- a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py +++ b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py @@ -362,7 +362,6 @@ def test_peer_buffer_poison_triggers_world_consistent_fatal_cleanup(monkeypatch) "Disagg KV cache transfer buffer is poisoned; process restart is required", requests=None, charge_budget=False, - fatal_is_collective_aligned=True, ) @@ -377,8 +376,6 @@ def test_preclassified_fatal_error_keeps_adp_response_collectives_aligned(): executor.executor_request_queue = Mock() executor.executor_request_queue.get_request_queue.return_value = raw_queue executor.active_requests = [] - executor._pending_transfer_responses = [] - executor._pending_response_terminations = [] executor.gather_all_responses = False executor.enable_attention_dp = True executor.dist = SimpleNamespace(rank=1, world_size=2) @@ -386,11 +383,7 @@ def test_preclassified_fatal_error_keeps_adp_response_collectives_aligned(): executor._terminate_request = Mock() PyExecutor._handle_errors( - executor, - "poisoned transfer buffer", - requests=None, - charge_budget=False, - fatal_is_collective_aligned=True, + executor, "poisoned transfer buffer", requests=None, charge_budget=False ) executor._error_budget.consume.assert_not_called()