From 3245fc3ecd76e2fb610f42f2422102e2430c28fe Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Thu, 27 Aug 2026 16:33:44 -0700 Subject: [PATCH 1/3] Revert "[None][fix] Keep ADP ranks in collective lockstep on request errors and fail fast on desync (#16687)" This reverts commit 09b77e8d5aa784e89bae9b044afd26b242e0ad51. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 144 ++---------------- .../test_disagg_inflight_cancel_gate.py | 9 +- 2 files changed, 14 insertions(+), 139 deletions(-) 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/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() From 59d25122cb0851ebee609c278bc72fa1a53a0ff8 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 18 Sep 2026 10:33:51 -0700 Subject: [PATCH 2/3] test: prepare reproducible B200 disaggregated stress A/B Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../disaggregated/slurm/stress_ab/README.md | 191 ++++++ .../slurm/stress_ab/build_runtime.py | 248 +++++++ .../slurm/stress_ab/launch.slurm | 137 ++++ .../disaggregated/slurm/stress_ab/run_ab.py | 618 ++++++++++++++++++ .../slurm/stress_ab/summarize.py | 88 +++ .../slurm/stress_ab/test_launch.py | 128 ++++ .../slurm/stress_ab/test_run_ab.py | 319 +++++++++ .../slurm/stress_ab/test_summarize.py | 74 +++ .../defs/disaggregated/test_disaggregated.py | 5 +- 9 files changed, 1807 insertions(+), 1 deletion(-) create mode 100644 examples/disaggregated/slurm/stress_ab/README.md create mode 100644 examples/disaggregated/slurm/stress_ab/build_runtime.py create mode 100644 examples/disaggregated/slurm/stress_ab/launch.slurm create mode 100644 examples/disaggregated/slurm/stress_ab/run_ab.py create mode 100644 examples/disaggregated/slurm/stress_ab/summarize.py create mode 100644 examples/disaggregated/slurm/stress_ab/test_launch.py create mode 100644 examples/disaggregated/slurm/stress_ab/test_run_ab.py create mode 100644 examples/disaggregated/slurm/stress_ab/test_summarize.py diff --git a/examples/disaggregated/slurm/stress_ab/README.md b/examples/disaggregated/slurm/stress_ab/README.md new file mode 100644 index 000000000000..e78e978eca0f --- /dev/null +++ b/examples/disaggregated/slurm/stress_ab/README.md @@ -0,0 +1,191 @@ + + +# 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 does not establish that #16687 caused the bug, +and the revert PR must not be merged as a production fix on this evidence. + +| 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. + +## 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. + +**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 + +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..7ff82e17ce28 --- /dev/null +++ b/examples/disaggregated/slurm/stress_ab/build_runtime.py @@ -0,0 +1,248 @@ +# 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.""" + +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" + + +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("--output", type=Path, required=True) + parser.add_argument("--build-root", type=Path, required=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") + 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") != CONTROL: + parser.error(f"source must be the exact control commit {CONTROL}") + 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") + command = [ + sys.executable, + str(source / "scripts/build_wheel.py"), + "--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.dry_run: + print(json.dumps({"source_sha": CONTROL, "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) + # CONTROL's out-of-tree staging omits this setup.py dependency. Keep source immutable. + 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) + build_log = output / "build.log" + manifest = { + "schema_version": 1, + "source_sha": CONTROL, + "clean_source": True, + "submodules": submodules.splitlines(), + "image": args.image, + "image_digest": args.image_digest, + "build_command": command, + "build_log": str(build_log), + "staging_addition": { + "source": str(extra_requirement), + "destination": str(staged_requirement), + "sha256": _sha256(extra_requirement), + }, + "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") + with build_log.open("w") as stream: + process = subprocess.Popen( + command, cwd=source, stdout=stream, stderr=subprocess.STDOUT, start_new_session=True + ) + try: + returncode = process.wait(timeout=args.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 + 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. + runtime_prefix = build_root / f"venv-{sys.version_info.major}.{sys.version_info.minor}" + runtime_python = runtime_prefix / "bin/python3" + try: + entrypoints = _prepare_cli_wrappers(runtime_python) + identity = runtime_identity(runtime_python) + 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..a37fa1e18d00 --- /dev/null +++ b/examples/disaggregated/slurm/stress_ab/launch.slurm @@ -0,0 +1,137 @@ +#!/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; } +for path in "$AB_PROJECT_ROOT" "$AB_HARNESS" "$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" "$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="$AB_HARNESS/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" <<'PY' +import json, os, socket, subprocess, sys +from datetime import datetime, timezone +from pathlib import Path +path, harness, image, digest, timeout = 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(), + '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"; 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..8bdbfc0eeb21 --- /dev/null +++ b/examples/disaggregated/slurm/stress_ab/run_ab.py @@ -0,0 +1,618 @@ +# 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.""" + +import argparse +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 runtime_identity + +CONTROL = "0f2c3a95f9415045bdf06a7230759475692483b6" +TREATMENT = "3245fc3ecd76e2fb610f42f2422102e2430c28fe" +EXECUTOR = "tensorrt_llm/_torch/pyexecutor/py_executor.py" +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) -> None: + _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") == CONTROL, "wheel source SHA is not the pinned control") + _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) + _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)] + changed = _command(git + ["diff", "--name-only", CONTROL, TREATMENT]) + _require( + set(changed.splitlines()) + == {EXECUTOR, "tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py"}, + "comparison has unexpected source changes", + ) + blobs = { + arm: subprocess.check_output(git + ["show", f"{sha}:{EXECUTOR}"]) + for arm, sha in (("control", CONTROL), ("treatment", 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) + _require( + (base / EXECUTOR).read_bytes() == blobs["control"], "wheel Python is not control source" + ) + _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, blob in blobs.items(): + target = base.parent / arm + shutil.copytree(base, target, copy_function=os.link) + destination = target / EXECUTOR + destination.unlink() # Other files are hardlinked; never write through to the baseline. + destination.write_bytes(blob) + for relative, digest in binaries.items(): + _require( + _sha(target / relative) == digest, f"compiled artifact differs in {arm}: {relative}" + ) + _require( + (target / "bin" / "trtllm-serve").is_file(), "wheel lacks trtllm-serve entry point" + ) + return { + "provenance": provenance, + "install_command": install, + "binaries": binaries, + "executor_sha256": {arm: hashlib.sha256(blob).hexdigest() for arm, blob in blobs.items()}, + "harness_sha": _command(git + ["rev-parse", "HEAD"]), + "harness_diff": _command(git + ["diff", "HEAD"]), + "harness_files": { + str(path.relative_to(args.harness)): _sha(path) + for path in [ + args.harness / "tests/integration/defs/disaggregated/test_disaggregated.py", + args.harness / "tests/integration/defs/conftest.py", + ] + }, + } + + +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_hash: str, env: dict, log: Path) -> dict: + code = """ +import hashlib, importlib.metadata, json, pathlib +import tensorrt_llm +import tensorrt_llm._torch.pyexecutor.py_executor as executor +import tensorrt_llm.bindings as bindings +print('AB_RUNTIME=' + json.dumps({'package': str(pathlib.Path(tensorrt_llm.__file__).resolve()), + 'executor': str(pathlib.Path(executor.__file__).resolve()), + 'executor_sha256': hashlib.sha256(pathlib.Path(executor.__file__).read_bytes()).hexdigest(), + 'bindings': str(pathlib.Path(bindings.__file__).resolve()), + 'versions': {name: importlib.metadata.version(name) for name in + ('tensorrt-llm', 'torch', 'aiperf', 'nixl-cu13', 'pytest')}})) +""" + try: + text = _command( + [sys.executable, "-c", code], 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", "executor", "bindings"): + _require( + Path(data[field]).is_relative_to(runtime), f"{field} imported outside selected runtime" + ) + _require( + data["executor_sha256"] == expected_hash, "imported executor differs from pinned source" + ) + _require(data["versions"]["aiperf"] == "0.8.0", "AIPerf must match the verified 0.8.0 schema") + 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) -> 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]) + _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["executor_sha256"][arm], env, trial / "preflight.log" + ) + 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"])) + 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("--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() + _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": CONTROL, + "treatment": TREATMENT, + "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": CONTROL, + "treatment": TREATMENT, + "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()) + 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) 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..f000230822d6 --- /dev/null +++ b/examples/disaggregated/slurm/stress_ab/test_launch.py @@ -0,0 +1,128 @@ +# 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] +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) + 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, + ) + 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_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..21637d345fbe --- /dev/null +++ b/examples/disaggregated/slurm/stress_ab/test_run_ab.py @@ -0,0 +1,319 @@ +# 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) + + 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": run_ab.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, "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, "_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 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/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, From 12ddc952a8db03496b960b583f8b2f7f719df270 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Fri, 18 Sep 2026 11:43:09 -0700 Subject: [PATCH 3/3] test: support pinned current-head stress A/B profiles Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../disaggregated/slurm/stress_ab/README.md | 69 +++++- .../slurm/stress_ab/build_runtime.py | 221 +++++++++++++++--- .../slurm/stress_ab/launch.slurm | 22 +- .../disaggregated/slurm/stress_ab/run_ab.py | 203 +++++++++++----- .../slurm/stress_ab/test_launch.py | 23 ++ .../slurm/stress_ab/test_run_ab.py | 159 ++++++++++++- 6 files changed, 599 insertions(+), 98 deletions(-) diff --git a/examples/disaggregated/slurm/stress_ab/README.md b/examples/disaggregated/slurm/stress_ab/README.md index e78e978eca0f..931b25b6485f 100644 --- a/examples/disaggregated/slurm/stress_ab/README.md +++ b/examples/disaggregated/slurm/stress_ab/README.md @@ -3,10 +3,11 @@ # 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 does not establish that #16687 caused the bug, +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. -| Identity | Pinned revision | +| Historical default identity | Pinned revision | | --- | --- | | Control A: historical source with #16687 | `0f2c3a95f9415045bdf06a7230759475692483b6` | | Treatment B: exact executor revert | `3245fc3ecd76e2fb610f42f2422102e2430c28fe` | @@ -19,6 +20,59 @@ 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 @@ -88,6 +142,13 @@ 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 @@ -132,6 +193,10 @@ 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" \ diff --git a/examples/disaggregated/slurm/stress_ab/build_runtime.py b/examples/disaggregated/slurm/stress_ab/build_runtime.py index 7ff82e17ce28..81e5989c9b89 100644 --- a/examples/disaggregated/slurm/stress_ab/build_runtime.py +++ b/examples/disaggregated/slurm/stress_ab/build_runtime.py @@ -14,6 +14,8 @@ """Build the pinned control once and record provenance for both A/B arms.""" +from __future__ import annotations + import argparse import hashlib import json @@ -27,6 +29,95 @@ 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: @@ -109,8 +200,11 @@ 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) @@ -119,13 +213,15 @@ def main() -> int: 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") != CONTROL: - parser.error(f"source must be the exact control commit {CONTROL}") + 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") @@ -137,9 +233,21 @@ def main() -> int: 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") - command = [ + # 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", @@ -152,8 +260,25 @@ def main() -> int: "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": CONTROL, "build_command": command}, indent=2)) + 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") @@ -164,46 +289,79 @@ def main() -> int: parser.error("container launcher must set matching TLLM_AB_IMAGE and TLLM_AB_IMAGE_DIGEST") output.mkdir(parents=True) build_root.mkdir(parents=True) - # CONTROL's out-of-tree staging omits this setup.py dependency. Keep source immutable. - 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 = [] + 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": CONTROL, + "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_addition": { - "source": str(extra_requirement), - "destination": str(staged_requirement), - "sha256": _sha256(extra_requirement), - }, + "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: - process = subprocess.Popen( - command, cwd=source, stdout=stream, stderr=subprocess.STDOUT, start_new_session=True - ) - try: - returncode = process.wait(timeout=args.timeout) - except subprocess.TimeoutExpired: - os.killpg(process.pid, signal.SIGTERM) + 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: - process.wait(timeout=30) + returncode = process.wait(timeout=timeout) except subprocess.TimeoutExpired: - os.killpg(process.pid, signal.SIGKILL) - process.wait() - returncode = 124 + 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")) @@ -215,11 +373,18 @@ def main() -> int: 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. - runtime_prefix = build_root / f"venv-{sys.version_info.major}.{sys.version_info.minor}" - runtime_python = runtime_prefix / "bin/python3" 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 ): diff --git a/examples/disaggregated/slurm/stress_ab/launch.slurm b/examples/disaggregated/slurm/stress_ab/launch.slurm index a37fa1e18d00..3f56d451d779 100644 --- a/examples/disaggregated/slurm/stress_ab/launch.slurm +++ b/examples/disaggregated/slurm/stress_ab/launch.slurm @@ -37,12 +37,20 @@ set -Eeuo pipefail [[ $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; } -for path in "$AB_PROJECT_ROOT" "$AB_HARNESS" "$AB_WHEEL" "$AB_PROVENANCE" "$AB_PYTHON" "$AB_MODELS_ROOT" "$AB_RUN_ROOT"; do +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" "$AB_WHEEL" "$AB_PROVENANCE" "$AB_PYTHON" "$AB_RUN_ROOT"; do +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 ]] || { @@ -65,23 +73,25 @@ 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="$AB_HARNESS/examples/disaggregated/slurm/stress_ab" +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" <<'PY' +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 = sys.argv[1:] +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', } @@ -112,7 +122,7 @@ for arm in control treatment treatment control; do "$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"; then + --order "$arm" --trial-timeout "$trial_timeout" ${profile_args[@]+"${profile_args[@]}"}; then step_status=0 else step_status=$? diff --git a/examples/disaggregated/slurm/stress_ab/run_ab.py b/examples/disaggregated/slurm/stress_ab/run_ab.py index 8bdbfc0eeb21..839e68afb1b8 100644 --- a/examples/disaggregated/slurm/stress_ab/run_ab.py +++ b/examples/disaggregated/slurm/stress_ab/run_ab.py @@ -13,7 +13,10 @@ # 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 @@ -28,11 +31,8 @@ import xml.etree.ElementTree as ET from pathlib import Path -from build_runtime import runtime_identity +from build_runtime import load_profile, runtime_identity, validate_source_profile -CONTROL = "0f2c3a95f9415045bdf06a7230759475692483b6" -TREATMENT = "3245fc3ecd76e2fb610f42f2422102e2430c28fe" -EXECUTOR = "tensorrt_llm/_torch/pyexecutor/py_executor.py" CASE = "test_disaggregated_stress_test[input8k-output1k-conc512-gpt_oss_120b_eagle_trtllm_stress]" SELECTOR = f"disaggregated/test_disaggregated.py::{CASE}" @@ -87,10 +87,17 @@ def _environment(runtime: Path, harness: Path, models: Path, run_id: str) -> dic return env -def _validate_provenance(wheel: Path, provenance: dict) -> None: +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") == CONTROL, "wheel source SHA is not the pinned control") + _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") @@ -138,21 +145,20 @@ def _validate_dependencies(provenance: dict) -> dict: def _prepare_runtime(args: argparse.Namespace, provenance: dict) -> dict: - _validate_provenance(args.wheel, provenance) + _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)] - changed = _command(git + ["diff", "--name-only", CONTROL, TREATMENT]) - _require( - set(changed.splitlines()) - == {EXECUTOR, "tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py"}, - "comparison has unexpected source changes", - ) + validate_source_profile(args.harness, args.profile) + harness = _validate_harness(args.harness, args.profile) blobs = { - arm: subprocess.check_output(git + ["show", f"{sha}:{EXECUTOR}"]) - for arm, sha in (("control", CONTROL), ("treatment", TREATMENT)) + 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) @@ -170,9 +176,11 @@ def _prepare_runtime(args: argparse.Namespace, provenance: dict) -> dict: ] with (args.output / "pip.log").open("w") as log: subprocess.run(install, check=True, stdout=log, stderr=subprocess.STDOUT, timeout=600) - _require( - (base / EXECUTOR).read_bytes() == blobs["control"], "wheel Python is not control source" - ) + 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) @@ -180,32 +188,73 @@ def _prepare_runtime(args: argparse.Namespace, provenance: dict) -> dict: 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, blob in blobs.items(): - target = base.parent / arm - shutil.copytree(base, target, copy_function=os.link) - destination = target / EXECUTOR - destination.unlink() # Other files are hardlinked; never write through to the baseline. - destination.write_bytes(blob) - for relative, digest in binaries.items(): - _require( - _sha(target / relative) == digest, f"compiled artifact differs in {arm}: {relative}" - ) - _require( - (target / "bin" / "trtllm-serve").is_file(), "wheel lacks trtllm-serve entry point" - ) + for arm, files in blobs.items(): + _overlay_runtime(base, base.parent / arm, files, binaries) return { "provenance": provenance, "install_command": install, "binaries": binaries, - "executor_sha256": {arm: hashlib.sha256(blob).hexdigest() for arm, blob in blobs.items()}, - "harness_sha": _command(git + ["rev-parse", "HEAD"]), + "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": { - str(path.relative_to(args.harness)): _sha(path) - for path in [ - args.harness / "tests/integration/defs/disaggregated/test_disaggregated.py", - args.harness / "tests/integration/defs/conftest.py", - ] + 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", + ) }, } @@ -334,22 +383,30 @@ def _cleanliness(run_id: str) -> dict: return {"clean": not gpu and not owned, "gpu_processes": gpu.splitlines(), "owned_pids": owned} -def _runtime_preflight(runtime: Path, expected_hash: str, env: dict, log: Path) -> dict: +def _runtime_preflight( + runtime: Path, expected_hashes: dict, env: dict, log: Path, dependency_versions: dict +) -> dict: code = """ -import hashlib, importlib.metadata, json, pathlib +import hashlib, importlib, importlib.metadata, json, pathlib, sys import tensorrt_llm -import tensorrt_llm._torch.pyexecutor.py_executor as executor 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()), - 'executor': str(pathlib.Path(executor.__file__).resolve()), - 'executor_sha256': hashlib.sha256(pathlib.Path(executor.__file__).read_bytes()).hexdigest(), - 'bindings': str(pathlib.Path(bindings.__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', 'nixl-cu13', 'pytest')}})) + ('tensorrt-llm', 'torch', 'aiperf', 'lm_eval', 'nixl-cu13', 'pytest')}})) """ try: text = _command( - [sys.executable, "-c", code], cwd=runtime, env=env, stderr=subprocess.STDOUT + [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) @@ -362,14 +419,24 @@ def _runtime_preflight(runtime: Path, expected_hash: str, env: dict, log: Path) if line.startswith("AB_RUNTIME=") ) ) - for field in ("package", "executor", "bindings"): + for field in ("package", "bindings"): _require( Path(data[field]).is_relative_to(runtime), f"{field} imported outside selected runtime" ) - _require( - data["executor_sha256"] == expected_hash, "imported executor differs from pinned source" - ) - _require(data["versions"]["aiperf"] == "0.8.0", "AIPerf must match the verified 0.8.0 schema") + _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 @@ -405,7 +472,7 @@ def _accounting(path: Path, expected: int = 60000) -> dict: return {**counts, "error_rate": counts["errors"] / considered} -def _classify(trial: Path, rc: int, timed_out: bool) -> dict: +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") @@ -426,7 +493,7 @@ def _classify(trial: Path, rc: int, timed_out: bool) -> dict: 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]) + accounting = _accounting(profiles[0], expected_requests) _require( failed or accounting["error_rate"] <= 0.05, "pytest passed despite excessive error rate" ) @@ -485,7 +552,11 @@ def _run_trial(args: argparse.Namespace, arm: str, index: int, identity: dict, r result["before"] = before _require(before["clean"], "live GPU workload or previous owned process before trial") result["runtime"] = _runtime_preflight( - runtime, identity["executor_sha256"][arm], env, trial / "preflight.log" + 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( @@ -506,7 +577,14 @@ def _run_trial(args: argparse.Namespace, arm: str, index: int, identity: dict, r _stop_group(process) if result["pytest_returncode"] is None: result["pytest_returncode"] = process.returncode - result.update(_classify(trial, result["pytest_returncode"], result["timed_out"])) + 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: @@ -537,10 +615,12 @@ 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") @@ -553,8 +633,9 @@ def _main() -> int: { "selector": SELECTOR, "order": order, - "control": CONTROL, - "treatment": TREATMENT, + "control": args.profile["control"], + "treatment": args.profile["treatment"], + "profile": args.profile, "timeout": args.trial_timeout, "status": "plan_only_not_runtime_validation", }, @@ -572,8 +653,9 @@ def _main() -> int: "run_id": run_id, "order": order, "selector": SELECTOR, - "control": CONTROL, - "treatment": TREATMENT, + "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.", @@ -581,6 +663,7 @@ def _main() -> int: _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() @@ -606,7 +689,7 @@ def _main() -> int: if "fail" in statuses else "pass" ) - except (ValueError, OSError, subprocess.SubprocessError) as error: + except (ValueError, OSError, subprocess.SubprocessError, StopIteration, SyntaxError) as error: summary["reason"] = str(error) finally: summary["finished"] = time.time() diff --git a/examples/disaggregated/slurm/stress_ab/test_launch.py b/examples/disaggregated/slurm/stress_ab/test_launch.py index f000230822d6..636e75760ee0 100644 --- a/examples/disaggregated/slurm/stress_ab/test_launch.py +++ b/examples/disaggregated/slurm/stress_ab/test_launch.py @@ -42,6 +42,13 @@ def _launch(self, mode: str, time_limit: str = "16:00:00") -> tuple[int, list, d 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': @@ -68,6 +75,8 @@ def _launch(self, mode: str, time_limit: str = "16:00:00") -> tuple[int, list, d 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), @@ -83,6 +92,14 @@ def _launch(self, mode: str, time_limit: str = "16:00:00") -> tuple[int, list, d 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, @@ -112,6 +129,12 @@ def test_test_failure_continues_and_remains_nonzero(self): 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) diff --git a/examples/disaggregated/slurm/stress_ab/test_run_ab.py b/examples/disaggregated/slurm/stress_ab/test_run_ab.py index 21637d345fbe..3fb318f36ef3 100644 --- a/examples/disaggregated/slurm/stress_ab/test_run_ab.py +++ b/examples/disaggregated/slurm/stress_ab/test_run_ab.py @@ -51,6 +51,9 @@ def test_success_requires_one_real_test_and_complete_records(self) -> None: 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( @@ -114,7 +117,7 @@ def test_provenance_rejects_wrong_wheel_source_and_digest(self) -> None: provenance = { "status": "built", "source_unchanged_after_build": True, - "source_sha": run_ab.CONTROL, + "source_sha": build_runtime.CONTROL, "clean_source": True, "wheel": wheel.name, "wheel_sha256": run_ab._sha(wheel), @@ -152,7 +155,9 @@ def test_import_path_must_be_the_selected_runtime(self) -> None: } 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, "a", {}, self.root / "preflight.log") + 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( @@ -189,6 +194,7 @@ def test_contamination_blocks_next_arm_and_preserves_invalid_summary(self) -> No 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={}), @@ -228,6 +234,155 @@ def test_timeout_cleanup_only_stops_owned_process_group(self) -> None: 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()