From f0e0920105380cbfcb381fed58a617c8c095ff0f Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Mon, 20 Jul 2026 03:39:59 +0000 Subject: [PATCH 01/21] [None][test] Add cache transceiver precheck driver for disagg perf sanity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Standalone MPI driver that transfers deterministic KV data ctx->gen through the cache transceiver using the SAME instance topology, cache_transceiver_config, and KV shape (from the model's config.json) as a disagg perf-sanity test's yaml. Supports asymmetric parallelism (pattern seeded per (request, global layer), constant along the KV-head axis) and attention DP (per-dp-rank request pairing), multi-instance (ctx x gen) pairing via file-based rendezvous + per-pair ZMQ sessions, and reports PASS/TIMEOUT/TRANSFER_ERROR/MISMATCH/INIT_ERROR per case with per-GPU bandwidth. Config resolution is a pure-stdlib module with CPU-only unit tests. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> [None][test] Use SELFKONLY cache type for MLA in cache transceiver precheck Real MLA serving creates the KV cache manager with CacheType.SELFKONLY (kv_factor=1: a single latent plane, no V) — see _torch/pyexecutor/_util.py. The precheck driver used CacheType.SELF unconditionally, doubling the per-token bytes for MLA models, which skews buffer chunking behavior and reported bandwidth relative to the real test. Fill/verify already read kv_factor from the buffer view shape, so no other change is needed. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../cache_transceiver_precheck/README.md | 77 ++ .../precheck_config.py | 346 ++++++ .../run_precheck.py | 997 ++++++++++++++++++ .../test_cache_transceiver_precheck_config.py | 216 ++++ 4 files changed, 1636 insertions(+) create mode 100644 tests/scripts/perf-sanity/cache_transceiver_precheck/README.md create mode 100644 tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py create mode 100644 tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py create mode 100644 tests/unittest/others/test_cache_transceiver_precheck_config.py diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md new file mode 100644 index 000000000000..5e44248b147d --- /dev/null +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md @@ -0,0 +1,77 @@ +# Disagg Perf-Sanity Cache-Transceiver Precheck + +A fast go/no-go network check that runs **inside every disaggregated +perf-sanity CI stage, before the real test**. It brings up one MPI instance +per ctx/gen server with the *same* topology, node placement, UCX environment, +and `cache_transceiver_config` as the real test, transfers deterministic KV +data ctx → gen through the cache transceiver, and verifies the received +bytes. If the transfer hangs, errors, or mismatches, the stage fails +immediately with a specific verdict — the (expensive) model bring-up never +starts. + +## Parity with the real test — by construction + +| Requirement | How it is guaranteed | +|---|---| +| Same UCX env vars (incl. the `unset UCX_TLS` cases) | `jenkins/scripts/perf/submit.py` builds the precheck commands from the **same** `ucx_tls_cmd` + `$CTX/GEN_WORKER_ENV_VARS` strings as the worker steps; `slurm_precheck_run.sh` sources the same `slurm_env_setup.sh` (the `UCX_TLS=tcp` fixup) as `slurm_run.sh`. | +| Same instance count / parallelism | One precheck `srun` per ctx/gen server with the same `-N/--ntasks/--ntasks-per-node/--mpi=pmix` and the same node slices (`-w`) as the real server steps (`slurm_launch_draft.sh`). TP/PP/CP/attention-DP come from the same `worker_config`. | +| Same transceiver config | `CacheTransceiverConfig(**yaml["worker_config"][role]["cache_transceiver_config"])` — the yaml block is passed through verbatim (backend, `max_tokens_in_buffer`, timeouts, ...). | + +Asymmetric layouts (ctx dep4 → gen dep16, ctx pp8 → gen tp32, ...) are +supported: data is seeded per (request, **global** layer) and constant along +the KV-head axis, so the receiver regenerates its expected slice locally under +any TP/PP resharding. KV shape (layers/heads/head_dim, MLA vs GQA, MTP nextn +layers) is read from the real model's `config.json` under `$LLM_MODELS_ROOT`. + +## Enabling + +Off by default. Either: + +- per test yaml (recommended for gradual rollout): + + ```yaml + cache_transceiver_precheck: + enabled: true + # optional overrides (defaults in precheck_config.PRECHECK_DEFAULTS): + # request_lengths: [1024, 8192] + # num_requests: 2 + # step_timeout_s: 900 + ``` + +- or globally at launch-script generation time: `TRTLLM_DISAGG_CT_PRECHECK=1`. + +## Failure output + +The sbatch log gets a summary block: per-instance verdicts +(`status/*.status`), the tail of each failing step log, and UCX red-flag lines +(`sw-emul` host-staged tcp fallback, UCX ERROR/WARN). Full artifacts under +`/cache_transceiver_precheck/`: + +``` +status/{ctx,gen}_.json # per-case detail: PASS/TIMEOUT/TRANSFER_ERROR/ + # MISMATCH/INIT_ERROR + reason + UCX env snapshot +status/{ctx,gen}_.status # one-line verdict (parsed by the launch script) +logs/{ctx,gen}_.log # merged per-rank logs (UCX_PROTO_INFO=used table) +csv/gen_/rank_*_recv.csv # C++ transceiver per-request bandwidth +``` + +## Manual runs + +```bash +# Inspect what a yaml resolves to (no GPU needed): +python3 run_precheck.py --role gen --server-idx 0 --dry-run \ + --config ../disaggregated/.yaml --work-dir /tmp/ct --llm-src + +# On a SLURM allocation: one srun per instance, e.g. ctx dep4 + gen dep8: +srun -N1 --ntasks=4 --mpi=pmix python3 run_precheck.py --role ctx --server-idx 0 \ + --config --work-dir --llm-src & +srun -N2 --ntasks=8 --mpi=pmix python3 run_precheck.py --role gen --server-idx 0 \ + --config --work-dir --llm-src & +wait +``` + +Unit tests (CPU-only): `pytest tests/unittest/others/test_cache_transceiver_precheck_config.py`. + +Lineage: adapted from `examples/disaggregated/slurm/cache_transceiver_test` +(the UCX tuning harness), reduced to a single sweep and extended to +asymmetric parallelism, attention DP, and multi-instance pairing. diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py new file mode 100644 index 000000000000..543f81e97791 --- /dev/null +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py @@ -0,0 +1,346 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Config resolution for the disagg cache-transceiver precheck. + +Pure-stdlib (+pyyaml) module: NO torch / tensorrt_llm imports, so it can be +unit-tested on a CPU-only machine and imported by the launch tooling. + +The precheck mirrors the *exact* disaggregated perf-sanity test configuration: +instance counts and per-side parallelism come from the same yaml +(`tests/scripts/perf-sanity/disaggregated/*.yaml`) the real test runs with, +and the transceiver is built from the yaml's own `cache_transceiver_config` +block, so config parity is by construction rather than by copying values. +""" + +import json +import os + +# Optional per-yaml overrides live under a `cache_transceiver_precheck:` block. +PRECHECK_DEFAULTS = { + # Request lengths to transfer. None -> derived: [1024, benchmark ISL], + # clamped to cache_transceiver_config.max_tokens_in_buffer and to + # max_request_length below. + "request_lengths": None, + # Cap for the DERIVED lengths only (long-ISL tests would otherwise move + # multi-GB per request in the precheck); explicit request_lengths in the + # yaml's cache_transceiver_precheck block are used as-is. + "max_request_length": 32768, + # Measured requests per (peer, request_length) after warmup. + "num_requests": 2, + "warmup_requests": 1, + # Upper bound on concurrently in-flight transfer pairs per chunk. + "max_concurrent_pairs": 8, + # signal.alarm / hang-detector budget for one chunk of transfers. + "chunk_timeout_s": 180, + # How long the gen side waits for the ctx rendezvous files (covers ctx + # KV-pool allocation + NIXL/UCX transceiver bring-up). + "rendezvous_timeout_s": 600, + # Verify the received bytes against the deterministic fill pattern. + "verify_data": True, +} + +# Fallback KV shape when the model directory cannot be resolved: the precheck +# still exercises the exact network path, just with a synthetic cache shape. +FALLBACK_KV_SHAPE = { + "num_layers": 32, + "num_kv_heads": 8, + "head_dim": 128, + "is_mla": False, + "source": "fallback", +} + + +def _side_cfg(cfg, role): + worker_config = cfg.get("worker_config", {}) or {} + side = worker_config.get(role, {}) or {} + if not side: + raise ValueError(f"worker_config.{role} missing in disagg yaml") + return side + + +def _parallel(side): + tp = int(side.get("tensor_parallel_size", 1)) + pp = int(side.get("pipeline_parallel_size", 1)) + cp = int(side.get("context_parallel_size", 1)) + adp = bool(side.get("enable_attention_dp", False)) + world = tp * pp * cp + # With attention DP the request-level data parallel width is the attention + # TP width (each dp rank owns whole requests / full KV heads). + dp_size = tp if adp else 1 + return { + "tp": tp, + "pp": pp, + "cp": cp, + "enable_attention_dp": adp, + "world_size": world, + "dp_size": dp_size, + } + + +def _spec_nextn(side): + spec = side.get("speculative_config") or {} + decoding_type = str(spec.get("decoding_type", "")).upper() + if decoding_type == "MTP": + return int(spec.get("num_nextn_predict_layers", 0) or 0) + return 0 + + +def resolve_plan(cfg, benchmark_mode="e2e"): + """Build the shared precheck plan both roles must agree on. + + `cfg` is the parsed disagg perf-sanity yaml. Returns a plain dict; the + `fingerprint` field is exchanged over the rendezvous channel so a ctx/gen + disagreement (e.g. mismatched yamls) fails fast with a clear error. + """ + hardware = cfg.get("hardware", {}) or {} + benchmark = cfg.get("benchmark", {}) or {} + + yaml_mode = str(benchmark.get("mode", "")) + if benchmark_mode == "gen_only" and "gen_only_no_context" in yaml_mode: + # No ctx workers at launch -> no KV transfer in the real test either. + return {"skip": True, "skip_reason": "gen_only_no_context mode has no KV transfer"} + + num_ctx_servers = int(hardware.get("num_ctx_servers", 0) or 0) + num_gen_servers = int(hardware.get("num_gen_servers", 0) or 0) + if num_ctx_servers < 1 or num_gen_servers < 1: + raise ValueError( + f"hardware.num_ctx_servers/num_gen_servers must be >= 1, got " + f"{num_ctx_servers}/{num_gen_servers}" + ) + gpus_per_node = int(hardware.get("gpus_per_node", 0) or 0) + if gpus_per_node < 1: + raise ValueError("hardware.gpus_per_node is required") + + ctx_side = _side_cfg(cfg, "ctx") + gen_side = _side_cfg(cfg, "gen") + ctx = _parallel(ctx_side) + gen = _parallel(gen_side) + + ctx_xcvr = ctx_side.get("cache_transceiver_config") or {} + gen_xcvr = gen_side.get("cache_transceiver_config") or {} + if not ctx_xcvr.get("backend") and not gen_xcvr.get("backend"): + return {"skip": True, "skip_reason": "no cache_transceiver_config.backend in yaml"} + if ctx_xcvr.get("backend") != gen_xcvr.get("backend"): + raise ValueError( + f"ctx/gen cache_transceiver_config.backend mismatch: " + f"{ctx_xcvr.get('backend')} vs {gen_xcvr.get('backend')}" + ) + + ctx_nextn = _spec_nextn(ctx_side) + gen_nextn = _spec_nextn(gen_side) + + knobs = dict(PRECHECK_DEFAULTS) + knobs.update(cfg.get("cache_transceiver_precheck", {}) or {}) + + tokens_per_block = int( + (ctx_side.get("kv_cache_config") or {}).get("tokens_per_block", 32) or 32 + ) + max_tokens_in_buffer = ctx_xcvr.get("max_tokens_in_buffer") or gen_xcvr.get( + "max_tokens_in_buffer" + ) + + req_lens = knobs["request_lengths"] + if not req_lens: + isl = int(benchmark.get("input_length", 1024) or 1024) + req_lens = [1024, min(isl, int(knobs["max_request_length"]))] + if max_tokens_in_buffer: + req_lens = [min(int(r), int(max_tokens_in_buffer)) for r in req_lens] + req_lens = sorted({max(int(r), tokens_per_block) for r in req_lens}) + + # Transfer pairs: with attention DP each dp rank owns whole requests, so + # cover every dp rank on both sides; without DP a single request already + # involves every rank of the instance (KV is sharded across TP/PP). + n_pairs = max(ctx["dp_size"], gen["dp_size"], 1) + chunk_size = max(1, min(n_pairs, int(knobs["max_concurrent_pairs"]))) + + plan = { + "skip": False, + "num_ctx_servers": num_ctx_servers, + "num_gen_servers": num_gen_servers, + "gpus_per_node": gpus_per_node, + "ctx": ctx, + "gen": gen, + "ctx_cache_transceiver_config": ctx_xcvr, + "gen_cache_transceiver_config": gen_xcvr, + "ctx_num_nextn_predict_layers": ctx_nextn, + "gen_num_nextn_predict_layers": gen_nextn, + "ctx_kv_dtype": str((ctx_side.get("kv_cache_config") or {}).get("dtype", "auto")), + "gen_kv_dtype": str((gen_side.get("kv_cache_config") or {}).get("dtype", "auto")), + "tokens_per_block": tokens_per_block, + "request_lengths": req_lens, + "num_requests": int(knobs["num_requests"]), + "warmup_requests": int(knobs["warmup_requests"]), + "n_pairs": n_pairs, + "chunk_size": chunk_size, + "chunk_timeout_s": int(knobs["chunk_timeout_s"]), + "rendezvous_timeout_s": int(knobs["rendezvous_timeout_s"]), + "verify_data": bool(knobs["verify_data"]), + } + plan["fingerprint"] = plan_fingerprint(plan) + return plan + + +def plan_fingerprint(plan): + """Stable string both sides must agree on before transferring.""" + keys = ( + "num_ctx_servers", + "num_gen_servers", + "ctx", + "gen", + "ctx_cache_transceiver_config", + "gen_cache_transceiver_config", + "ctx_num_nextn_predict_layers", + "gen_num_nextn_predict_layers", + "tokens_per_block", + "request_lengths", + "num_requests", + "warmup_requests", + "n_pairs", + "chunk_size", + ) + return json.dumps({k: plan[k] for k in keys}, sort_keys=True) + + +def side_plan(plan, role): + """Per-role view: this role's parallelism + transceiver/kv config.""" + other = "gen" if role == "ctx" else "ctx" + return { + "role": role, + "parallel": plan[role], + "peer_parallel": plan[other], + "cache_transceiver_config": plan[f"{role}_cache_transceiver_config"], + "kv_dtype": plan[f"{role}_kv_dtype"], + "num_nextn_predict_layers": plan[f"{role}_num_nextn_predict_layers"], + "num_peers": plan["num_gen_servers" if role == "ctx" else "num_ctx_servers"], + } + + +def pair_participates(plan, role, tp_rank, pair_idx): + """Whether this rank takes part in transfer pair `pair_idx`. + + Attention-DP side: pair k belongs to dp rank k % dp_size (every pp stage + of that dp rank participates). Non-DP side: KV is sharded across the whole + instance, so every rank participates in every pair. + """ + side = plan[role] + if not side["enable_attention_dp"]: + return True + return tp_rank == pair_idx % side["dp_size"] + + +def owned_pairs(plan, role, tp_rank, chunk_pairs): + return [k for k in chunk_pairs if pair_participates(plan, role, tp_rank, k)] + + +def max_owned_per_chunk(plan, role): + """Max concurrently owned pairs per rank (KV pool sizing).""" + side = plan[role] + chunk = plan["chunk_size"] + if not side["enable_attention_dp"]: + return chunk + return (chunk + side["dp_size"] - 1) // side["dp_size"] + + +def chunks(plan): + """Pair indices grouped into concurrency-bounded chunks.""" + pairs = list(range(plan["n_pairs"])) + size = plan["chunk_size"] + return [pairs[i : i + size] for i in range(0, len(pairs), size)] + + +# --------------------------------------------------------------------------- # +# Model KV shape resolution +# --------------------------------------------------------------------------- # +def _load_model_path_dict(llm_src): + """Import MODEL_PATH_DICT from tests/integration/defs/perf/_model_paths.py.""" + path = os.path.join(llm_src, "tests", "integration", "defs", "perf", "_model_paths.py") + namespace = {} + with open(path) as f: + exec(compile(f.read(), path, "exec"), namespace) # noqa: S102 - repo-local constants file + return namespace["MODEL_PATH_DICT"] + + +def resolve_model_dir(cfg, llm_src=None, llm_models_root=None): + """Resolve the local model directory the same way test_perf_sanity does. + + metadata.model_name -> MODEL_PATH_DICT -> $LLM_MODELS_ROOT/, + falling back to metadata.model_dir_name under LLM_MODELS_ROOT. Returns + None when nothing resolvable exists (precheck then uses FALLBACK_KV_SHAPE). + """ + llm_models_root = llm_models_root or os.environ.get("LLM_MODELS_ROOT", "") + if not llm_models_root: + return None + metadata = cfg.get("metadata", {}) or {} + candidates = [] + model_name = metadata.get("model_name", "") + if llm_src and model_name: + try: + rel = _load_model_path_dict(llm_src).get(model_name) + if rel: + candidates.append(os.path.join(llm_models_root, rel)) + except (OSError, KeyError, SyntaxError): + pass + if metadata.get("model_dir_name"): + candidates.append(os.path.join(llm_models_root, metadata["model_dir_name"])) + for cand in candidates: + if os.path.isfile(os.path.join(cand, "config.json")): + return cand + return None + + +def model_kv_shape(model_dir): + """KV cache shape (per-token layout) from the model's config.json. + + Handles MLA checkpoints (kv_lora_rank present: one latent 'head' of + kv_lora_rank + qk_rope_head_dim) and GQA/MHA. `num_layers` excludes the + MTP nextn layers -- those are added by the KV cache manager via + spec_config, mirroring real serving. + """ + if not model_dir: + return dict(FALLBACK_KV_SHAPE) + try: + with open(os.path.join(model_dir, "config.json")) as f: + hf_cfg = json.load(f) + except (OSError, json.JSONDecodeError): + return dict(FALLBACK_KV_SHAPE) + if isinstance(hf_cfg.get("text_config"), dict): + hf_cfg = hf_cfg["text_config"] + + num_layers = hf_cfg.get("num_hidden_layers") + if num_layers is None: + return dict(FALLBACK_KV_SHAPE) + + if hf_cfg.get("kv_lora_rank"): # MLA (DeepSeek-family, Kimi K2, ...) + return { + "num_layers": int(num_layers), + "num_kv_heads": 1, + "head_dim": int(hf_cfg["kv_lora_rank"]) + int(hf_cfg.get("qk_rope_head_dim", 0)), + "is_mla": True, + "source": "config.json (MLA)", + } + + num_heads = hf_cfg.get("num_attention_heads", 1) + num_kv_heads = hf_cfg.get("num_key_value_heads", num_heads) + head_dim = hf_cfg.get("head_dim") + if not head_dim: + hidden = hf_cfg.get("hidden_size") + head_dim = hidden // num_heads if hidden and num_heads else 128 + return { + "num_layers": int(num_layers), + "num_kv_heads": int(num_kv_heads), + "head_dim": int(head_dim), + "is_mla": False, + "source": "config.json", + } diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py new file mode 100644 index 000000000000..58a45e30c954 --- /dev/null +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -0,0 +1,997 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Per-rank driver for the disagg perf-sanity cache-transceiver precheck. + +Runs BEFORE the real disaggregated perf-sanity servers, with the SAME +instance layout (one MPI world per ctx/gen server, same node/GPU topology), +the SAME UCX environment (the launch script reuses the worker env prefix +verbatim), and the SAME `cache_transceiver_config` (built directly from the +test's disagg yaml). It allocates a KV cache shaped like the real model, +transfers deterministic data ctx -> gen through the transceiver for every +(ctx server, gen server) pair, verifies the received bytes, and fails the +stage with a specific error (TIMEOUT / TRANSFER_ERROR / MISMATCH / ...) +before any model is loaded, so network/UCX misconfiguration is caught in +minutes instead of after a full model bring-up. + +Asymmetric parallelism (e.g. ctx dep4 -> gen dep16, ctx pp8 -> gen tp32) is +supported: the fill pattern is seeded per (request, GLOBAL layer) and is +constant along the KV-head axis, so any TP resharding or PP re-splitting on +the receiving side regenerates the identical expected bytes locally. (This +deliberately cannot detect head-permutation bugs -- it is a network +precheck, not a transceiver-correctness test.) + +Derived from examples/disaggregated/slurm/cache_transceiver_test (the UCX +tuning harness), reduced to a single go/no-go sweep and extended to +asymmetric layouts, attention DP, and multi-instance pairing. + +Rendezvous is file-based under --work-dir (a shared filesystem): each ctx +leader binds one ZMQ REP socket per gen peer and publishes host:port in +rendezvous/ctx{ci}_gen{gj}.addr; gen leaders connect with REQ sockets. Only +tiny control payloads travel over ZMQ -- KV data goes through the +transceiver under test. +""" + +import argparse +import json +import os +import pickle +import signal +import socket +import sys +import time + +CUR_DIR = os.path.dirname(os.path.abspath(__file__)) +if CUR_DIR not in sys.path: + sys.path.insert(0, CUR_DIR) + +import precheck_config as pcfg # noqa: E402 + +# Request-id strides (must keep rids unique across the whole precheck). +RID_PAIR_STRIDE = 1 +RID_MAX_PAIRS = 4096 +RID_REP_STRIDE = RID_MAX_PAIRS +RID_MAX_REPS = 64 +RID_LEN_STRIDE = RID_REP_STRIDE * RID_MAX_REPS +RID_MAX_LENS = 64 +RID_PEER_STRIDE = RID_LEN_STRIDE * RID_MAX_LENS + + +class _Timeout(Exception): + pass + + +class _TransferError(Exception): + pass + + +class _PeerAbort(Exception): + pass + + +def _alarm_handler(signum, frame): + raise _Timeout() + + +def make_rid(ctx_idx, gen_idx, num_ctx, li, rep, pair): + peer = gen_idx * num_ctx + ctx_idx + return 1 + peer * RID_PEER_STRIDE + li * RID_LEN_STRIDE + rep * RID_REP_STRIDE + pair + + +def seed_for(rid, global_layer): + # Per (request, GLOBAL layer); rank-independent so any receiving layout + # can regenerate its local slice. + return (rid * 1_000_003 + global_layer * 31) & 0x7FFFFFFF + + +# --------------------------------------------------------------------------- # +# KV fill / verify (heavy imports stay inside functions: --dry-run and unit +# tests must work without torch / tensorrt_llm installed) +# --------------------------------------------------------------------------- # +def _pattern_like(view, seed): + """Deterministic tensor matching `view`, constant along the head axis. + + `view` is an HND block slice: [nblocks, kv_factor, heads, tokens, dim]. + Generated on CPU (bit-identical across nodes), head axis generated as 1 + and expanded, so ctx/gen sides with different local head counts (TP + resharding) or different local layer sets (PP re-splitting) still agree. + """ + import torch + + nb, kv, heads, tok, dim = view.shape + g = torch.Generator(device="cpu").manual_seed(int(seed)) + rnd = torch.rand((nb, kv, 1, tok, dim), dtype=torch.float32, generator=g) + return rnd.to(view.dtype).expand(nb, kv, heads, tok, dim) + + +def _request_block_views(kvm, rid): + """Yield (global_layer, buffer, valid_block_indices) for this rank.""" + for global_layer in kvm.pp_layers: + blocks = kvm.get_batch_cache_indices([rid], layer_idx=global_layer)[0] + valid = [b for b in blocks if b >= 0] + if not valid: + continue + buf = kvm.get_buffers(global_layer, kv_layout="HND") + yield global_layer, buf, valid + + +def fill_request(kvm, rid): + for global_layer, buf, valid in _request_block_views(kvm, rid): + view = buf[valid] + buf[valid] = _pattern_like(view, seed_for(rid, global_layer)).to(view.device) + + +def verify_request(kvm, rid): + """Returns (ok, detail) comparing received blocks to the expected pattern.""" + import torch + + for global_layer, buf, valid in _request_block_views(kvm, rid): + recv = buf[valid] + exp = _pattern_like(recv, seed_for(rid, global_layer)).to(recv.device) + if not torch.equal(recv.float(), exp.float()): + bad = (recv.float() != exp.float()).sum().item() + return False, f"layer={global_layer} mismatched_elements={bad}/{recv.numel()}" + return True, "" + + +def build_kv_cache_manager(kv_shape, plan, side, mapping, max_req_len): + import tensorrt_llm.bindings + import tensorrt_llm.bindings.executor as trtllm_executor + from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager + + DataType = tensorrt_llm.bindings.DataType + CacheTypeCpp = tensorrt_llm.bindings.internal.batch_manager.CacheType + dtype_map = { + "fp8": DataType.FP8, + "fp16": DataType.HALF, + "half": DataType.HALF, + "bf16": DataType.BF16, + } + dtype_str = side["kv_dtype"].lower() + dtype = dtype_map.get(dtype_str) + if dtype is None: + print(f"[precheck] kv dtype {dtype_str!r} not mapped, using BF16", flush=True) + dtype = DataType.BF16 + + spec_config = None + if side["num_nextn_predict_layers"] > 0: + from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig + + spec_config = MTPDecodingConfig(num_nextn_predict_layers=side["num_nextn_predict_layers"]) + + tpb = plan["tokens_per_block"] + padded_len = ((max_req_len + tpb - 1) // tpb) * tpb + owned = pcfg.max_owned_per_chunk(plan, side["role"]) + max_tokens = owned * padded_len + 2 * tpb # concurrent pairs + headroom + + # Real MLA serving uses SELFKONLY (kv_factor=1: one latent plane, no V) — + # see _torch/pyexecutor/_util.py; SELF would double the per-token bytes. + cache_type = CacheTypeCpp.SELFKONLY if kv_shape["is_mla"] else CacheTypeCpp.SELF + return KVCacheManager( + trtllm_executor.KvCacheConfig(max_tokens=max_tokens, enable_block_reuse=False), + cache_type, + num_layers=kv_shape["num_layers"], + num_kv_heads=kv_shape["num_kv_heads"], + head_dim=kv_shape["head_dim"], + tokens_per_block=tpb, + max_seq_len=padded_len, + max_batch_size=max(4, owned + 1), + mapping=mapping, + dtype=dtype, + spec_config=spec_config, + ) + + +def make_request(is_ctx, rid, req_len, runtime, ctx_params=None): + """Build a ctx or gen LlmRequest (mirrors the UCX-tuning harness).""" + import tensorrt_llm.bindings + from tensorrt_llm import DisaggregatedParams + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestType + from tensorrt_llm.sampling_params import SamplingParams + + sampling = SamplingParams() + common = dict( + request_id=rid, + max_new_tokens=1, + input_tokens=list(range(req_len)), + sampling_config=tensorrt_llm.bindings.SamplingConfig(sampling._get_sampling_config()), + is_streaming=False, + ) + if is_ctx: + req = LlmRequest(llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY, **common) + if runtime == "PYTHON": + req.py_disaggregated_params = DisaggregatedParams( + request_type="context_only", disagg_request_id=rid + ) + return req + if runtime == "PYTHON": + req = LlmRequest(llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, **common) + req.py_disaggregated_params = DisaggregatedParams( + request_type="generation_only", + disagg_request_id=rid, + ctx_request_id=rid, + ctx_dp_rank=ctx_params.ctx_dp_rank, + ctx_info_endpoint=ctx_params.disagg_info_endpoint, + first_gen_tokens=ctx_params.first_gen_tokens, + draft_tokens=ctx_params.draft_tokens, + ) + return req + return LlmRequest( + llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, + context_phase_params=ctx_params, + **common, + ) + + +def add_sequence(kvm, req, prompt_len): + kvm.impl.add_sequence_batch([(req.py_request_id, prompt_len, 1)], [req]) + + +def free_sequence(kvm, req): + kvm.impl.remove_sequence(req.py_request_id, req, True) + + +def _wait_gen_complete(xcvr, req, runtime, llm_request_state): + """Block until this gen request's receive finishes (or errors). + + PYTHON transceiver: check_gen_transfer_status(None) blocks for all. C++: + the int API can return before THIS request completes on a cold link, so + poll for a terminal state (bounded by signal.alarm + hang detector). + """ + if runtime == "PYTHON": + xcvr.check_gen_transfer_status(None) + return + terminal = ( + llm_request_state.DISAGG_GENERATION_TRANS_COMPLETE, + llm_request_state.DISAGG_TRANS_ERROR, + ) + while req.state not in terminal: + xcvr.check_gen_transfer_status(1) + if req.state in terminal: + break + time.sleep(0.001) + + +# --------------------------------------------------------------------------- # +# Rendezvous +# --------------------------------------------------------------------------- # +def addr_path(work_dir, ctx_idx, gen_idx): + return os.path.join(work_dir, "rendezvous", f"ctx{ctx_idx}_gen{gen_idx}.addr") + + +def write_addr(path, payload): + os.makedirs(os.path.dirname(path), exist_ok=True) + tmp = f"{path}.tmp.{os.getpid()}" + with open(tmp, "w") as f: + json.dump(payload, f) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, path) + + +def wait_for_addr(path, timeout_s): + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if os.path.exists(path): + try: + with open(path) as f: + return json.load(f) + except (OSError, json.JSONDecodeError): + pass # mid-rename/NFS staleness; retry + time.sleep(1.0) + raise _Timeout(f"rendezvous file {path} not published within {timeout_s}s") + + +# --------------------------------------------------------------------------- # +# Status recording +# --------------------------------------------------------------------------- # +class StatusRecorder: + """Leader-side result sink. + + Rewritten on every record so a SIGKILL'd run still leaves the completed + cases + the in-flight failure on disk. + """ + + def __init__(self, work_dir, role, server_idx, is_leader): + self.role = role + self.server_idx = server_idx + self.is_leader = is_leader + self.cases = [] + status_dir = os.path.join(work_dir, "status") + if is_leader: + os.makedirs(status_dir, exist_ok=True) + self.json_path = os.path.join(status_dir, f"{role}_{server_idx}.json") + self.text_path = os.path.join(status_dir, f"{role}_{server_idx}.status") + env_keys = ("UCX_", "NIXL_", "TRTLLM_", "TLLM_") + self.env = {k: v for k, v in sorted(os.environ.items()) if k.startswith(env_keys)} + + def record(self, peer, req_len, status, reason=""): + if not self.is_leader: + return + self.cases.append({"peer": peer, "req_len": req_len, "status": status, "reason": reason}) + self._flush(final=False) + + def finalize(self, extra=None): + if not self.is_leader: + return + self._flush(final=True, extra=extra) + + def failed_cases(self): + return [c for c in self.cases if c["status"] not in ("PASS", "SKIP")] + + def _flush(self, final, extra=None): + failed = self.failed_cases() + overall = "PASS" if (final and not failed) else ("FAIL" if failed else "RUNNING") + doc = { + "role": self.role, + "server_idx": self.server_idx, + "overall": overall, + "cases": self.cases, + "env": self.env, + } + if extra: + doc.update(extra) + tmp = f"{self.json_path}.tmp" + with open(tmp, "w") as f: + json.dump(doc, f, indent=2) + os.replace(tmp, self.json_path) + with open(f"{self.text_path}.tmp", "w") as f: + if failed: + first = failed[0] + # One-line summary for the launch-script console output: root + # cause only (first line, bounded); the full reason incl. any + # backtrace stays in the .json. + reason = " | ".join(first["reason"].splitlines()[:2])[:400] + f.write( + f"FAIL {self.role}_{self.server_idx}: {len(failed)} case(s) failed; " + f"first: peer={first['peer']} req_len={first['req_len']} " + f"{first['status']}: {reason}\n" + ) + elif final: + f.write(f"PASS {self.role}_{self.server_idx}: {len(self.cases)} case(s)\n") + else: + f.write(f"RUNNING {self.role}_{self.server_idx}\n") + os.replace(f"{self.text_path}.tmp", self.text_path) + + +def parse_bandwidth_gbps(csv_dir, rank, tag="recv"): + """Median per-request bandwidth in GB/s (bytes/1e9), best-effort. + + Parsed from the C++ transceiver CSV this rank wrote via + TRTLLM_KVCACHE_TIME_OUTPUT_PATH. + """ + import csv as csv_mod + import statistics + + path = os.path.join(csv_dir, f"rank_{rank}_{tag}.csv") + try: + with open(path) as f: + rows = list(csv_mod.DictReader(f)) + col = next((c for c in (rows[0] or {}) if "Bandwidth" in c), None) + vals = [float(r[col]) / 8.0 for r in rows if col and r.get(col)] + return statistics.median(vals) if vals else None + except (OSError, ValueError, IndexError, StopIteration): + return None + + +# --------------------------------------------------------------------------- # +# Chunk execution +# --------------------------------------------------------------------------- # +class PrecheckRunner: + """One MPI world = one ctx or gen server instance.""" + + def __init__(self, args, plan, side, comm): + from mpi4py import MPI # noqa: F401 - ensures MPI initialized + + self.args = args + self.plan = plan + self.side = side + self.role = side["role"] + self.is_ctx = self.role == "ctx" + self.server_idx = args.server_idx + self.comm = comm + self.rank = comm.Get_rank() + self.is_leader = self.rank == 0 + self.work_dir = args.work_dir + self.recorder = StatusRecorder(self.work_dir, self.role, self.server_idx, self.is_leader) + self.zmq_ctx = None + self.kvm = None + self.xcvr = None + self.runtime = "CPP" + self.mapping = None + self.llm_request_state = None + self.csv_dir = os.path.join(self.work_dir, "csv", f"{self.role}_{self.server_idx}") + + # ---- consensus helpers ------------------------------------------------- + def _consensus_error(self, local_err): + """All ranks agree whether anyone failed; returns the shared reason.""" + errs = self.comm.allgather("" if local_err is None else repr(local_err)) + bad = [(r, e) for r, e in enumerate(errs) if e] + if not bad: + return None + ranks = [r for r, _ in bad] + return f"rank(s) {ranks}: {bad[0][1]}" + + # ---- setup ------------------------------------------------------------- + def setup(self, kv_shape, max_req_len): + import tensorrt_llm + import tensorrt_llm.bindings + from tensorrt_llm._torch.distributed import Distributed + from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import create_kv_cache_transceiver + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState + from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig + from tensorrt_llm.mapping import Mapping + + self.llm_request_state = LlmRequestState + par = self.side["parallel"] + self.mapping = Mapping( + world_size=par["world_size"], + rank=self.rank, + gpus_per_node=self.plan["gpus_per_node"], + tp_size=par["tp"], + pp_size=par["pp"], + cp_size=par["cp"], + enable_attention_dp=par["enable_attention_dp"], + ) + os.makedirs(self.csv_dir, exist_ok=True) + os.environ["TRTLLM_KVCACHE_TIME_OUTPUT_PATH"] = self.csv_dir + + self.kvm = build_kv_cache_manager(kv_shape, self.plan, self.side, self.mapping, max_req_len) + # Built VERBATIM from the disagg yaml's cache_transceiver_config so + # backend/max_tokens_in_buffer/timeouts match the real test exactly. + cache_cfg = CacheTransceiverConfig(**self.side["cache_transceiver_config"]) + AttentionTypeCpp = tensorrt_llm.bindings.internal.batch_manager.AttentionType + attention_type = AttentionTypeCpp.MLA if kv_shape["is_mla"] else AttentionTypeCpp.DEFAULT + dist_obj = Distributed.get(self.mapping) + self.xcvr = create_kv_cache_transceiver( + self.mapping, dist_obj, self.kvm, attention_type, cache_cfg + ) + if self.xcvr is None: + raise RuntimeError("cache transceiver disabled by config") + # create_kv_cache_transceiver resolves 'auto' in-place (no model + # preference on this path -> C++), so read the effective runtime back. + self.runtime = cache_cfg.transceiver_runtime or "CPP" + + # ---- per-chunk transfer logic ------------------------------------------- + def _pair_rid(self, peer_idx, li, rep, pair): + ctx_idx = self.server_idx if self.is_ctx else peer_idx + gen_idx = peer_idx if self.is_ctx else self.server_idx + return make_rid(ctx_idx, gen_idx, self.plan["num_ctx_servers"], li, rep, pair) + + def _owned(self, chunk): + return pcfg.owned_pairs(self.plan, self.role, self.mapping.tp_rank, chunk) + + def ctx_run_chunk(self, peer_idx, li, req_len, rep, chunk): + """Fill + send owned pairs. + + Returns {pair: context_phase_params} on the leader (params from each + pair's owning dp rank, pp stage 0). + """ + import tensorrt_llm + + owned = self._owned(chunk) + reqs, local_err = {}, None + try: + for pair in owned: + rid = self._pair_rid(peer_idx, li, rep, pair) + req = make_request(True, rid, req_len, self.runtime) + add_sequence(self.kvm, req, req_len) + fill_request(self.kvm, rid) + tensorrt_llm.logger.info( + f"[ctx{self.server_idx} r{self.rank}] rid={rid} len={req_len}: send START" + ) + self.xcvr.respond_and_send_async(req) + reqs[pair] = req + except Exception as e: # noqa: BLE001 - relayed to gen, then raised + local_err = e + reason = self._consensus_error(local_err) + + # Params for pair k come from its owning dp rank at pp stage 0 with + # attention DP; without DP every rank sends the same request, and the + # instance leader's params are the ones the real server would return. + contrib = {} + if local_err is None: + if self.side["parallel"]["enable_attention_dp"]: + if self.mapping.pp_rank == 0: + contrib = {p: r.context_phase_params for p, r in reqs.items()} + elif self.is_leader: + contrib = {p: r.context_phase_params for p, r in reqs.items()} + gathered = self.comm.gather(contrib, root=0) + params_by_pair = {} + if self.is_leader and reason is None: + for d in gathered: + params_by_pair.update(d or {}) + missing = [p for p in chunk if p not in params_by_pair] + if missing: + reason = f"missing context_phase_params for pairs {missing}" + + if reason is not None: + self._free_all(reqs) + raise _TransferError(f"ctx send setup failed: {reason}") + return params_by_pair, reqs + + def ctx_finish_chunk(self, reqs): + """Wait for all in-flight sends of this chunk, then free.""" + local_err = None + try: + self.xcvr.check_context_transfer_status(None) # block-all + bad = [ + p for p, r in reqs.items() if r.state == self.llm_request_state.DISAGG_TRANS_ERROR + ] + if bad: + local_err = _TransferError(f"ctx DISAGG_TRANS_ERROR on pairs {bad}") + except Exception as e: # noqa: BLE001 + local_err = e + finally: + self._free_all(reqs) + reason = self._consensus_error(local_err) + if reason is not None: + raise _TransferError(f"ctx transfer failed: {reason}") + + def gen_run_chunk(self, peer_idx, li, req_len, rep, chunk, params_by_pair): + """Receive + verify owned pairs. Returns (ok, mismatch_detail).""" + import torch + + import tensorrt_llm + + owned = self._owned(chunk) + reqs, local_err = {}, None + try: + for pair in owned: + rid = self._pair_rid(peer_idx, li, rep, pair) + req = make_request( + False, rid, req_len, self.runtime, ctx_params=params_by_pair[pair] + ) + add_sequence(self.kvm, req, req_len) + tensorrt_llm.logger.info( + f"[gen{self.server_idx} r{self.rank}] rid={rid} len={req_len}: recv START" + ) + self.xcvr.request_and_receive_async(req) + reqs[pair] = req + except Exception as e: # noqa: BLE001 + local_err = e + reason = self._consensus_error(local_err) + if reason is not None: + self._free_all(reqs) + raise _TransferError(f"gen receive setup failed: {reason}") + + mismatch = "" + try: + for pair, req in reqs.items(): + _wait_gen_complete(self.xcvr, req, self.runtime, self.llm_request_state) + torch.cuda.synchronize() # receive may land on a side stream + bad = [ + p for p, r in reqs.items() if r.state == self.llm_request_state.DISAGG_TRANS_ERROR + ] + if bad: + local_err = _TransferError(f"gen DISAGG_TRANS_ERROR on pairs {bad}") + elif self.plan["verify_data"]: + for pair, req in reqs.items(): + ok, detail = verify_request(self.kvm, req.py_request_id) + if not ok: + mismatch = f"pair={pair} {detail}" + break + except Exception as e: # noqa: BLE001 + local_err = e + finally: + self._free_all(reqs) + reason = self._consensus_error(local_err) + if reason is not None: + raise _TransferError(f"gen transfer failed: {reason}") + mismatches = [m for m in self.comm.allgather(mismatch) if m] + return (not mismatches, "; ".join(mismatches[:4])) + + def _free_all(self, reqs): + for req in reqs.values(): + try: + free_sequence(self.kvm, req) + except Exception: # noqa: BLE001 - teardown best-effort + pass + + # ---- ZMQ helpers --------------------------------------------------------- + def _zmq(self): + import zmq + + if self.zmq_ctx is None: + self.zmq_ctx = zmq.Context.instance() + return zmq, self.zmq_ctx + + def _leader_send_recv(self, sock, obj): + """REQ round-trip on the gen leader; broadcast the reply to all ranks.""" + reply = None + err = None + if self.is_leader: + try: + sock.send(pickle.dumps(obj)) + reply = pickle.loads(sock.recv()) + except Exception as e: # noqa: BLE001 + err = repr(e) + err, reply = self.comm.bcast((err, reply), root=0) + if err: + raise _TransferError(f"ZMQ control channel failed: {err}") + return reply + + +def _schedule(plan): + """Deterministic (li, req_len, rep, chunk) schedule both sides iterate.""" + out = [] + total_reps = plan["warmup_requests"] + plan["num_requests"] + for li, req_len in enumerate(plan["request_lengths"]): + for rep in range(total_reps): + for chunk in pcfg.chunks(plan): + out.append((li, req_len, rep, chunk)) + return out + + +def hello_timeout_s(plan, num_peers): + """Timeout budget for session handshakes. + + Handshakes are serialized across peers (one gen talks to one ctx at a + time), so waiting for a peer's hello/welcome can legitimately span other + peers' full sessions -- budget rendezvous + per-peer slack. + """ + return plan["rendezvous_timeout_s"] + num_peers * 300 + + +# --------------------------------------------------------------------------- # +# ctx / gen session loops +# --------------------------------------------------------------------------- # +def ctx_serve_peer(runner, sock, peer_idx, arm, disarm): + """Serve one gen peer's full schedule on a dedicated REP socket.""" + plan = runner.plan + comm = runner.comm + + def leader_recv(): + msg, err = None, None + if runner.is_leader: + try: + msg = pickle.loads(sock.recv()) + except Exception as e: # noqa: BLE001 + err = repr(e) + err, msg = comm.bcast((err, msg), root=0) + if err: + raise _TransferError(f"ZMQ recv from gen_{peer_idx} failed: {err}") + return msg + + def leader_reply(obj): + if runner.is_leader: + sock.send(pickle.dumps(obj)) + + arm(f"hello gen_{peer_idx}", seconds=hello_timeout_s(plan, runner.side["num_peers"])) + msg = leader_recv() + if msg[0] != "hello" or msg[1].get("fingerprint") != plan["fingerprint"]: + leader_reply(("abort", "plan fingerprint mismatch (ctx/gen yaml disagree)")) + raise _TransferError(f"handshake with gen_{peer_idx} failed: {msg[:1]}") + leader_reply(("welcome", {"fingerprint": plan["fingerprint"]})) + + for li, req_len, rep, chunk in _schedule(plan): + arm(f"gen_{peer_idx} len={req_len} rep={rep}") + msg = leader_recv() + if msg[0] == "abort": + raise _PeerAbort(f"gen_{peer_idx} aborted: {msg[1]}") + if msg[0] != "go" or (msg[1]["li"], msg[1]["rep"]) != (li, rep): + raise _TransferError( + f"schedule desync with gen_{peer_idx}: expected li={li} rep={rep}, got {msg}" + ) + try: + params_by_pair, reqs = runner.ctx_run_chunk(peer_idx, li, req_len, rep, chunk) + except _TransferError as e: + leader_reply(("abort", str(e))) + raise + leader_reply(("params", params_by_pair)) + runner.ctx_finish_chunk(reqs) + + arm(f"bye gen_{peer_idx}") + msg = leader_recv() + if msg[0] != "done": + raise _TransferError(f"expected done from gen_{peer_idx}, got {msg[:1]}") + leader_reply(("bye", {})) + disarm() + + +def gen_run_peer(runner, peer_idx, arm, disarm): + """Run the full schedule against ctx server `peer_idx`.""" + plan = runner.plan + comm = runner.comm + hello_s = hello_timeout_s(plan, runner.side["num_peers"]) + + # Rendezvous with instance-wide consensus: if only the leader raised here, + # the other ranks would deadlock in the next bcast. + sock, err = None, None + arm(f"rendezvous ctx_{peer_idx}", seconds=hello_s) + if runner.is_leader: + try: + addr = wait_for_addr( + addr_path(runner.work_dir, peer_idx, runner.server_idx), + plan["rendezvous_timeout_s"], + ) + zmq, zctx = runner._zmq() + sock = zctx.socket(zmq.REQ) + sock.setsockopt(zmq.LINGER, 0) + sock.setsockopt(zmq.RCVTIMEO, hello_s * 1000) + sock.connect(f"tcp://{addr['host']}:{addr['port']}") + except Exception as e: # noqa: BLE001 - shared via bcast below + err = repr(e) + err = comm.bcast(err, root=0) + if err: + raise _TransferError(f"rendezvous with ctx_{peer_idx} failed: {err}") + + try: + arm(f"hello ctx_{peer_idx}", seconds=hello_s) + reply = runner._leader_send_recv( + sock, ("hello", {"gen_idx": runner.server_idx, "fingerprint": plan["fingerprint"]}) + ) + if reply[0] == "abort": + raise _TransferError(f"ctx_{peer_idx} aborted handshake: {reply[1]}") + if reply[0] != "welcome": + raise _TransferError(f"unexpected handshake reply from ctx_{peer_idx}: {reply[:1]}") + # Established sessions are dedicated: chunk replies are prompt. + if runner.is_leader: + import zmq as zmq_mod + + sock.setsockopt(zmq_mod.RCVTIMEO, (plan["chunk_timeout_s"] + 30) * 1000) + + case_ok = {} + for li, req_len, rep, chunk in _schedule(plan): + arm(f"ctx_{peer_idx} len={req_len} rep={rep}") + reply = runner._leader_send_recv( + sock, ("go", {"li": li, "rep": rep, "chunk": chunk[0]}) + ) + if reply[0] == "abort": + raise _TransferError(f"ctx_{peer_idx} aborted: {reply[1]}") + ok, detail = runner.gen_run_chunk(peer_idx, li, req_len, rep, chunk, reply[1]) + if rep >= plan["warmup_requests"]: + prev_ok, prev_detail = case_ok.get(req_len, (True, "")) + case_ok[req_len] = (prev_ok and ok, prev_detail or detail) + + arm(f"bye ctx_{peer_idx}") + runner._leader_send_recv(sock, ("done", {})) + disarm() + + for req_len, (ok, detail) in case_ok.items(): + runner.recorder.record( + f"ctx_{peer_idx}", + req_len, + "PASS" if ok else "MISMATCH", + "" if ok else detail, + ) + finally: + if sock is not None: + sock.close(linger=0) + + +# --------------------------------------------------------------------------- # +# main +# --------------------------------------------------------------------------- # +def parse_args(argv=None): + ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + ap.add_argument("--role", required=True, choices=["ctx", "gen"]) + ap.add_argument("--server-idx", type=int, required=True) + ap.add_argument("--config", required=True, help="disagg perf-sanity yaml path") + ap.add_argument("--work-dir", required=True, help="shared dir for rendezvous/status") + ap.add_argument("--benchmark-mode", default="e2e", choices=["e2e", "gen_only"]) + ap.add_argument("--llm-src", default="", help="repo root (model path dict lookup)") + ap.add_argument("--dry-run", action="store_true", help="print the resolved plan and exit") + return ap.parse_args(argv) + + +def load_plan(args): + import yaml + + with open(args.config) as f: + cfg = yaml.safe_load(f) + plan = pcfg.resolve_plan(cfg, benchmark_mode=args.benchmark_mode) + if plan.get("skip"): + return plan, None + model_dir = pcfg.resolve_model_dir(cfg, llm_src=args.llm_src or None) + role_side = pcfg.side_plan(plan, args.role) if not args.dry_run else None + kv_shape = pcfg.model_kv_shape(model_dir) + plan["_kv_shape"] = kv_shape + plan["_model_dir"] = model_dir + return plan, role_side + + +def main(argv=None): + args = parse_args(argv) + plan, side = load_plan(args) + + if args.dry_run: + print(json.dumps({k: v for k, v in plan.items()}, indent=2, default=str)) + return 0 + if plan.get("skip"): + print(f"[precheck] SKIP: {plan['skip_reason']}", flush=True) + return 0 + + # UCX_PROTO_INFO=used is log-only (does not change transport selection): + # it makes UCX >= 1.21 print the chosen GPU<->GPU protocol table, which the + # failure summary uses to spot host-staged tcp fallbacks. + os.environ.setdefault("UCX_PROTO_INFO", "used") + + import torch + from mpi4py import MPI + + import tensorrt_llm + + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + world = comm.Get_size() + expected_world = side["parallel"]["world_size"] + if world != expected_world: + raise RuntimeError( + f"MPI world size {world} != {args.role} world size {expected_world}; " + f"the precheck srun must use the same topology as the real " + f"{args.role} server step." + ) + torch.cuda.set_device(rank % torch.cuda.device_count()) + tensorrt_llm.logger.set_level("info") + + ucx_env = " ".join(f"{k}={v}" for k, v in sorted(os.environ.items()) if k.startswith("UCX_")) + print( + f"[precheck {args.role}_{args.server_idx} r{rank}] UCX env: {ucx_env or ''}", + flush=True, + ) + + runner = PrecheckRunner(args, plan, side, comm) + kv_shape = plan["_kv_shape"] + if runner.is_leader: + print( + f"[precheck {args.role}_{args.server_idx}] kv_shape={kv_shape} " + f"model_dir={plan['_model_dir']} pairs={plan['n_pairs']} " + f"req_lens={plan['request_lengths']}", + flush=True, + ) + + # --- watchdog: signal.alarm for Python-level stalls + HangDetector for + # GIL-released native hangs (dumps stacks, records TIMEOUT, SIGKILLs so + # `srun --kill-on-bad-exit` tears the step down; the external `timeout` + # around the srun is the guaranteed backstop for GIL-held hangs). + signal.signal(signal.SIGALRM, _alarm_handler) + from tensorrt_llm._torch.pyexecutor.hang_detector import HangDetector + + current_cell = {"what": "startup"} + + def _on_hang(): + runner.recorder.record("-", 0, "TIMEOUT", f"hang detected during {current_cell['what']}") + runner.recorder.finalize() + sys.stderr.write( + f"[precheck {args.role}_{args.server_idx} r{rank}] WATCHDOG_KILL " + f"{current_cell['what']}\n" + ) + sys.stderr.flush() + os.kill(os.getpid(), signal.SIGKILL) + + # The detector must outlast the LONGEST legitimate wait (peer handshakes + # are serialized across sessions); per-cell alarms below are the tighter + # bound for actual transfer work. + hang_detector = HangDetector( + timeout=hello_timeout_s(plan, side["num_peers"]) + plan["chunk_timeout_s"] + 60, + on_detected=_on_hang, + ) + hang_detector.start() + + def arm(what, seconds=None): + current_cell["what"] = what + signal.alarm(seconds or plan["chunk_timeout_s"]) + hang_detector.checkpoint() + + def disarm(): + signal.alarm(0) + hang_detector.cancel_task() + + # --- setup: KV pool + transceiver (same config as the real test) --------- + setup_err = None + try: + arm("kv pool + transceiver setup") + runner.setup(kv_shape, max_req_len=max(plan["request_lengths"])) + disarm() + except Exception as e: # noqa: BLE001 - recorded and gated below + disarm() + setup_err = e + reason = runner._consensus_error(setup_err) + if reason is not None: + runner.recorder.record("-", 0, "INIT_ERROR", f"transceiver setup failed: {reason}") + runner.recorder.finalize() + print( + f"[precheck {args.role}_{args.server_idx} r{rank}] INIT_ERROR: {reason}", + file=sys.stderr, + flush=True, + ) + return 1 + + # --- sessions ------------------------------------------------------------- + num_peers = side["num_peers"] + try: + if args.role == "ctx": + # One dedicated REP socket per gen peer (avoids REQ interleaving + # across sessions on a shared socket), published via addr files. + socks = {} + if runner.is_leader: + zmq, zctx = runner._zmq() + host = os.environ.get("SLURMD_NODENAME") or socket.gethostname() + for gj in range(num_peers): + s = zctx.socket(zmq.REP) + s.setsockopt(zmq.LINGER, 0) + # Generous: gen peers are serialized across ctx servers. + s.setsockopt(zmq.RCVTIMEO, hello_timeout_s(plan, num_peers) * 1000) + port = s.bind_to_random_port("tcp://*") + write_addr( + addr_path(runner.work_dir, args.server_idx, gj), + {"host": host, "port": port}, + ) + socks[gj] = s + for gj in range(num_peers): + try: + ctx_serve_peer( + runner, socks.get(gj) if runner.is_leader else None, gj, arm, disarm + ) + runner.recorder.record(f"gen_{gj}", 0, "PASS", "served all transfers") + except _Timeout: + disarm() + runner.recorder.record( + f"gen_{gj}", + 0, + "TIMEOUT", + f"exceeded {plan['chunk_timeout_s']}s during {current_cell['what']}", + ) + except _PeerAbort as e: + disarm() + runner.recorder.record(f"gen_{gj}", 0, "TRANSFER_ERROR", str(e)) + except Exception as e: # noqa: BLE001 - per-peer isolation + disarm() + runner.recorder.record(f"gen_{gj}", 0, "TRANSFER_ERROR", repr(e)) + else: + for ci in range(num_peers): + try: + gen_run_peer(runner, ci, arm, disarm) + except _Timeout: + disarm() + runner.recorder.record( + f"ctx_{ci}", + 0, + "TIMEOUT", + f"exceeded {plan['chunk_timeout_s']}s during {current_cell['what']}", + ) + except Exception as e: # noqa: BLE001 - per-peer isolation + disarm() + runner.recorder.record(f"ctx_{ci}", 0, "TRANSFER_ERROR", repr(e)) + finally: + disarm() + hang_detector.stop() + + # --- teardown + result ------------------------------------------------------ + bw = None + if args.role == "gen": + local_bw = parse_bandwidth_gbps(runner.csv_dir, rank) + bws = [b for b in comm.gather(local_bw, root=0) or [] if b] + if runner.is_leader and bws: + bw = sorted(bws)[len(bws) // 2] + if runner.xcvr is not None and hasattr(runner.xcvr, "shutdown"): + try: + runner.xcvr.shutdown() + except Exception: # noqa: BLE001 - teardown best-effort + pass + + failed_local = 1 if runner.recorder.failed_cases() else 0 + failed = comm.allreduce(failed_local, op=MPI.MAX) + runner.recorder.finalize(extra={"per_gpu_bw_gbps": bw} if bw else None) + comm.Barrier() + if runner.is_leader: + verdict = "FAIL" if failed else "PASS" + bw_note = f" per-GPU BW ~{bw:.1f} GB/s" if bw else "" + print(f"[precheck {args.role}_{args.server_idx}] {verdict}{bw_note}", flush=True) + return 1 if failed else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py new file mode 100644 index 000000000000..492e895b1369 --- /dev/null +++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py @@ -0,0 +1,216 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""CPU-only tests for the disagg cache-transceiver precheck config resolution. + +Target: tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py +""" + +import json +import os +import sys + +import pytest + +_PRECHECK_DIR = os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "..", + "..", + "..", + "tests", + "scripts", + "perf-sanity", + "cache_transceiver_precheck", +) +sys.path.insert(0, os.path.abspath(_PRECHECK_DIR)) + +import precheck_config as pcfg # noqa: E402 + + +def _disagg_yaml(ctx_extra=None, gen_extra=None, **overrides): + """Minimal disagg perf-sanity yaml shaped like the checked-in configs.""" + ctx = { + "tensor_parallel_size": 4, + "pipeline_parallel_size": 1, + "context_parallel_size": 1, + "enable_attention_dp": True, + "kv_cache_config": {"dtype": "fp8"}, + "cache_transceiver_config": {"max_tokens_in_buffer": 16384, "backend": "NIXL"}, + "speculative_config": {"decoding_type": "MTP", "num_nextn_predict_layers": 1}, + } + gen = { + "tensor_parallel_size": 16, + "pipeline_parallel_size": 1, + "context_parallel_size": 1, + "enable_attention_dp": True, + "kv_cache_config": {"dtype": "fp8"}, + "cache_transceiver_config": {"max_tokens_in_buffer": 16384, "backend": "NIXL"}, + "speculative_config": {"decoding_type": "MTP", "num_nextn_predict_layers": 1}, + } + ctx.update(ctx_extra or {}) + gen.update(gen_extra or {}) + cfg = { + "metadata": {"model_name": "deepseek_r1_0528_fp4_v2"}, + "benchmark": {"mode": "e2e", "input_length": 8192, "output_length": 1024}, + "hardware": {"gpus_per_node": 4, "num_ctx_servers": 1, "num_gen_servers": 1}, + "worker_config": {"ctx": ctx, "gen": gen}, + } + cfg.update(overrides) + return cfg + + +def test_resolve_plan_adp_asymmetric(): + plan = pcfg.resolve_plan(_disagg_yaml()) + assert not plan["skip"] + assert plan["ctx"] == { + "tp": 4, + "pp": 1, + "cp": 1, + "enable_attention_dp": True, + "world_size": 4, + "dp_size": 4, + } + assert plan["gen"]["world_size"] == 16 and plan["gen"]["dp_size"] == 16 + # Cover every gen dp rank. + assert plan["n_pairs"] == 16 + assert plan["chunk_size"] == 8 + assert plan["request_lengths"] == [1024, 8192] + assert plan["ctx_num_nextn_predict_layers"] == 1 + assert plan["ctx_cache_transceiver_config"]["backend"] == "NIXL" + + +def test_request_lengths_clamped_by_buffer_and_cap(): + cfg = _disagg_yaml(benchmark={"mode": "e2e", "input_length": 131072}) + cfg["worker_config"]["ctx"]["cache_transceiver_config"]["max_tokens_in_buffer"] = 131104 + cfg["worker_config"]["gen"]["cache_transceiver_config"]["max_tokens_in_buffer"] = 131104 + plan = pcfg.resolve_plan(cfg) + # Derived ISL is capped by max_request_length (default 32768). + assert plan["request_lengths"] == [1024, 32768] + + cfg["worker_config"]["ctx"]["cache_transceiver_config"]["max_tokens_in_buffer"] = 4096 + cfg["worker_config"]["gen"]["cache_transceiver_config"]["max_tokens_in_buffer"] = 4096 + plan = pcfg.resolve_plan(cfg) + assert plan["request_lengths"] == [1024, 4096] + + # Explicit yaml override is used as-is (not capped). + cfg["cache_transceiver_precheck"] = {"request_lengths": [64000]} + cfg["worker_config"]["ctx"]["cache_transceiver_config"]["max_tokens_in_buffer"] = 131104 + cfg["worker_config"]["gen"]["cache_transceiver_config"]["max_tokens_in_buffer"] = 131104 + plan = pcfg.resolve_plan(cfg) + assert plan["request_lengths"] == [64000] + + +def test_gen_only_no_context_skips(): + cfg = _disagg_yaml(benchmark={"mode": "gen_only_no_context", "input_length": 1024}) + plan = pcfg.resolve_plan(cfg, benchmark_mode="gen_only") + assert plan["skip"] + # e2e over the same yaml still runs (ctx servers are launched there). + assert not pcfg.resolve_plan(cfg, benchmark_mode="e2e")["skip"] + + +def test_backend_mismatch_raises(): + cfg = _disagg_yaml( + gen_extra={"cache_transceiver_config": {"backend": "UCX", "max_tokens_in_buffer": 16384}} + ) + with pytest.raises(ValueError, match="backend mismatch"): + pcfg.resolve_plan(cfg) + + +def test_pair_participation_and_chunks(): + plan = pcfg.resolve_plan(_disagg_yaml()) + # ADP ctx (dp4): pair k belongs to tp_rank k % 4. + assert pcfg.pair_participates(plan, "ctx", 1, 5) + assert not pcfg.pair_participates(plan, "ctx", 0, 5) + # ADP gen (dp16): 1:1. + assert pcfg.pair_participates(plan, "gen", 5, 5) + assert not pcfg.pair_participates(plan, "gen", 4, 5) + assert pcfg.chunks(plan) == [list(range(8)), list(range(8, 16))] + # ctx rank owns 2 pairs per chunk of 8; gen rank owns at most 1. + assert pcfg.max_owned_per_chunk(plan, "ctx") == 2 + assert pcfg.max_owned_per_chunk(plan, "gen") == 1 + + # Non-ADP side participates everywhere and owns the whole chunk. + plan_pp = pcfg.resolve_plan( + _disagg_yaml( + ctx_extra={ + "tensor_parallel_size": 1, + "pipeline_parallel_size": 8, + "enable_attention_dp": False, + } + ) + ) + assert plan_pp["ctx"]["dp_size"] == 1 and plan_pp["n_pairs"] == 16 + assert pcfg.pair_participates(plan_pp, "ctx", 0, 11) + assert pcfg.max_owned_per_chunk(plan_pp, "ctx") == plan_pp["chunk_size"] + + +def test_fingerprint_role_agnostic(): + plan_a = pcfg.resolve_plan(_disagg_yaml()) + plan_b = pcfg.resolve_plan(_disagg_yaml()) + assert plan_a["fingerprint"] == plan_b["fingerprint"] + changed = _disagg_yaml() + changed["worker_config"]["gen"]["tensor_parallel_size"] = 8 + assert pcfg.resolve_plan(changed)["fingerprint"] != plan_a["fingerprint"] + + +def test_model_kv_shape_mla_and_gqa(tmp_path): + mla = tmp_path / "mla" + mla.mkdir() + (mla / "config.json").write_text( + json.dumps( + { + "num_hidden_layers": 61, + "kv_lora_rank": 512, + "qk_rope_head_dim": 64, + "num_attention_heads": 128, + } + ) + ) + shape = pcfg.model_kv_shape(str(mla)) + assert shape == { + "num_layers": 61, + "num_kv_heads": 1, + "head_dim": 576, + "is_mla": True, + "source": "config.json (MLA)", + } + + gqa = tmp_path / "gqa" + gqa.mkdir() + (gqa / "config.json").write_text( + json.dumps( + { + "num_hidden_layers": 32, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "hidden_size": 4096, + } + ) + ) + shape = pcfg.model_kv_shape(str(gqa)) + assert shape["num_kv_heads"] == 8 and shape["head_dim"] == 128 and not shape["is_mla"] + + # Unresolvable model dir -> synthetic fallback (precheck still runs). + assert pcfg.model_kv_shape(None)["source"] == "fallback" + assert pcfg.model_kv_shape(str(tmp_path / "missing"))["source"] == "fallback" + + +def test_side_plan_views(): + plan = pcfg.resolve_plan(_disagg_yaml()) + ctx_view = pcfg.side_plan(plan, "ctx") + gen_view = pcfg.side_plan(plan, "gen") + assert ctx_view["parallel"]["world_size"] == 4 + assert ctx_view["peer_parallel"]["world_size"] == 16 + assert ctx_view["num_peers"] == 1 and gen_view["num_peers"] == 1 + assert gen_view["cache_transceiver_config"]["max_tokens_in_buffer"] == 16384 From 05323f95dcfb7293f4abf9a61ef0a4092e523b81 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Mon, 20 Jul 2026 03:40:16 +0000 Subject: [PATCH 02/21] [None][infra] Gate disagg perf sanity stages on a cache transceiver precheck Before starting the real ctx/gen servers, the generated launch script now optionally runs one precheck srun per server instance with the same node slices (deterministic SLURM_JOB_NODELIST slicing + srun -w, also applied to the real server steps), the same MPI topology, and the same UCX env: submit.py builds the precheck commands from the exact ucx_tls_cmd + CTX/GEN_WORKER_ENV_VARS strings used by the worker steps, and the shared slurm_env_setup.sh (extracted from slurm_run.sh) applies the same LD_LIBRARY_PATH / 'unset UCX_TLS=tcp' / PMIX_MCA_gds fixups. The gate logic lives in slurm_ct_precheck_gate.sh as functions; submit.py splices it into the launch script ahead of slurm_launch_draft.sh, which just calls run_cache_transceiver_precheck. On precheck failure the stage aborts before any model bring-up, with per-instance verdicts, the first error lines + tail of each failing step log, and UCX red-flag lines (sw-emul tcp fallback), and writes a synthetic junit xml to $jobWorkspace/results-ct-precheck.xml so the failure shows up as a test entry in the Jenkins report (uploadResults already collects results*.xml). On by default; disable per yaml via cache_transceiver_precheck.enabled: false, or globally with TRTLLM_DISAGG_CT_PRECHECK=0 at generation time (=1 force-enables; the env var overrides the yaml either way). The local perf flow (jenkins/scripts/perf/local) gets the same exports. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> [None][infra] Make the precheck run script work with the local perf flow layout slurm_precheck_run.sh hardcoded the CI layout (cd $resourcePathNode; llmSrcNode=$resourcePathNode/TensorRT-LLM/src), but the local perf flow (jenkins/scripts/perf/local/submit.py) exports llmSrcNode directly and has no resourcePathNode -- with set -u the precheck step died instantly. Honor whichever layout is present, and fail with a clear message when neither is exported. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../disaggregated/slurm_ct_precheck_gate.sh | 161 ++++++++++++++++++ .../perf/disaggregated/slurm_launch_draft.sh | 38 ++++- .../perf/disaggregated/slurm_precheck_run.sh | 39 +++++ jenkins/scripts/perf/local/submit.py | 41 ++++- jenkins/scripts/perf/submit.py | 48 +++++- jenkins/scripts/slurm_env_setup.sh | 41 +++++ jenkins/scripts/slurm_run.sh | 31 +--- .../cache_transceiver_precheck/README.md | 12 +- 8 files changed, 376 insertions(+), 35 deletions(-) create mode 100644 jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh create mode 100644 jenkins/scripts/perf/disaggregated/slurm_precheck_run.sh create mode 100644 jenkins/scripts/slurm_env_setup.sh diff --git a/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh b/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh new file mode 100644 index 000000000000..65e9158b7973 --- /dev/null +++ b/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh @@ -0,0 +1,161 @@ +# Cache-transceiver precheck gate for the disaggregated perf-sanity launch +# script. submit.py splices this file into the generated launch script ahead +# of slurm_launch_draft.sh, which calls run_cache_transceiver_precheck after +# computing the per-server node slices. Kept as functions in a separate file +# so the gate logic can be sourced and exercised standalone. +# +# Expects the launch-script globals: srunArgs, numGenServers, numCtxServers, +# nodesPerGenServer/nodesPerCtxServer, gpusPerNodePerGenServer/ +# gpusPerNodePerCtxServer, genNodeLists/ctxNodeLists, testOutputDir, +# jobWorkspace, pytestCommandCTXPrecheck/pytestCommandGENPrecheck, +# precheckRunScript, ctPrecheckEnabled, ctPrecheckTimeout, stageName, +# and the cleanup_on_failure function. + +# Escape text for embedding in XML; junit parsers also choke on raw control +# bytes that MPI/UCX logs may contain, so strip those too. +ct_xml_escape() { + sed -e 's/&/\&/g' -e 's//\>/g' -e 's/"/\"/g' \ + | tr -d '\000-\010\013\014\016-\037' +} + +# First few root-cause-shaped lines of a step log — the tail alone can miss +# the real error when it happened early and retry spam follows. +ct_first_errors() { + grep -m 5 -nE "Traceback \(most recent call last\)|MPI_ABORT|MPIR_Err|srun: error|Segmentation fault|CUDA error|RuntimeError|AssertionError|INIT_ERROR|TRANSFER_ERROR" \ + "$1" 2>/dev/null || true +} + +# Console summary for a failed precheck: per-instance verdicts, first error +# lines + tail of each failing step log, and UCX red-flag lines. +# Uses: precheckDir, precheckNames. +ct_print_failure_summary() { + echo "====================================================================" + echo "CACHE TRANSCEIVER PRECHECK FAILED - the disaggregated test will NOT run" + echo "Instance verdicts:" + cat "$precheckDir"/status/*.status 2>/dev/null \ + || echo "(no status files - steps died before transceiver setup; see logs below)" + echo "" + echo "Failing step logs:" + for k in "${!precheckNames[@]}"; do + statusFile="$precheckDir/status/${precheckNames[$k]}.status" + if [ -f "$statusFile" ] && grep -q "^PASS" "$statusFile"; then + continue + fi + stepLog="$precheckDir/logs/${precheckNames[$k]}.log" + echo "----- ${precheckNames[$k]} ($stepLog) -----" + echo "First error lines (line-numbered):" + ct_first_errors "$stepLog" + echo "Tail:" + tail -n 60 "$stepLog" 2>/dev/null || true + done + echo "" + echo "UCX red flags (host-staged tcp fallback / UCX errors), if any:" + grep -hE "sw-emul|UCX +(ERROR|WARN)" "$precheckDir"/logs/*.log 2>/dev/null | sort -u | head -20 || true + echo "Full artifacts: $precheckDir (status/*.json, logs/, csv/)" + echo "====================================================================" +} + +# Synthetic junit result so the failure shows up as a test entry in the +# Jenkins test report: uploadResults scps $jobWorkspace/results*.xml back and +# junit() ingests them. Best-effort — callers must not let this mask the real +# failure path. Uses: precheckDir, precheckNames, jobWorkspace, stageName. +ct_write_junit_xml() { + local junitXml="$jobWorkspace/results-ct-precheck.xml" + local suiteName + suiteName="$(printf '%s' "${stageName:-${SLURM_JOB_NAME:-disagg_perf_sanity}}" | ct_xml_escape)" + local junitFailures=0 + local junitCases="" + local name statusFile stepLog verdict detail + for k in "${!precheckNames[@]}"; do + name="${precheckNames[$k]}" + statusFile="$precheckDir/status/$name.status" + if [ -f "$statusFile" ] && grep -q "^PASS" "$statusFile"; then + junitCases+=""$'\n' + continue + fi + junitFailures=$((junitFailures + 1)) + verdict="$( (head -n 1 "$statusFile" 2>/dev/null \ + || echo "NO_STATUS: step died before writing a verdict (see log)") | ct_xml_escape)" + stepLog="$precheckDir/logs/$name.log" + detail="$( { echo "First error lines:"; ct_first_errors "$stepLog"; echo ""; \ + echo "Log tail ($stepLog):"; tail -n 60 "$stepLog" 2>/dev/null; } | ct_xml_escape)" + junitCases+="" + junitCases+="$detail"$'\n' + done + if { + echo '' + echo "" + printf '%s' "$junitCases" + echo "" + } > "$junitXml" 2>/dev/null; then + echo "Synthetic junit result written to $junitXml (will appear in the Jenkins test report)" + else + echo "WARNING: could not write synthetic junit xml to $junitXml (non-fatal)" + fi +} + +# Run one precheck srun per ctx/gen server instance with the same node +# slices, MPI topology, and UCX env as the real server steps +# (pytestCommand{CTX,GEN}Precheck embed the same ucx_tls_cmd + worker env var +# strings). On failure: console summary + synthetic junit, then the stage +# aborts via cleanup_on_failure — before any model bring-up. +run_cache_transceiver_precheck() { + if [ "${ctPrecheckEnabled:-0}" != "1" ] || [ "${TRTLLM_DISAGG_BENCHMARK_GEN_ONLY:-0}" = "1" ]; then + return 0 + fi + echo "Starting cache transceiver precheck..." + precheckDir="$testOutputDir/cache_transceiver_precheck" + mkdir -p "$precheckDir/logs" + precheckPids=() + precheckNames=() + local i gen_world_size ctx_world_size + for i in $(seq 0 $((numGenServers - 1))); do + gen_world_size=$((nodesPerGenServer * gpusPerNodePerGenServer)) + export DISAGG_SERVING_TYPE="GEN_PRECHECK_$i" + export pytestCommand="$pytestCommandGENPrecheck --server-idx $i" + timeout -k 60 "${ctPrecheckTimeout:-900}" \ + srun "${srunArgs[@]}" --mpi=pmix --kill-on-bad-exit=1 \ + -N $nodesPerGenServer \ + -w "${genNodeLists[$i]}" \ + --ntasks=$gen_world_size \ + --ntasks-per-node=$gpusPerNodePerGenServer \ + bash $precheckRunScript &> "$precheckDir/logs/gen_$i.log" & + precheckPids+=($!) + precheckNames+=("gen_$i") + sleep 5 # Wait for pyxis container namespace initialization to avoid race condition + done + for i in $(seq 0 $((numCtxServers - 1))); do + ctx_world_size=$((nodesPerCtxServer * gpusPerNodePerCtxServer)) + export DISAGG_SERVING_TYPE="CTX_PRECHECK_$i" + export pytestCommand="$pytestCommandCTXPrecheck --server-idx $i" + timeout -k 60 "${ctPrecheckTimeout:-900}" \ + srun "${srunArgs[@]}" --mpi=pmix --kill-on-bad-exit=1 \ + -N $nodesPerCtxServer \ + -w "${ctxNodeLists[$i]}" \ + --ntasks=$ctx_world_size \ + --ntasks-per-node=$gpusPerNodePerCtxServer \ + bash $precheckRunScript &> "$precheckDir/logs/ctx_$i.log" & + precheckPids+=($!) + precheckNames+=("ctx_$i") + sleep 5 # Wait for pyxis container namespace initialization to avoid race condition + done + + local precheckFailed=0 k rc + for k in "${!precheckPids[@]}"; do + if wait "${precheckPids[$k]}"; then + echo "Precheck step ${precheckNames[$k]} passed" + else + rc=$? + echo "Precheck step ${precheckNames[$k]} FAILED (exit $rc; 124 = external timeout)" + precheckFailed=1 + fi + done + + if [ "$precheckFailed" -eq 1 ]; then + ct_print_failure_summary + ct_write_junit_xml + cleanup_on_failure "Cache transceiver precheck failed. See summary above and $precheckDir" + fi + echo "Cache transceiver precheck PASSED:" + cat "$precheckDir"/status/*.status 2>/dev/null || true +} diff --git a/jenkins/scripts/perf/disaggregated/slurm_launch_draft.sh b/jenkins/scripts/perf/disaggregated/slurm_launch_draft.sh index a3712ff6d114..9b8dd9d75504 100644 --- a/jenkins/scripts/perf/disaggregated/slurm_launch_draft.sh +++ b/jenkins/scripts/perf/disaggregated/slurm_launch_draft.sh @@ -17,6 +17,38 @@ if ! srun "${srunArgs[@]}" $installScript &> $jobWorkspace/install.log; then fi echo "Installation completed on all nodes" +# Deterministic node slices per server: gen servers take the first nodes, +# then ctx servers (same order the steps are started in). Both the cache +# transceiver precheck and the real server steps pin to these slices with +# `srun -w`, so the precheck exercises exactly the node pairs / NICs the +# real disaggregated test will use. +mapfile -t allNodes < <(scontrol show hostnames "$SLURM_JOB_NODELIST") +nodeCursor=0 +genNodeLists=() +for i in $(seq 0 $((numGenServers - 1))); do + slice=("${allNodes[@]:$nodeCursor:$nodesPerGenServer}") + genNodeLists+=("$(IFS=,; echo "${slice[*]}")") + nodeCursor=$((nodeCursor + nodesPerGenServer)) +done +ctxNodeLists=() +if [ "${TRTLLM_DISAGG_BENCHMARK_GEN_ONLY:-0}" != "1" ]; then + for i in $(seq 0 $((numCtxServers - 1))); do + slice=("${allNodes[@]:$nodeCursor:$nodesPerCtxServer}") + ctxNodeLists+=("$(IFS=,; echo "${slice[*]}")") + nodeCursor=$((nodeCursor + nodesPerCtxServer)) + done +fi +if [ "$nodeCursor" -gt "${#allNodes[@]}" ]; then + cleanup_on_failure "Node slicing needs $nodeCursor nodes but the job only has ${#allNodes[@]} ($SLURM_JOB_NODELIST)" +fi + +# Cache transceiver network precheck: same instance count / node slices / +# MPI topology / UCX env as the real ctx+gen server steps. On failure the +# stage aborts HERE, with per-instance verdicts + a synthetic junit entry, +# before any model bring-up. Functions come from slurm_ct_precheck_gate.sh, +# spliced in above this draft by submit.py. No-op unless ctPrecheckEnabled=1. +run_cache_transceiver_precheck + # Start gen servers echo "Starting gen servers..." for i in $(seq 0 $((numGenServers - 1))); do @@ -25,10 +57,11 @@ for i in $(seq 0 $((numGenServers - 1))); do export pytestCommand="$pytestCommandGENWorker" srun "${srunArgs[@]}" --mpi=pmix --kill-on-bad-exit=1 \ -N $nodesPerGenServer \ + -w "${genNodeLists[$i]}" \ --ntasks=$gen_world_size \ --ntasks-per-node=$gpusPerNodePerGenServer \ $runScript &> $testOutputDir/gen_server_$i.log & - echo "Started gen server $i" + echo "Started gen server $i on ${genNodeLists[$i]}" sleep 5 # Wait for pyxis container namespace initialization to avoid race condition done @@ -41,10 +74,11 @@ if [ "${TRTLLM_DISAGG_BENCHMARK_GEN_ONLY:-0}" != "1" ]; then export pytestCommand="$pytestCommandCTXWorker" srun "${srunArgs[@]}" --mpi=pmix --kill-on-bad-exit=1 \ -N $nodesPerCtxServer \ + -w "${ctxNodeLists[$i]}" \ --ntasks=$ctx_world_size \ --ntasks-per-node=$gpusPerNodePerCtxServer \ $runScript &> $testOutputDir/ctx_server_$i.log & - echo "Started ctx server $i" + echo "Started ctx server $i on ${ctxNodeLists[$i]}" sleep 5 # Wait for pyxis container namespace initialization to avoid race condition done else diff --git a/jenkins/scripts/perf/disaggregated/slurm_precheck_run.sh b/jenkins/scripts/perf/disaggregated/slurm_precheck_run.sh new file mode 100644 index 000000000000..1dbafd5c51bf --- /dev/null +++ b/jenkins/scripts/perf/disaggregated/slurm_precheck_run.sh @@ -0,0 +1,39 @@ +#!/bin/bash + +# Per-rank container entrypoint for the disagg cache-transceiver PRECHECK +# step (launched by slurm_launch_draft.sh BEFORE the real ctx/gen servers). +# +# It intentionally mirrors slurm_run.sh's runtime environment: the same +# slurm_env_setup.sh is sourced (LD_LIBRARY_PATH, the `unset UCX_TLS=tcp` +# fixup, PMIX_MCA_gds), and $pytestCommand carries the same +# `unset/export UCX_TLS ...` prefix and worker env vars as the real ctx/gen +# worker steps (built from the same strings by jenkins/scripts/perf/submit.py), +# so a precheck PASS/FAIL is representative of the network environment the +# real test will run in. Unlike slurm_run.sh it skips coverage/perf-report +# handling -- the precheck is a plain MPI program, not a pytest run. + +set -xEeuo pipefail +trap 'rc=$?; echo "Error in file ${BASH_SOURCE[0]} on line $LINENO: $BASH_COMMAND (exit $rc)"; exit $rc' ERR + +# CI (jenkins/L0_Test.groovy) exports resourcePathNode with the /tmp +# extraction layout; the local flow (jenkins/scripts/perf/local/submit.py) +# exports llmSrcNode pointing at the repo directly and has no +# resourcePathNode. Honor whichever layout is present. +if [ -n "${resourcePathNode:-}" ]; then + cd "$resourcePathNode" + llmSrcNode=$resourcePathNode/TensorRT-LLM/src +fi +: "${llmSrcNode:?either resourcePathNode or llmSrcNode must be exported}" + +source "$llmSrcNode/jenkins/scripts/slurm_env_setup.sh" +slurm_setup_runtime_env + +echo "Precheck rank ${SLURM_PROCID:-?} (${DISAGG_SERVING_TYPE:-unknown}) command: $pytestCommand" + +set +e +eval $pytestCommand +precheck_exit_code=$? +set -e + +echo "Rank${SLURM_PROCID:-?} cache-transceiver precheck finished with exit code $precheck_exit_code" +exit $precheck_exit_code diff --git a/jenkins/scripts/perf/local/submit.py b/jenkins/scripts/perf/local/submit.py index 0d56b3cfa068..9ca0c8894100 100755 --- a/jenkins/scripts/perf/local/submit.py +++ b/jenkins/scripts/perf/local/submit.py @@ -910,6 +910,37 @@ def main(): ] ) + # Cache-transceiver network precheck (mirrors jenkins/scripts/perf/ + # submit.py): reuses the exact ucx_tls_cmd + worker env strings of the + # real worker steps. On by default; disabled per-yaml + # (cache_transceiver_precheck.enabled: false); TRTLLM_DISAGG_CT_PRECHECK + # =0/1 at generation time overrides the yaml either way. + precheck_cfg = config.get("cache_transceiver_precheck", {}) or {} + precheck_env = os.environ.get("TRTLLM_DISAGG_CT_PRECHECK") + if precheck_env is not None: + precheck_enabled = precheck_env == "1" + else: + precheck_enabled = bool(precheck_cfg.get("enabled", True)) + precheck_mode = "gen_only" if benchmark_mode == "gen_only" else "e2e" + precheck_cmd = ( + "python3 $llmSrcNode/tests/scripts/perf-sanity/cache_transceiver_precheck/" + "run_precheck.py --config $configYamlPath " + "--work-dir $testOutputDir/cache_transceiver_precheck " + f"--benchmark-mode {precheck_mode} --llm-src $llmSrcNode" + ) + script_prefix_lines.extend( + [ + f"export ctPrecheckEnabled={1 if precheck_enabled else 0}", + f"export ctPrecheckTimeout={int(precheck_cfg.get('step_timeout_s', 900))}", + "export precheckRunScript=$llmSrcNode/jenkins/scripts/perf/" + "disaggregated/slurm_precheck_run.sh", + f'export pytestCommandCTXPrecheck="{ucx_tls_cmd} $CTX_WORKER_ENV_VARS' + f' {precheck_cmd} --role ctx"', + f'export pytestCommandGENPrecheck="{ucx_tls_cmd} $GEN_WORKER_ENV_VARS' + f' {precheck_cmd} --role gen"', + ] + ) + # Add srun args for disagg srun_args_lines.extend( [ @@ -1006,9 +1037,17 @@ def main(): draft_launch_lines = remove_whitespace_lines(draft_launch_lines) draft_launch_content = "\n".join(draft_launch_lines) + # The disagg draft calls run_cache_transceiver_precheck; splice in the + # gate function library (sibling of the draft) ahead of it. + gate_content = "" + if runtime_mode == "disaggregated": + gate_sh = os.path.join(os.path.dirname(draft_launch_sh), "slurm_ct_precheck_gate.sh") + with open(gate_sh, "r") as f: + gate_content = "\n".join(remove_whitespace_lines(f.read().split("\n"))) + "\n\n" + # Combine and write launch script script_prefix = "\n".join(script_prefix_lines) - final_script = f"{script_prefix}\n\n{srun_args}\n\n{draft_launch_content}" + final_script = f"{script_prefix}\n\n{srun_args}\n\n{gate_content}{draft_launch_content}" with open(launch_sh, "w") as f: f.write(final_script) diff --git a/jenkins/scripts/perf/submit.py b/jenkins/scripts/perf/submit.py index da49367c1fdc..5ddb40ca197c 100755 --- a/jenkins/scripts/perf/submit.py +++ b/jenkins/scripts/perf/submit.py @@ -626,6 +626,42 @@ def main(): f"export testOutputDir={test_output_dir}", ] ) + + # Cache-transceiver network precheck: runs BEFORE the real ctx/gen + # servers with the same instance topology, and reuses the exact + # $ucx_tls_cmd / $CTX_WORKER_ENV_VARS / $GEN_WORKER_ENV_VARS strings + # of the worker steps so the UCX environment matches by construction. + # On by default. Disabled per-yaml (cache_transceiver_precheck. + # enabled: false); TRTLLM_DISAGG_CT_PRECHECK=0/1 at generation time + # overrides the yaml either way (global kill switch). + precheck_cfg = config.get("cache_transceiver_precheck", {}) or {} + precheck_env = os.environ.get("TRTLLM_DISAGG_CT_PRECHECK") + if precheck_env is not None: + precheck_enabled = precheck_env == "1" + else: + precheck_enabled = bool(precheck_cfg.get("enabled", True)) + config_yaml_node = f"$llmSrcNode/{os.path.relpath(config_yaml, args.llm_src)}" + precheck_cmd = ( + "python3 $llmSrcNode/tests/scripts/perf-sanity/cache_transceiver_precheck/" + f"run_precheck.py --config {config_yaml_node} " + "--work-dir $testOutputDir/cache_transceiver_precheck " + f"--benchmark-mode {benchmark_mode} --llm-src $llmSrcNode" + ) + script_prefix_lines.extend( + [ + f"export ctPrecheckEnabled={1 if precheck_enabled else 0}", + f"export ctPrecheckTimeout={int(precheck_cfg.get('step_timeout_s', 900))}", + # Suite name for the synthetic junit xml the launch script + # writes when the precheck fails. + f'export stageName="{args.stage_name}"', + "export precheckRunScript=$llmSrcNode/jenkins/scripts/perf/" + "disaggregated/slurm_precheck_run.sh", + f'export pytestCommandCTXPrecheck="{ucx_tls_cmd} $CTX_WORKER_ENV_VARS' + f' {precheck_cmd} --role ctx"', + f'export pytestCommandGENPrecheck="{ucx_tls_cmd} $GEN_WORKER_ENV_VARS' + f' {precheck_cmd} --role gen"', + ] + ) srun_args_lines.extend( [ "--container-env=DISAGG_SERVING_TYPE", @@ -645,8 +681,18 @@ def main(): draft_launch_lines = remove_whitespace_lines(draft_launch_content.split("\n")) draft_launch_content = "\n".join(draft_launch_lines) + # The disagg draft calls run_cache_transceiver_precheck; splice in the + # gate function library (sibling of the draft) ahead of it. + gate_content = "" + if runtime_mode == "disaggregated": + gate_sh = os.path.join( + os.path.dirname(args.draft_launch_sh), "slurm_ct_precheck_gate.sh" + ) + with open(gate_sh, "r") as f: + gate_content = "\n".join(remove_whitespace_lines(f.read().split("\n"))) + "\n" + with open(args.launch_sh, "w") as f: - f.write(f"{script_prefix}\n{srun_args}\n{draft_launch_content}") + f.write(f"{script_prefix}\n{srun_args}\n{gate_content}{draft_launch_content}") print(f"Launch script generated at: {args.launch_sh}") print(f"Launch script:\n{script_prefix}\n{srun_args}\n{draft_launch_content}") diff --git a/jenkins/scripts/slurm_env_setup.sh b/jenkins/scripts/slurm_env_setup.sh new file mode 100644 index 000000000000..6b51c7c1b8fe --- /dev/null +++ b/jenkins/scripts/slurm_env_setup.sh @@ -0,0 +1,41 @@ +#!/bin/bash + +# Shared runtime-environment setup for SLURM test steps. Sourced by +# slurm_run.sh (real test steps) AND by +# jenkins/scripts/perf/disaggregated/slurm_precheck_run.sh (the cache +# transceiver precheck), so both run with IDENTICAL library paths and +# UCX/PMIx fixups -- the precheck must observe exactly the network +# environment the real disaggregated workers will use. Keep any change here +# valid for both callers. + +slurm_setup_runtime_env() { + # Prepend the installed tensorrt_llm wheel's libs to LD_LIBRARY_PATH. + local containerPipLLMLibPath + containerPipLLMLibPath=$(pip3 show tensorrt_llm | grep "Location" | awk -F ":" '{ gsub(/ /, "", $2); print $2"/tensorrt_llm/libs"}') + containerPipLLMLibPath=$(echo "$containerPipLLMLibPath" | sed 's/[[:space:]]+/_/g') + local containerLDLibPath=$LD_LIBRARY_PATH + containerLDLibPath=$(echo "$containerLDLibPath" | sed 's/[[:space:]]+/_/g') + if [[ "$containerLDLibPath" != *"$containerPipLLMLibPath"* ]]; then + containerLDLibPath="$containerPipLLMLibPath:$containerLDLibPath" + containerLDLibPath="${containerLDLibPath%:}" + fi + export LD_LIBRARY_PATH=$containerLDLibPath + + # Slurm ENROOT/pyxis may inject UCX_TLS=tcp from the host MPI stack + # (intended for host-only MPI jobs). That disables CUDA transports and + # breaks NIXL GPU memory registration. Unset it so UCX can auto-select. + if [ "${UCX_TLS:-}" = "tcp" ]; then + unset UCX_TLS + echo "Unset UCX_TLS (cluster injected UCX_TLS=tcp)" + fi + + # Force PMIx to use the in-memory hash GDS instead of ds12/ds21 + # shared-memory. Under `srun --mpi=pmix` with the DLFW 26.04 OpenMPI + # build, the shared-memory GDS modes can fail to publish UCX worker + # addresses across nodes, producing: + # pml_ucx.c:178 Error: Failed to receive UCX worker address: Not found (-13) + # pml_ucx.c:482 Error: Failed to resolve UCX endpoint for rank N + # See https://github.com/open-mpi/ompi/issues/6981. Setting this is a + # no-op when PMIx isn't used. + export PMIX_MCA_gds=hash +} diff --git a/jenkins/scripts/slurm_run.sh b/jenkins/scripts/slurm_run.sh index e7d5f1a56d2c..25622daff391 100755 --- a/jenkins/scripts/slurm_run.sh +++ b/jenkins/scripts/slurm_run.sh @@ -45,32 +45,11 @@ else sleep 30 fi -containerPipLLMLibPath=$(pip3 show tensorrt_llm | grep "Location" | awk -F ":" '{ gsub(/ /, "", $2); print $2"/tensorrt_llm/libs"}') -containerPipLLMLibPath=$(echo "$containerPipLLMLibPath" | sed 's/[[:space:]]+/_/g') -containerLDLibPath=$LD_LIBRARY_PATH -containerLDLibPath=$(echo "$containerLDLibPath" | sed 's/[[:space:]]+/_/g') -if [[ "$containerLDLibPath" != *"$containerPipLLMLibPath"* ]]; then - containerLDLibPath="$containerPipLLMLibPath:$containerLDLibPath" - containerLDLibPath="${containerLDLibPath%:}" -fi -export LD_LIBRARY_PATH=$containerLDLibPath - -# Slurm ENROOT/pyxis may inject UCX_TLS=tcp from the host MPI stack (intended for -# host-only MPI jobs). That disables CUDA transports and breaks NIXL GPU memory -# registration. Unset it so UCX can auto-select. -if [ "${UCX_TLS:-}" = "tcp" ]; then - unset UCX_TLS - echo "Unset UCX_TLS (cluster injected UCX_TLS=tcp)" -fi - -# Force PMIx to use the in-memory hash GDS instead of ds12/ds21 shared-memory. -# Under `srun --mpi=pmix` with the DLFW 26.04 OpenMPI build, the shared-memory -# GDS modes can fail to publish UCX worker addresses across nodes, producing: -# pml_ucx.c:178 Error: Failed to receive UCX worker address: Not found (-13) -# pml_ucx.c:482 Error: Failed to resolve UCX endpoint for rank N -# See https://github.com/open-mpi/ompi/issues/6981. Setting this is a no-op -# when PMIx isn't used. -export PMIX_MCA_gds=hash +# Library path + UCX/PMIx fixups shared with the disagg cache-transceiver +# precheck (slurm_precheck_run.sh) -- keeping them in one place guarantees the +# precheck observes the same network environment as the real test steps. +source "$llmSrcNode/jenkins/scripts/slurm_env_setup.sh" +slurm_setup_runtime_env echo "Library Path:" echo "$LD_LIBRARY_PATH" env | sort diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md index 5e44248b147d..bf3323098b6d 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md @@ -23,22 +23,24 @@ the KV-head axis, so the receiver regenerates its expected slice locally under any TP/PP resharding. KV shape (layers/heads/head_dim, MLA vs GQA, MTP nextn layers) is read from the real model's `config.json` under `$LLM_MODELS_ROOT`. -## Enabling +## Enabling / disabling -Off by default. Either: +**On by default** for every disaggregated perf-sanity test. To opt out: -- per test yaml (recommended for gradual rollout): +- per test yaml: ```yaml cache_transceiver_precheck: - enabled: true + enabled: false # optional overrides (defaults in precheck_config.PRECHECK_DEFAULTS): # request_lengths: [1024, 8192] # num_requests: 2 # step_timeout_s: 900 ``` -- or globally at launch-script generation time: `TRTLLM_DISAGG_CT_PRECHECK=1`. +- or globally at launch-script generation time: `TRTLLM_DISAGG_CT_PRECHECK=0` + (kill switch; `=1` force-enables). The env var, when set, overrides the yaml + either way. ## Failure output From c98289d4556d13f3711bbe0f964d1df101eac272 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Mon, 20 Jul 2026 00:36:55 -0700 Subject: [PATCH 03/21] [None][infra] Print the test-results artifact URL when a Slurm stage fails The failure branch of the Slurm job tracking script only said 'Pytest failed'; the Artifactory URL of the full test output (results-*.tar.gz) was printed much later by uploadResults, far from where people read the failure. Echo the URL right at the failure line so the console points straight at the full logs. Applies to every Slurm-run stage, not just perf sanity. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- jenkins/L0_Test.groovy | 2 ++ 1 file changed, 2 insertions(+) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index b6ecea081559..087c9b99e741 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -2043,6 +2043,8 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG echo "Pytest succeed in Slurm job \$jobId" else echo "Pytest failed in Slurm job \$jobId" + echo "Full test output (logs not shown above) is uploaded after stage teardown to:" + echo " https://urm.nvidia.com/artifactory/${UPLOAD_PATH}/test-results/results-${stageName}*.tar.gz" fi echo "Status: \$STATUS | Exit_code \$EXIT_CODE" exit 0 From 94a522c59fb5a1b5618e7037e04e033d936a0278 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Mon, 20 Jul 2026 02:07:42 -0700 Subject: [PATCH 04/21] [None][test] Drop pickle from the precheck ZMQ control channel The precheck's ctx leaders bind ZMQ REP sockets reachable from the cluster network and previously pickle.loads()'d whatever arrived -- unpickling network-supplied bytes is arbitrary code execution. Control messages are now JSON with an appended HMAC-SHA256 tag, verified with a constant-time compare before parsing. Each (ctx, gen) session gets a fresh 32-byte key generated by the ctx leader and shared only through the work-dir rendezvous addr file (created 0600 on the shared filesystem -- the job's trust domain), never over the network, so a network-only attacker can neither forge nor tamper frames. ContextPhaseParams crosses the wire as its primitive fields with base64 opaque_state (mirroring disaggregated_params.py / executor/result.py) instead of being pickled. Adds CPU-only unit tests for the wire format (roundtrip, tampered frame, wrong key, short frame, addr file permissions). Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../run_precheck.py | 119 +++++++++++++++--- .../test_cache_transceiver_precheck_config.py | 42 +++++++ 2 files changed, 141 insertions(+), 20 deletions(-) diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index 58a45e30c954..74c06d56a314 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -37,16 +37,28 @@ asymmetric layouts, attention DP, and multi-instance pairing. Rendezvous is file-based under --work-dir (a shared filesystem): each ctx -leader binds one ZMQ REP socket per gen peer and publishes host:port in -rendezvous/ctx{ci}_gen{gj}.addr; gen leaders connect with REQ sockets. Only -tiny control payloads travel over ZMQ -- KV data goes through the -transceiver under test. +leader binds one ZMQ REP socket per gen peer and publishes host:port plus a +per-session HMAC key in rendezvous/ctx{ci}_gen{gj}.addr; gen leaders connect +with REQ sockets. Only tiny control payloads travel over ZMQ -- KV data goes +through the transceiver under test. + +Control messages are JSON with an appended HMAC-SHA256 tag -- NEVER pickle: +the REP port is reachable from the cluster network, and unpickling +network-supplied bytes is arbitrary code execution. The key travels only via +the work-dir addr file (filesystem permissions = the job's trust domain), so +a network-only attacker can neither read it nor forge/tamper messages. +ContextPhaseParams crosses the wire as its primitive fields (opaque_state +base64-encoded), mirroring DisaggregatedParams <-> ContextPhaseParams in +tensorrt_llm/disaggregated_params.py and executor/result.py. """ import argparse +import base64 +import hashlib +import hmac import json import os -import pickle +import secrets import signal import socket import sys @@ -95,6 +107,55 @@ def seed_for(rid, global_layer): return (rid * 1_000_003 + global_layer * 31) & 0x7FFFFFFF +# --------------------------------------------------------------------------- # +# Control-channel wire format: HMAC-SHA256-authenticated JSON (never pickle -- +# the ZMQ port is reachable from the cluster network). +# --------------------------------------------------------------------------- # +_HMAC_TAG_LEN = hashlib.sha256().digest_size + + +def pack_msg(obj, key): + """JSON-encode `obj` and append an HMAC-SHA256 tag.""" + data = json.dumps(obj, separators=(",", ":")).encode() + return data + hmac.new(key, data, hashlib.sha256).digest() + + +def unpack_msg(raw, key): + """Verify the HMAC tag, then JSON-decode. Raises _TransferError on forgery.""" + if len(raw) <= _HMAC_TAG_LEN: + raise _TransferError(f"control frame too short ({len(raw)} bytes)") + data, tag = raw[:-_HMAC_TAG_LEN], raw[-_HMAC_TAG_LEN:] + if not hmac.compare_digest(hmac.new(key, data, hashlib.sha256).digest(), tag): + raise _TransferError("control frame failed HMAC verification (tampered or wrong key)") + return json.loads(data) + + +def params_to_wire(p): + """ContextPhaseParams -> JSON-safe dict (fields per executor/result.py).""" + return { + "first_gen_tokens": list(p.first_gen_tokens or []), + "req_id": p.req_id, + "opaque_state": base64.b64encode(p.opaque_state or b"").decode(), + "draft_tokens": list(p.draft_tokens) if p.draft_tokens is not None else None, + "ctx_dp_rank": p.ctx_dp_rank, + "ctx_info_endpoint": p.disagg_info_endpoint, + } + + +def params_from_wire(d): + """Inverse of params_to_wire (ctor per disaggregated_params.py).""" + import tensorrt_llm.bindings.executor as tllme + + return tllme.ContextPhaseParams( + list(d["first_gen_tokens"]), + int(d["req_id"]), + base64.b64decode(d["opaque_state"]), + d["draft_tokens"], + d["ctx_dp_rank"], + d["ctx_info_endpoint"], + ) + + # --------------------------------------------------------------------------- # # KV fill / verify (heavy imports stay inside functions: --dry-run and unit # tests must work without torch / tensorrt_llm installed) @@ -271,9 +332,12 @@ def addr_path(work_dir, ctx_idx, gen_idx): def write_addr(path, payload): + """Atomically publish an addr file. It carries the session HMAC key, so + restrict it to the owning user before it becomes visible.""" os.makedirs(os.path.dirname(path), exist_ok=True) tmp = f"{path}.tmp.{os.getpid()}" with open(tmp, "w") as f: + os.fchmod(f.fileno(), 0o600) json.dump(payload, f) f.flush() os.fsync(f.fileno()) @@ -607,14 +671,14 @@ def _zmq(self): self.zmq_ctx = zmq.Context.instance() return zmq, self.zmq_ctx - def _leader_send_recv(self, sock, obj): + def _leader_send_recv(self, sock, obj, key): """REQ round-trip on the gen leader; broadcast the reply to all ranks.""" reply = None err = None if self.is_leader: try: - sock.send(pickle.dumps(obj)) - reply = pickle.loads(sock.recv()) + sock.send(pack_msg(obj, key)) + reply = unpack_msg(sock.recv(), key) except Exception as e: # noqa: BLE001 err = repr(e) err, reply = self.comm.bcast((err, reply), root=0) @@ -647,7 +711,7 @@ def hello_timeout_s(plan, num_peers): # --------------------------------------------------------------------------- # # ctx / gen session loops # --------------------------------------------------------------------------- # -def ctx_serve_peer(runner, sock, peer_idx, arm, disarm): +def ctx_serve_peer(runner, sock, peer_idx, arm, disarm, key): """Serve one gen peer's full schedule on a dedicated REP socket.""" plan = runner.plan comm = runner.comm @@ -656,7 +720,7 @@ def leader_recv(): msg, err = None, None if runner.is_leader: try: - msg = pickle.loads(sock.recv()) + msg = unpack_msg(sock.recv(), key) except Exception as e: # noqa: BLE001 err = repr(e) err, msg = comm.bcast((err, msg), root=0) @@ -666,7 +730,7 @@ def leader_recv(): def leader_reply(obj): if runner.is_leader: - sock.send(pickle.dumps(obj)) + sock.send(pack_msg(obj, key)) arm(f"hello gen_{peer_idx}", seconds=hello_timeout_s(plan, runner.side["num_peers"])) msg = leader_recv() @@ -689,7 +753,8 @@ def leader_reply(obj): except _TransferError as e: leader_reply(("abort", str(e))) raise - leader_reply(("params", params_by_pair)) + # JSON object keys are strings; the gen side converts back to int. + leader_reply(("params", {str(p): params_to_wire(v) for p, v in params_by_pair.items()})) runner.ctx_finish_chunk(reqs) arm(f"bye gen_{peer_idx}") @@ -708,7 +773,7 @@ def gen_run_peer(runner, peer_idx, arm, disarm): # Rendezvous with instance-wide consensus: if only the leader raised here, # the other ranks would deadlock in the next bcast. - sock, err = None, None + sock, key, err = None, None, None arm(f"rendezvous ctx_{peer_idx}", seconds=hello_s) if runner.is_leader: try: @@ -716,6 +781,9 @@ def gen_run_peer(runner, peer_idx, arm, disarm): addr_path(runner.work_dir, peer_idx, runner.server_idx), plan["rendezvous_timeout_s"], ) + # Session HMAC key: shared only through the work-dir addr file, + # never over the network. + key = bytes.fromhex(addr["key"]) zmq, zctx = runner._zmq() sock = zctx.socket(zmq.REQ) sock.setsockopt(zmq.LINGER, 0) @@ -730,7 +798,9 @@ def gen_run_peer(runner, peer_idx, arm, disarm): try: arm(f"hello ctx_{peer_idx}", seconds=hello_s) reply = runner._leader_send_recv( - sock, ("hello", {"gen_idx": runner.server_idx, "fingerprint": plan["fingerprint"]}) + sock, + ("hello", {"gen_idx": runner.server_idx, "fingerprint": plan["fingerprint"]}), + key, ) if reply[0] == "abort": raise _TransferError(f"ctx_{peer_idx} aborted handshake: {reply[1]}") @@ -746,17 +816,18 @@ def gen_run_peer(runner, peer_idx, arm, disarm): for li, req_len, rep, chunk in _schedule(plan): arm(f"ctx_{peer_idx} len={req_len} rep={rep}") reply = runner._leader_send_recv( - sock, ("go", {"li": li, "rep": rep, "chunk": chunk[0]}) + sock, ("go", {"li": li, "rep": rep, "chunk": chunk[0]}), key ) if reply[0] == "abort": raise _TransferError(f"ctx_{peer_idx} aborted: {reply[1]}") - ok, detail = runner.gen_run_chunk(peer_idx, li, req_len, rep, chunk, reply[1]) + params_by_pair = {int(p): params_from_wire(v) for p, v in reply[1].items()} + ok, detail = runner.gen_run_chunk(peer_idx, li, req_len, rep, chunk, params_by_pair) if rep >= plan["warmup_requests"]: prev_ok, prev_detail = case_ok.get(req_len, (True, "")) case_ok[req_len] = (prev_ok and ok, prev_detail or detail) arm(f"bye ctx_{peer_idx}") - runner._leader_send_recv(sock, ("done", {})) + runner._leader_send_recv(sock, ("done", {}), key) disarm() for req_len, (ok, detail) in case_ok.items(): @@ -915,7 +986,9 @@ def disarm(): if args.role == "ctx": # One dedicated REP socket per gen peer (avoids REQ interleaving # across sessions on a shared socket), published via addr files. - socks = {} + # Each session gets a fresh HMAC key, shared only through the + # work-dir addr file (0600). + socks, keys = {}, {} if runner.is_leader: zmq, zctx = runner._zmq() host = os.environ.get("SLURMD_NODENAME") or socket.gethostname() @@ -925,15 +998,21 @@ def disarm(): # Generous: gen peers are serialized across ctx servers. s.setsockopt(zmq.RCVTIMEO, hello_timeout_s(plan, num_peers) * 1000) port = s.bind_to_random_port("tcp://*") + keys[gj] = secrets.token_bytes(32) write_addr( addr_path(runner.work_dir, args.server_idx, gj), - {"host": host, "port": port}, + {"host": host, "port": port, "key": keys[gj].hex()}, ) socks[gj] = s for gj in range(num_peers): try: ctx_serve_peer( - runner, socks.get(gj) if runner.is_leader else None, gj, arm, disarm + runner, + socks.get(gj) if runner.is_leader else None, + gj, + arm, + disarm, + keys.get(gj), ) runner.recorder.record(f"gen_{gj}", 0, "PASS", "served all transfers") except _Timeout: diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py index 492e895b1369..fe15ed0e481c 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_config.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py @@ -214,3 +214,45 @@ def test_side_plan_views(): assert ctx_view["peer_parallel"]["world_size"] == 16 assert ctx_view["num_peers"] == 1 and gen_view["num_peers"] == 1 assert gen_view["cache_transceiver_config"]["max_tokens_in_buffer"] == 16384 + + +class TestControlWireFormat: + """run_precheck's HMAC-JSON control frames (importable without torch).""" + + def _rp(self): + import run_precheck as rp + + return rp + + def test_roundtrip(self): + rp = self._rp() + key = b"\x01" * 32 + msg = ["go", {"li": 0, "rep": 1, "chunk": 2}] + assert rp.unpack_msg(rp.pack_msg(msg, key), key) == msg + + def test_tampered_frame_rejected(self): + rp = self._rp() + key = b"\x01" * 32 + raw = rp.pack_msg(["hello", {}], key) + bad = raw[:-1] + bytes([raw[-1] ^ 0xFF]) + with pytest.raises(rp._TransferError): + rp.unpack_msg(bad, key) + + def test_wrong_key_rejected(self): + rp = self._rp() + raw = rp.pack_msg(["hello", {}], b"\x01" * 32) + with pytest.raises(rp._TransferError): + rp.unpack_msg(raw, b"\x02" * 32) + + def test_short_frame_rejected(self): + rp = self._rp() + with pytest.raises(rp._TransferError): + rp.unpack_msg(b"tiny", b"\x01" * 32) + + def test_addr_file_owner_only(self, tmp_path): + rp = self._rp() + path = str(tmp_path / "rendezvous" / "ctx0_gen0.addr") + rp.write_addr(path, {"host": "h", "port": 1, "key": "aa"}) + assert (os.stat(path).st_mode & 0o777) == 0o600 + with open(path) as f: + assert json.load(f)["key"] == "aa" From b870a9f45b7f106346f149330cd2b37dae7d8593 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Mon, 20 Jul 2026 02:10:56 -0700 Subject: [PATCH 05/21] [None][test] Support KVCacheManagerV2 in the cache transceiver precheck Follows the proven pattern from examples/disaggregated/slurm/cache_transceiver_test (which exercises V2 with the Python transceiver end-to-end): a KvCacheConfigV2 dataclass stand-in, prepare_context/resize_context (ctx) and prepare_disagg_gen_init (gen) for allocation, free_resources() for release (also frees the IndexMapper slot), and is_disagg=True for in-flight-transfer slot headroom. fill/verify are unchanged -- get_buffers()/get_batch_cache_indices() alias the real pool on both V1 and V2. vocab_size (a V2 ctor input) now comes from the model's config.json. Manager version and transceiver runtime resolve exactly like serving (PR #15823): explicit yaml values win; an absent use_kv_cache_manager_v2 means 'auto' and adopts the model class's get_model_defaults() (e.g. DeepseekV4 defaults to V2), and transceiver_runtime 'auto' resolves through the REAL llm_utils._resolve_transceiver_runtime_auto against model_cls.get_preferred_transceiver_runtime() with the same NIXL gate. The model class comes from config.json architectures via MODEL_CLASS_MAPPING, mirroring the automodel path. Effective values are printed by the leader and recorded in the status json. V2 only works with the Python transceiver (see cache_transceiver_test/report.py), so a V2 + CPP pairing is rejected up front with a clear INIT_ERROR instead of a C++ binding type error. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../cache_transceiver_precheck/README.md | 1 + .../precheck_config.py | 15 ++ .../run_precheck.py | 190 +++++++++++++++++- .../test_cache_transceiver_precheck_config.py | 38 ++++ 4 files changed, 233 insertions(+), 11 deletions(-) diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md index bf3323098b6d..e18a7396e94a 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md @@ -16,6 +16,7 @@ starts. | Same UCX env vars (incl. the `unset UCX_TLS` cases) | `jenkins/scripts/perf/submit.py` builds the precheck commands from the **same** `ucx_tls_cmd` + `$CTX/GEN_WORKER_ENV_VARS` strings as the worker steps; `slurm_precheck_run.sh` sources the same `slurm_env_setup.sh` (the `UCX_TLS=tcp` fixup) as `slurm_run.sh`. | | Same instance count / parallelism | One precheck `srun` per ctx/gen server with the same `-N/--ntasks/--ntasks-per-node/--mpi=pmix` and the same node slices (`-w`) as the real server steps (`slurm_launch_draft.sh`). TP/PP/CP/attention-DP come from the same `worker_config`. | | Same transceiver config | `CacheTransceiverConfig(**yaml["worker_config"][role]["cache_transceiver_config"])` — the yaml block is passed through verbatim (backend, `max_tokens_in_buffer`, timeouts, ...). | +| Same KV cache manager version + transceiver runtime | Explicit per-side `kv_cache_config.use_kv_cache_manager_v2` wins; absent means "auto" and resolves against the model class's `get_model_defaults()`, and `transceiver_runtime: auto` resolves via `get_preferred_transceiver_runtime()` (NIXL-gated) — both through the same llm_utils code serving uses. V2 requires the Python transceiver (the C++ one only supports V1); a V2+CPP combination fails fast with INIT_ERROR. | Asymmetric layouts (ctx dep4 → gen dep16, ctx pp8 → gen tp32, ...) are supported: data is seeded per (request, **global** layer) and constant along diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py index 543f81e97791..ca4f8128bd5d 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py @@ -178,6 +178,16 @@ def resolve_plan(cfg, benchmark_mode="e2e"): "gen_num_nextn_predict_layers": gen_nextn, "ctx_kv_dtype": str((ctx_side.get("kv_cache_config") or {}).get("dtype", "auto")), "gen_kv_dtype": str((gen_side.get("kv_cache_config") or {}).get("dtype", "auto")), + # Tri-state, matching KvCacheConfig's pydantic default: explicit + # True/False from the yaml wins; absent means "auto", which the + # driver resolves against the model class's get_model_defaults() at + # runtime — exactly like serving (_resolve_kv_cache_manager_v2_auto). + "ctx_use_kv_cache_manager_v2": (ctx_side.get("kv_cache_config") or {}).get( + "use_kv_cache_manager_v2", "auto" + ), + "gen_use_kv_cache_manager_v2": (gen_side.get("kv_cache_config") or {}).get( + "use_kv_cache_manager_v2", "auto" + ), "tokens_per_block": tokens_per_block, "request_lengths": req_lens, "num_requests": int(knobs["num_requests"]), @@ -223,6 +233,7 @@ def side_plan(plan, role): "cache_transceiver_config": plan[f"{role}_cache_transceiver_config"], "kv_dtype": plan[f"{role}_kv_dtype"], "num_nextn_predict_layers": plan[f"{role}_num_nextn_predict_layers"], + "use_kv_cache_manager_v2": plan[f"{role}_use_kv_cache_manager_v2"], "num_peers": plan["num_gen_servers" if role == "ctx" else "num_ctx_servers"], } @@ -322,12 +333,15 @@ def model_kv_shape(model_dir): if num_layers is None: return dict(FALLBACK_KV_SHAPE) + vocab_size = int(hf_cfg.get("vocab_size") or 0) or None # V2 ctor input + if hf_cfg.get("kv_lora_rank"): # MLA (DeepSeek-family, Kimi K2, ...) return { "num_layers": int(num_layers), "num_kv_heads": 1, "head_dim": int(hf_cfg["kv_lora_rank"]) + int(hf_cfg.get("qk_rope_head_dim", 0)), "is_mla": True, + "vocab_size": vocab_size, "source": "config.json (MLA)", } @@ -342,5 +356,6 @@ def model_kv_shape(model_dir): "num_kv_heads": int(num_kv_heads), "head_dim": int(head_dim), "is_mla": False, + "vocab_size": vocab_size, "source": "config.json", } diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index 74c06d56a314..288487931b78 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -54,6 +54,7 @@ import argparse import base64 +import dataclasses import hashlib import hmac import json @@ -206,7 +207,106 @@ def verify_request(kvm, rid): return True, "" -def build_kv_cache_manager(kv_shape, plan, side, mapping, max_req_len): +@dataclasses.dataclass +class _KvCacheConfigV2: + """KvCacheConfig stand-in for KVCacheManagerV2. + + Mirrors examples/disaggregated/slurm/cache_transceiver_test (and the + reference dataclass in + tests/unittest/disaggregated/test_cache_transceiver_single_process.py): + KVCacheManagerV2 reads these fields off the config object directly, so + every attribute it accesses must exist here. + """ + + max_tokens: "int | None" = None + enable_block_reuse: bool = False + max_attention_window: "list | None" = None + sink_token_length: "int | None" = None + free_gpu_memory_fraction: "float | None" = None + host_cache_size: "int | None" = None + disk_cache_size: "int | None" = None + disk_cache_path: "str | None" = None + onboard_blocks: bool = True + cross_kv_cache_fraction: "float | None" = None + secondary_offload_min_priority: "int | None" = None + event_buffer_max_size: int = 0 + kv_cache_event_hash_algo: str = "auto" + max_gpu_total_bytes: "int | None" = None + enable_partial_reuse: bool = False + copy_on_partial_reuse: bool = False + dtype: str = "auto" + pool_ratio: "list | None" = None + avg_seq_len: "int | None" = None + block_reuse_policy: str = "all_reusable" + enable_swa_scratch_reuse: bool = False + disk_prefetch_num_reqs: int = 4 + max_util_for_resume: float = 0.95 + + +def _lookup_model_cls(model_dir): + """Model class from config.json architectures, like serving's automodel path.""" + try: + with open(os.path.join(model_dir, "config.json")) as f: + hf_cfg = json.load(f) + except (OSError, json.JSONDecodeError, TypeError): + return None, None + archs = hf_cfg.get("architectures") or [] + hf_view = type("HFConfigView", (), hf_cfg) # attribute access for the pref hook + if not archs: + return None, hf_view + import tensorrt_llm._torch.models # noqa: F401 - populates the registry + + from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPING + + return MODEL_CLASS_MAPPING.get(archs[0]), hf_view + + +def resolve_model_prefs(model_dir, side, cache_cfg): + """Mirror serving's model-preference resolution (PR #15823 semantics). + + - use_kv_cache_manager_v2 == "auto" (yaml absent): adopt the model + class's get_model_defaults() value, default False + (llm_utils._resolve_kv_cache_manager_v2_auto). + - cache_cfg.transceiver_runtime == "auto": adopt + model_cls.get_preferred_transceiver_runtime(), NIXL-gated, via the + REAL llm_utils._resolve_transceiver_runtime_auto (mutates cache_cfg). + + Returns the effective use_v2 bool. + """ + import types + + model_cls, hf_view = _lookup_model_cls(model_dir) + + setting = side["use_kv_cache_manager_v2"] + if setting == "auto": + defaults = {} + if model_cls is not None: + try: + defaults = model_cls.get_model_defaults(None) or {} + except Exception as e: # noqa: BLE001 - model hooks may need llm_args + print(f"[precheck] WARNING: get_model_defaults failed ({e!r}); " + f"assuming V1", flush=True) + kv_defaults = defaults.get("kv_cache_config") or {} + use_v2 = kv_defaults.get("use_kv_cache_manager_v2", False) is True + else: + use_v2 = bool(setting) + + if getattr(cache_cfg, "transceiver_runtime", None) == "auto": + try: + from tensorrt_llm.llmapi.llm_utils import _resolve_transceiver_runtime_auto + + shim = types.SimpleNamespace(cache_transceiver_config=cache_cfg) + _resolve_transceiver_runtime_auto(shim, model_cls, hf_view) + except Exception as e: # noqa: BLE001 - fall back to the create() default (CPP) + print( + f"[precheck] WARNING: transceiver_runtime 'auto' resolution failed " + f"({e!r}); create_kv_cache_transceiver will fall back to CPP", + flush=True, + ) + return use_v2 + + +def build_kv_cache_manager(kv_shape, plan, side, mapping, max_req_len, use_v2): import tensorrt_llm.bindings import tensorrt_llm.bindings.executor as trtllm_executor from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager @@ -239,9 +339,7 @@ def build_kv_cache_manager(kv_shape, plan, side, mapping, max_req_len): # Real MLA serving uses SELFKONLY (kv_factor=1: one latent plane, no V) — # see _torch/pyexecutor/_util.py; SELF would double the per-token bytes. cache_type = CacheTypeCpp.SELFKONLY if kv_shape["is_mla"] else CacheTypeCpp.SELF - return KVCacheManager( - trtllm_executor.KvCacheConfig(max_tokens=max_tokens, enable_block_reuse=False), - cache_type, + common = dict( num_layers=kv_shape["num_layers"], num_kv_heads=kv_shape["num_kv_heads"], head_dim=kv_shape["head_dim"], @@ -252,6 +350,27 @@ def build_kv_cache_manager(kv_shape, plan, side, mapping, max_req_len): dtype=dtype, spec_config=spec_config, ) + if use_v2: + from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 + + # is_disagg=True doubles the IndexMapper capacity so in-flight + # transfers (TRANS_IN_PROGRESS) can hold slots, like real serving. + return KVCacheManagerV2( + _KvCacheConfigV2( + max_tokens=max_tokens, + enable_block_reuse=False, + max_attention_window=[padded_len], + ), + cache_type, + vocab_size=kv_shape.get("vocab_size") or 32000, + is_disagg=True, + **common, + ) + return KVCacheManager( + trtllm_executor.KvCacheConfig(max_tokens=max_tokens, enable_block_reuse=False), + cache_type, + **common, + ) def make_request(is_ctx, rid, req_len, runtime, ctx_params=None): @@ -295,11 +414,28 @@ def make_request(is_ctx, rid, req_len, runtime, ctx_params=None): ) -def add_sequence(kvm, req, prompt_len): +def add_sequence(kvm, req, prompt_len, use_v2): + """Allocate KV blocks (mirrors the cache_transceiver_test harness).""" + if use_v2: + if req.is_disagg_generation_init_state: + ok = kvm.prepare_disagg_gen_init(req) + else: + ok = kvm.prepare_context(req) and kvm.resize_context(req, prompt_len) + if not ok: + raise RuntimeError(f"V2 KV cache allocation failed for request {req.py_request_id}") + return kvm.impl.add_sequence_batch([(req.py_request_id, prompt_len, 1)], [req]) -def free_sequence(kvm, req): +def free_sequence(kvm, req, use_v2): + if use_v2: + import torch + + # free_resources() closes the kv_cache AND releases the IndexMapper + # slot (closing the cache alone leaks slots). + torch.cuda.current_stream().synchronize() + kvm.free_resources(req) + return kvm.impl.remove_sequence(req.py_request_id, req, True) @@ -473,6 +609,8 @@ def __init__(self, args, plan, side, comm): self.kvm = None self.xcvr = None self.runtime = "CPP" + # Resolved in setup(): "auto" needs the model class (get_model_defaults). + self.use_v2 = False self.mapping = None self.llm_request_state = None self.csv_dir = os.path.join(self.work_dir, "csv", f"{self.role}_{self.server_idx}") @@ -511,10 +649,25 @@ def setup(self, kv_shape, max_req_len): os.makedirs(self.csv_dir, exist_ok=True) os.environ["TRTLLM_KVCACHE_TIME_OUTPUT_PATH"] = self.csv_dir - self.kvm = build_kv_cache_manager(kv_shape, self.plan, self.side, self.mapping, max_req_len) # Built VERBATIM from the disagg yaml's cache_transceiver_config so # backend/max_tokens_in_buffer/timeouts match the real test exactly. cache_cfg = CacheTransceiverConfig(**self.side["cache_transceiver_config"]) + # Yaml-absent settings resolve against the model's preferences, like + # serving does (kv manager version + transceiver runtime). + self.use_v2 = resolve_model_prefs(self.plan.get("_model_dir"), self.side, cache_cfg) + # KVCacheManagerV2 only works with the Python transceiver (see + # cache_transceiver_test/report.py); reject the pairing up front with + # a clear INIT_ERROR instead of a C++ binding type error. + if self.use_v2 and cache_cfg.transceiver_runtime != "PYTHON": + raise RuntimeError( + "KVCacheManagerV2 requires cache_transceiver_config." + f"transceiver_runtime: PYTHON, got {cache_cfg.transceiver_runtime!r} " + "(the C++ transceiver only supports the V1 manager)" + ) + + self.kvm = build_kv_cache_manager( + kv_shape, self.plan, self.side, self.mapping, max_req_len, self.use_v2 + ) AttentionTypeCpp = tensorrt_llm.bindings.internal.batch_manager.AttentionType attention_type = AttentionTypeCpp.MLA if kv_shape["is_mla"] else AttentionTypeCpp.DEFAULT dist_obj = Distributed.get(self.mapping) @@ -550,7 +703,7 @@ def ctx_run_chunk(self, peer_idx, li, req_len, rep, chunk): for pair in owned: rid = self._pair_rid(peer_idx, li, rep, pair) req = make_request(True, rid, req_len, self.runtime) - add_sequence(self.kvm, req, req_len) + add_sequence(self.kvm, req, req_len, self.use_v2) fill_request(self.kvm, rid) tensorrt_llm.logger.info( f"[ctx{self.server_idx} r{self.rank}] rid={rid} len={req_len}: send START" @@ -617,7 +770,7 @@ def gen_run_chunk(self, peer_idx, li, req_len, rep, chunk, params_by_pair): req = make_request( False, rid, req_len, self.runtime, ctx_params=params_by_pair[pair] ) - add_sequence(self.kvm, req, req_len) + add_sequence(self.kvm, req, req_len, self.use_v2) tensorrt_llm.logger.info( f"[gen{self.server_idx} r{self.rank}] rid={rid} len={req_len}: recv START" ) @@ -659,7 +812,7 @@ def gen_run_chunk(self, peer_idx, li, req_len, rep, chunk, params_by_pair): def _free_all(self, reqs): for req in reqs.values(): try: - free_sequence(self.kvm, req) + free_sequence(self.kvm, req, self.use_v2) except Exception: # noqa: BLE001 - teardown best-effort pass @@ -979,6 +1132,15 @@ def disarm(): flush=True, ) return 1 + if runner.is_leader: + # Effective values after model-preference resolution — what serving + # would actually run with (PR #15823 semantics). + print( + f"[precheck {args.role}_{args.server_idx}] " + f"kv_cache_manager={'V2' if runner.use_v2 else 'V1'} " + f"transceiver_runtime={runner.runtime}", + flush=True, + ) # --- sessions ------------------------------------------------------------- num_peers = side["num_peers"] @@ -1063,7 +1225,13 @@ def disarm(): failed_local = 1 if runner.recorder.failed_cases() else 0 failed = comm.allreduce(failed_local, op=MPI.MAX) - runner.recorder.finalize(extra={"per_gpu_bw_gbps": bw} if bw else None) + extra = { + "kv_cache_manager": "V2" if runner.use_v2 else "V1", + "transceiver_runtime": runner.runtime, + } + if bw: + extra["per_gpu_bw_gbps"] = bw + runner.recorder.finalize(extra=extra) comm.Barrier() if runner.is_leader: verdict = "FAIL" if failed else "PASS" diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py index fe15ed0e481c..2287bc387d30 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_config.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py @@ -183,6 +183,7 @@ def test_model_kv_shape_mla_and_gqa(tmp_path): "num_kv_heads": 1, "head_dim": 576, "is_mla": True, + "vocab_size": None, "source": "config.json (MLA)", } @@ -256,3 +257,40 @@ def test_addr_file_owner_only(self, tmp_path): assert (os.stat(path).st_mode & 0o777) == 0o600 with open(path) as f: assert json.load(f)["key"] == "aa" + + +def test_use_kv_cache_manager_v2_flags(): + # Absent -> "auto" (the driver resolves it against the model's + # get_model_defaults at runtime, like serving). + plan = pcfg.resolve_plan(_disagg_yaml()) + assert plan["ctx_use_kv_cache_manager_v2"] == "auto" + assert plan["gen_use_kv_cache_manager_v2"] == "auto" + assert pcfg.side_plan(plan, "ctx")["use_kv_cache_manager_v2"] == "auto" + + # Explicit yaml values win, per side. + plan = pcfg.resolve_plan( + _disagg_yaml( + ctx_extra={"kv_cache_config": {"dtype": "fp8", "use_kv_cache_manager_v2": False}}, + gen_extra={"kv_cache_config": {"dtype": "fp8", "use_kv_cache_manager_v2": True}}, + ) + ) + assert plan["ctx_use_kv_cache_manager_v2"] is False + assert plan["gen_use_kv_cache_manager_v2"] is True + assert pcfg.side_plan(plan, "gen")["use_kv_cache_manager_v2"] is True + + +def test_model_kv_shape_vocab_size(tmp_path): + model_dir = tmp_path / "m" + model_dir.mkdir() + (model_dir / "config.json").write_text( + json.dumps( + { + "num_hidden_layers": 2, + "num_attention_heads": 8, + "num_key_value_heads": 8, + "head_dim": 128, + "vocab_size": 129280, + } + ) + ) + assert pcfg.model_kv_shape(str(model_dir))["vocab_size"] == 129280 From 7099f8148219d7a660cb45f72cc3f437857da329 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Mon, 20 Jul 2026 03:06:43 -0700 Subject: [PATCH 06/21] [None][test] Guard the precheck rendezvous against stale files from reused work dirs A reused --work-dir (Slurm requeue reruns the batch script with the same directories; manual reruns) could hold a previous run's rendezvous addr files: gen leaders would connect to a dead host:port and block until the hello timeout, misreporting TIMEOUT. Stale status files would also pollute the launch script's verdict aggregation. Three layers: - the gate script clears rendezvous/status artifacts before launching, - write_addr removes any stale file before atomically publishing, and - addr files are stamped with SLURM_JOB_ID; wait_for_addr skips files stamped by another run (empty job id -- manual non-Slurm runs -- disables the check), so even a crash path that skipped cleanup cannot hand a gen leader a dead endpoint. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../disaggregated/slurm_ct_precheck_gate.sh | 7 +++ .../run_precheck.py | 34 ++++++++++++-- .../test_cache_transceiver_precheck_config.py | 44 +++++++++++++++++++ 3 files changed, 81 insertions(+), 4 deletions(-) diff --git a/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh b/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh index 65e9158b7973..86cdbba90c61 100644 --- a/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh +++ b/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh @@ -106,6 +106,13 @@ run_cache_transceiver_precheck() { echo "Starting cache transceiver precheck..." precheckDir="$testOutputDir/cache_transceiver_precheck" mkdir -p "$precheckDir/logs" + # A reused work dir (Slurm requeue reruns this batch script with the same + # directories) may hold a previous run's rendezvous/status files: stale + # addr files would point gen leaders at dead ports, and stale status + # files would pollute the verdict aggregation below. The driver also + # job-id-stamps addr files as a second line of defense. + rm -f "$precheckDir"/rendezvous/*.addr "$precheckDir"/status/*.status \ + "$precheckDir"/status/*.json 2>/dev/null || true precheckPids=() precheckNames=() local i gen_world_size ctx_world_size diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index 288487931b78..a05149902fca 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -467,10 +467,28 @@ def addr_path(work_dir, ctx_idx, gen_idx): return os.path.join(work_dir, "rendezvous", f"ctx{ctx_idx}_gen{gen_idx}.addr") +def run_token(): + """Identity of THIS run, stamped into addr files and checked by readers. + + A reused --work-dir (Slurm requeue reruns the batch script with the same + directories; manual reruns) can hold addr files from a previous run -- + connecting to that stale host:port would block until the hello timeout + and misreport TIMEOUT. Within one precheck all instances share + SLURM_JOB_ID; empty (non-Slurm manual runs) disables the check. + """ + return os.environ.get("SLURM_JOB_ID", "") + + def write_addr(path, payload): - """Atomically publish an addr file. It carries the session HMAC key, so - restrict it to the owning user before it becomes visible.""" + """Atomically publish an addr file (clearing any stale one first). It + carries the session HMAC key, so restrict it to the owning user before + it becomes visible.""" os.makedirs(os.path.dirname(path), exist_ok=True) + try: + os.remove(path) # stale file from a previous run in a reused work dir + except FileNotFoundError: + pass + payload = dict(payload, job=run_token()) tmp = f"{path}.tmp.{os.getpid()}" with open(tmp, "w") as f: os.fchmod(f.fileno(), 0o600) @@ -481,14 +499,22 @@ def write_addr(path, payload): def wait_for_addr(path, timeout_s): + """Wait for THIS run's addr file; files stamped with another run's job id + are treated as stale and skipped (keep polling).""" + expect_job = run_token() deadline = time.monotonic() + timeout_s while time.monotonic() < deadline: if os.path.exists(path): try: with open(path) as f: - return json.load(f) + payload = json.load(f) except (OSError, json.JSONDecodeError): - pass # mid-rename/NFS staleness; retry + payload = None # mid-rename/NFS staleness; retry + if payload is not None: + stamped = payload.get("job", "") + if not expect_job or not stamped or stamped == expect_job: + return payload + # Stale addr from a previous run in a reused work dir. time.sleep(1.0) raise _Timeout(f"rendezvous file {path} not published within {timeout_s}s") diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py index 2287bc387d30..17f30babddf9 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_config.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py @@ -294,3 +294,47 @@ def test_model_kv_shape_vocab_size(tmp_path): ) ) assert pcfg.model_kv_shape(str(model_dir))["vocab_size"] == 129280 + + +class TestRendezvousStaleness: + """wait_for_addr must skip addr files stamped by a previous run.""" + + def _rp(self): + import run_precheck as rp + + return rp + + def test_same_job_accepted(self, tmp_path, monkeypatch): + rp = self._rp() + monkeypatch.setenv("SLURM_JOB_ID", "12345") + p = str(tmp_path / "rendezvous" / "ctx0_gen0.addr") + rp.write_addr(p, {"host": "h", "port": 1, "key": "aa"}) + got = rp.wait_for_addr(p, timeout_s=2) + assert got["job"] == "12345" and got["port"] == 1 + + def test_stale_job_skipped_until_timeout(self, tmp_path, monkeypatch): + rp = self._rp() + monkeypatch.setenv("SLURM_JOB_ID", "11111") + p = str(tmp_path / "rendezvous" / "ctx0_gen0.addr") + rp.write_addr(p, {"host": "h", "port": 1, "key": "aa"}) # stamped 11111 + monkeypatch.setenv("SLURM_JOB_ID", "22222") # new run + with pytest.raises(rp._Timeout): + rp.wait_for_addr(p, timeout_s=2) + + def test_no_job_id_accepts_any(self, tmp_path, monkeypatch): + rp = self._rp() + monkeypatch.setenv("SLURM_JOB_ID", "11111") + p = str(tmp_path / "rendezvous" / "ctx0_gen0.addr") + rp.write_addr(p, {"host": "h", "port": 1, "key": "aa"}) + monkeypatch.delenv("SLURM_JOB_ID") # manual non-slurm run + assert rp.wait_for_addr(p, timeout_s=2)["port"] == 1 + + def test_write_addr_replaces_stale_file(self, tmp_path, monkeypatch): + rp = self._rp() + monkeypatch.setenv("SLURM_JOB_ID", "11111") + p = str(tmp_path / "rendezvous" / "ctx0_gen0.addr") + rp.write_addr(p, {"host": "old", "port": 1, "key": "aa"}) + monkeypatch.setenv("SLURM_JOB_ID", "22222") + rp.write_addr(p, {"host": "new", "port": 2, "key": "bb"}) + got = rp.wait_for_addr(p, timeout_s=2) + assert got["host"] == "new" and got["job"] == "22222" From 723125e72e9445f6ed3cddbcaa6c8dd01a412e71 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Mon, 20 Jul 2026 04:49:53 -0700 Subject: [PATCH 07/21] [None][test] Keep ctx precheck instances alive until the gen finishes all peers With multiple ctx servers, the gen sent 'done' to each ctx right after that ctx's schedule, so ctx_0's process exited while the gen went on to ctx_1 -- leaving the gen's C++ transceiver holding connections to a dead NIXL agent, a state real serving never produces (ctx servers outlive the whole run). Observed on a ctx2xdep4 -> gen1xdep16 topology: the ctx_0 session passed, then ctx_1's first chunk stalled until the chunk timeout. The gen now defers 'done' to every ctx until ALL its peers' schedules have finished (gen_release_peer), and the ctx budgets its final wait like a handshake, since it can legitimately span the gen's remaining sessions. Session order stays deterministic (every ctx serves gens in index order, every gen visits ctxs in index order), so the deferred release cannot deadlock. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../run_precheck.py | 51 ++++++++++++++++--- 1 file changed, 45 insertions(+), 6 deletions(-) diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index a05149902fca..1a0ef574748f 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -936,7 +936,13 @@ def leader_reply(obj): leader_reply(("params", {str(p): params_to_wire(v) for p, v in params_by_pair.items()})) runner.ctx_finish_chunk(reqs) - arm(f"bye gen_{peer_idx}") + # The gen defers "done" until it has finished the schedules of ALL its + # ctx peers, so every ctx instance stays alive for the whole precheck -- + # matching real serving, where ctx servers outlive the entire run. (An + # early-exiting ctx leaves the gen's C++ transceiver holding connections + # to a dead agent, a state the real test never produces.) The wait can + # therefore span the gen's remaining sessions: budget like a handshake. + arm(f"bye gen_{peer_idx}", seconds=hello_timeout_s(plan, runner.side["num_peers"])) msg = leader_recv() if msg[0] != "done": raise _TransferError(f"expected done from gen_{peer_idx}, got {msg[:1]}") @@ -945,7 +951,13 @@ def leader_reply(obj): def gen_run_peer(runner, peer_idx, arm, disarm): - """Run the full schedule against ctx server `peer_idx`.""" + """Run the full schedule against ctx server `peer_idx`. + + Returns (sock, key) with the session STILL OPEN on success -- the caller + sends the deferred "done" only after every ctx peer's schedule finished, + keeping all ctx instances alive for the whole precheck (real-serving + lifecycle; see ctx_serve_peer). On failure the socket is closed here. + """ plan = runner.plan comm = runner.comm hello_s = hello_timeout_s(plan, runner.side["num_peers"]) @@ -1004,9 +1016,6 @@ def gen_run_peer(runner, peer_idx, arm, disarm): if rep >= plan["warmup_requests"]: prev_ok, prev_detail = case_ok.get(req_len, (True, "")) case_ok[req_len] = (prev_ok and ok, prev_detail or detail) - - arm(f"bye ctx_{peer_idx}") - runner._leader_send_recv(sock, ("done", {}), key) disarm() for req_len, (ok, detail) in case_ok.items(): @@ -1016,6 +1025,19 @@ def gen_run_peer(runner, peer_idx, arm, disarm): "PASS" if ok else "MISMATCH", "" if ok else detail, ) + return sock, key + except BaseException: + if sock is not None: + sock.close(linger=0) + raise + + +def gen_release_peer(runner, peer_idx, sock, key, arm, disarm): + """Deferred session teardown: send "done" and close (best-effort).""" + try: + arm(f"bye ctx_{peer_idx}") + runner._leader_send_recv(sock, ("done", {}), key) + disarm() finally: if sock is not None: sock.close(linger=0) @@ -1218,9 +1240,11 @@ def disarm(): disarm() runner.recorder.record(f"gen_{gj}", 0, "TRANSFER_ERROR", repr(e)) else: + open_sessions = [] for ci in range(num_peers): try: - gen_run_peer(runner, ci, arm, disarm) + sock, sess_key = gen_run_peer(runner, ci, arm, disarm) + open_sessions.append((ci, sock, sess_key)) except _Timeout: disarm() runner.recorder.record( @@ -1232,6 +1256,21 @@ def disarm(): except Exception as e: # noqa: BLE001 - per-peer isolation disarm() runner.recorder.record(f"ctx_{ci}", 0, "TRANSFER_ERROR", repr(e)) + # Deferred "done": released only after EVERY ctx peer's schedule + # finished, so all ctx instances stay alive for the whole + # precheck (real-serving lifecycle -- no dead-agent connections + # in anyone's transceiver while transfers are still running). + for ci, sock, sess_key in open_sessions: + try: + gen_release_peer(runner, ci, sock, sess_key, arm, disarm) + except _Timeout: + disarm() + runner.recorder.record( + f"ctx_{ci}", 0, "TIMEOUT", "releasing session (deferred done)" + ) + except Exception as e: # noqa: BLE001 - best-effort release + disarm() + runner.recorder.record(f"ctx_{ci}", 0, "TRANSFER_ERROR", repr(e)) finally: disarm() hang_detector.stop() From 73928d18d4ccf61123807fc240dd4baa68f1725b Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Mon, 20 Jul 2026 05:00:11 -0700 Subject: [PATCH 08/21] [None][test] Add stall-debugging logs to the precheck driver - gen: periodic progress while waiting for a receive (rid, request state, elapsed) so a stalled transfer shows whether it never matched (request-info phase) or the RDMA write never completed; per-chunk recv timing. - ctx: per-chunk send timing with final request states. - PRECHECK_DEBUG=1 raises the C++ (TLLM_LOG_LEVEL, set before the tensorrt_llm import) and Python logger levels for deep debugging. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../run_precheck.py | 36 ++++++++++++++++++- 1 file changed, 35 insertions(+), 1 deletion(-) diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index 1a0ef574748f..6f76162140e7 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -445,6 +445,9 @@ def _wait_gen_complete(xcvr, req, runtime, llm_request_state): PYTHON transceiver: check_gen_transfer_status(None) blocks for all. C++: the int API can return before THIS request completes on a cold link, so poll for a terminal state (bounded by signal.alarm + hang detector). + Logs periodic progress so a stalled transfer shows WHICH request is stuck + in WHICH state (the difference between "requests never matched" and + "RDMA write never completed"). """ if runtime == "PYTHON": xcvr.check_gen_transfer_status(None) @@ -453,10 +456,20 @@ def _wait_gen_complete(xcvr, req, runtime, llm_request_state): llm_request_state.DISAGG_GENERATION_TRANS_COMPLETE, llm_request_state.DISAGG_TRANS_ERROR, ) + t0 = time.monotonic() + next_report = t0 + 15.0 while req.state not in terminal: xcvr.check_gen_transfer_status(1) if req.state in terminal: break + now = time.monotonic() + if now >= next_report: + print( + f"[precheck] rid={req.py_request_id} recv still waiting: " + f"state={req.state} elapsed={now - t0:.0f}s", + flush=True, + ) + next_report = now + 15.0 time.sleep(0.001) @@ -766,6 +779,9 @@ def ctx_run_chunk(self, peer_idx, li, req_len, rep, chunk): def ctx_finish_chunk(self, reqs): """Wait for all in-flight sends of this chunk, then free.""" + import tensorrt_llm + + t0 = time.monotonic() local_err = None try: self.xcvr.check_context_transfer_status(None) # block-all @@ -778,6 +794,11 @@ def ctx_finish_chunk(self, reqs): local_err = e finally: self._free_all(reqs) + states = {p: str(r.state) for p, r in reqs.items()} + tensorrt_llm.logger.info( + f"[ctx{self.server_idx} r{self.rank}] chunk sends finished in " + f"{time.monotonic() - t0:.1f}s states={states}" + ) reason = self._consensus_error(local_err) if reason is not None: raise _TransferError(f"ctx transfer failed: {reason}") @@ -810,9 +831,15 @@ def gen_run_chunk(self, peer_idx, li, req_len, rep, chunk, params_by_pair): raise _TransferError(f"gen receive setup failed: {reason}") mismatch = "" + t0 = time.monotonic() try: for pair, req in reqs.items(): _wait_gen_complete(self.xcvr, req, self.runtime, self.llm_request_state) + if reqs: + tensorrt_llm.logger.info( + f"[gen{self.server_idx} r{self.rank}] chunk recvs finished in " + f"{time.monotonic() - t0:.1f}s" + ) torch.cuda.synchronize() # receive may land on a side stream bad = [ p for p, r in reqs.items() if r.state == self.llm_request_state.DISAGG_TRANS_ERROR @@ -1090,6 +1117,13 @@ def main(argv=None): # failure summary uses to spot host-staged tcp fallbacks. os.environ.setdefault("UCX_PROTO_INFO", "used") + # PRECHECK_DEBUG=1: verbose C++/Python transceiver logs for stall + # debugging. Must be set before importing tensorrt_llm (the C++ logger + # reads TLLM_LOG_LEVEL at init). + debug = os.environ.get("PRECHECK_DEBUG") == "1" + if debug: + os.environ.setdefault("TLLM_LOG_LEVEL", "DEBUG") + import torch from mpi4py import MPI @@ -1106,7 +1140,7 @@ def main(argv=None): f"{args.role} server step." ) torch.cuda.set_device(rank % torch.cuda.device_count()) - tensorrt_llm.logger.set_level("info") + tensorrt_llm.logger.set_level("debug" if debug else "info") ucx_env = " ".join(f"{k}={v}" for k, v in sorted(os.environ.items()) if k.startswith("UCX_")) print( From ea597c0d86cbea99ae42e30824559d3cf5919bee Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Mon, 20 Jul 2026 07:05:04 -0700 Subject: [PATCH 09/21] [None][test] Budget NIXL agent wire-up in the precheck's first rep Root cause of the pp8 ctx -> dep16 gen false positive (reproduced on three node sets while the REAL test with the same yaml passes): the C++ NIXL path pays a one-time fetchRemoteMD metadata exchange per (receiver rank, ctx rank) agent pair, serialized under a lock with an unbounded retry loop, and a cold cross-rack fetch was measured at 100-170s per pair. MLA + ctx pp8 makes every receiving rank wire up to all 8 ctx agents (dep16 -> 128 pairs over 4 nodes / 2 racks), so the first chunk blew the 180s chunk alarm; the watchdog SIGKILL mid-fetch also produced the UCX-keepalive segfaults seen on the ctx side. Real serving absorbs the wire-up as slow first requests (no hard alarm, observe-only kv_transfer_timeout), so the precheck must too. The first rep of the schedule (the warmup rep) now gets an extra wireup_timeout_s budget (default min(1800, 150 * max world size), yaml overridable) on both sides' chunk alarms, the handshake budget, the gen ZMQ backstop, and the submit.py external step timeout. Later reps keep the tight 180s alarm, preserving the fast-fail property once connections exist. Also fix rid construction: the C++ notification tag is rid & 0xFFF and all non-pair strides were multiples of 4096, so every request with the same pair index across reps/lengths shared one (agent, tag) key -- stray notifications from a previous rep could be consumed by the wrong request. Rids are now peer-session base + dense in-session sequence, keeping tags unique within any 4096-request window (sessions talk to distinct agents, so the peer stride cannot alias). Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- jenkins/scripts/perf/local/submit.py | 10 +++- jenkins/scripts/perf/submit.py | 9 +++- .../cache_transceiver_precheck/README.md | 15 +++++- .../precheck_config.py | 12 +++++ .../run_precheck.py | 50 ++++++++++++------- .../test_cache_transceiver_precheck_config.py | 37 ++++++++++++++ 6 files changed, 113 insertions(+), 20 deletions(-) diff --git a/jenkins/scripts/perf/local/submit.py b/jenkins/scripts/perf/local/submit.py index 9ca0c8894100..8f99af0b9f2e 100755 --- a/jenkins/scripts/perf/local/submit.py +++ b/jenkins/scripts/perf/local/submit.py @@ -928,10 +928,18 @@ def main(): "--work-dir $testOutputDir/cache_transceiver_precheck " f"--benchmark-mode {precheck_mode} --llm-src $llmSrcNode" ) + # The external step timeout must cover the driver's first-rep NIXL + # wire-up allowance (precheck_config: min(1800, 150 * max world)). + max_world = max( + hardware_config.get("gpus_per_ctx_server", 0) or 0, + hardware_config.get("gpus_per_gen_server", 0) or 0, + ) + default_step_timeout = 900 + min(1800, 150 * max_world) script_prefix_lines.extend( [ f"export ctPrecheckEnabled={1 if precheck_enabled else 0}", - f"export ctPrecheckTimeout={int(precheck_cfg.get('step_timeout_s', 900))}", + f"export ctPrecheckTimeout=" + f"{int(precheck_cfg.get('step_timeout_s', default_step_timeout))}", "export precheckRunScript=$llmSrcNode/jenkins/scripts/perf/" "disaggregated/slurm_precheck_run.sh", f'export pytestCommandCTXPrecheck="{ucx_tls_cmd} $CTX_WORKER_ENV_VARS' diff --git a/jenkins/scripts/perf/submit.py b/jenkins/scripts/perf/submit.py index 5ddb40ca197c..81aeca1ed029 100755 --- a/jenkins/scripts/perf/submit.py +++ b/jenkins/scripts/perf/submit.py @@ -647,10 +647,17 @@ def main(): "--work-dir $testOutputDir/cache_transceiver_precheck " f"--benchmark-mode {benchmark_mode} --llm-src $llmSrcNode" ) + # The external step timeout must cover the driver's first-rep NIXL + # wire-up allowance (precheck_config: min(1800, 150 * max world)). + max_world = max( + hardware_config["gpus_per_ctx_server"], hardware_config["gpus_per_gen_server"] + ) + default_step_timeout = 900 + min(1800, 150 * max_world) script_prefix_lines.extend( [ f"export ctPrecheckEnabled={1 if precheck_enabled else 0}", - f"export ctPrecheckTimeout={int(precheck_cfg.get('step_timeout_s', 900))}", + f"export ctPrecheckTimeout=" + f"{int(precheck_cfg.get('step_timeout_s', default_step_timeout))}", # Suite name for the synthetic junit xml the launch script # writes when the precheck fails. f'export stageName="{args.stage_name}"', diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md index e18a7396e94a..4ebfdd10656b 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md @@ -36,13 +36,26 @@ layers) is read from the real model's `config.json` under `$LLM_MODELS_ROOT`. # optional overrides (defaults in precheck_config.PRECHECK_DEFAULTS): # request_lengths: [1024, 8192] # num_requests: 2 - # step_timeout_s: 900 + # chunk_timeout_s: 180 + # wireup_timeout_s: 1800 # first-rep NIXL agent wire-up allowance + # step_timeout_s: 2700 # external srun timeout (default derives from topology) ``` - or globally at launch-script generation time: `TRTLLM_DISAGG_CT_PRECHECK=0` (kill switch; `=1` force-enables). The env var, when set, overrides the yaml either way. +## Timeouts + +The first rep of the schedule (the warmup rep) additionally budgets +`wireup_timeout_s` (default `min(1800, 150 * max world size)`): the C++ NIXL +path pays a one-time serialized `fetchRemoteMD` metadata exchange per +(receiver rank, ctx rank) agent pair, and cold cross-rack fetches were +measured at 100-170s each — real serving absorbs this as slow first requests, +so the precheck does too. Later reps run under the tight `chunk_timeout_s`, +which is what actually catches hangs. Set `PRECHECK_DEBUG=1` in the worker +env to raise the C++/Python transceiver log levels when debugging a stall. + ## Failure output The sbatch log gets a summary block: per-instance verdicts diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py index ca4f8128bd5d..d732caf28b07 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py @@ -44,6 +44,13 @@ "max_concurrent_pairs": 8, # signal.alarm / hang-detector budget for one chunk of transfers. "chunk_timeout_s": 180, + # Extra budget for the FIRST rep of the schedule, which pays the one-time + # NIXL agent metadata wire-up (fetchRemoteMD over the mgmt network): with + # MLA + ctx pp>1 every receiving rank must connect to every ctx rank, and + # a cold cross-rack fetch was measured at 100-170s per agent pair. Real + # serving absorbs this as slow first requests (no hard alarm), so the + # precheck must too. None -> derived: min(1800, 150 * max world size). + "wireup_timeout_s": None, # How long the gen side waits for the ctx rendezvous files (covers ctx # KV-pool allocation + NIXL/UCX transceiver bring-up). "rendezvous_timeout_s": 600, @@ -195,6 +202,11 @@ def resolve_plan(cfg, benchmark_mode="e2e"): "n_pairs": n_pairs, "chunk_size": chunk_size, "chunk_timeout_s": int(knobs["chunk_timeout_s"]), + "wireup_timeout_s": int( + knobs["wireup_timeout_s"] + if knobs["wireup_timeout_s"] is not None + else min(1800, 150 * max(ctx["world_size"], gen["world_size"])) + ), "rendezvous_timeout_s": int(knobs["rendezvous_timeout_s"]), "verify_data": bool(knobs["verify_data"]), } diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index 6f76162140e7..2c82ed7e3c15 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -71,14 +71,14 @@ import precheck_config as pcfg # noqa: E402 -# Request-id strides (must keep rids unique across the whole precheck). -RID_PAIR_STRIDE = 1 -RID_MAX_PAIRS = 4096 -RID_REP_STRIDE = RID_MAX_PAIRS -RID_MAX_REPS = 64 -RID_LEN_STRIDE = RID_REP_STRIDE * RID_MAX_REPS -RID_MAX_LENS = 64 -RID_PEER_STRIDE = RID_LEN_STRIDE * RID_MAX_LENS +# Request-id scheme: rids must be unique across the whole precheck AND +# dense within a (ctx, gen) session -- the C++ transceiver derives its +# notification tag from the LOW 12 BITS of the request id (tagFromRequestId, +# dataTransceiver.cpp), and notifications are matched by (remote agent, tag). +# A dense per-session sequence keeps tags unique among any 4096 consecutive +# requests of a session; the peer stride only separates sessions, which talk +# to distinct agents and therefore cannot alias tags with each other. +RID_PEER_STRIDE = 1 << 24 class _Timeout(Exception): @@ -97,9 +97,10 @@ def _alarm_handler(signum, frame): raise _Timeout() -def make_rid(ctx_idx, gen_idx, num_ctx, li, rep, pair): +def make_rid(ctx_idx, gen_idx, num_ctx, seq): + """Unique rid: peer-session base + dense in-session sequence number.""" peer = gen_idx * num_ctx + ctx_idx - return 1 + peer * RID_PEER_STRIDE + li * RID_LEN_STRIDE + rep * RID_REP_STRIDE + pair + return 1 + peer * RID_PEER_STRIDE + seq def seed_for(rid, global_layer): @@ -723,7 +724,9 @@ def setup(self, kv_shape, max_req_len): def _pair_rid(self, peer_idx, li, rep, pair): ctx_idx = self.server_idx if self.is_ctx else peer_idx gen_idx = peer_idx if self.is_ctx else self.server_idx - return make_rid(ctx_idx, gen_idx, self.plan["num_ctx_servers"], li, rep, pair) + total_reps = self.plan["warmup_requests"] + self.plan["num_requests"] + seq = (li * total_reps + rep) * self.plan["n_pairs"] + pair + return make_rid(ctx_idx, gen_idx, self.plan["num_ctx_servers"], seq) def _owned(self, chunk): return pcfg.owned_pairs(self.plan, self.role, self.mapping.tp_rank, chunk) @@ -909,9 +912,17 @@ def hello_timeout_s(plan, num_peers): Handshakes are serialized across peers (one gen talks to one ctx at a time), so waiting for a peer's hello/welcome can legitimately span other - peers' full sessions -- budget rendezvous + per-peer slack. + peers' full sessions -- budget rendezvous + per-peer slack (including + the peer's first-rep wire-up). """ - return plan["rendezvous_timeout_s"] + num_peers * 300 + return plan["rendezvous_timeout_s"] + num_peers * (300 + plan["wireup_timeout_s"]) + + +def chunk_timeout_s(plan, li, rep): + """Per-chunk budget: the first rep additionally pays the one-time NIXL + agent wire-up (see PRECHECK_DEFAULTS['wireup_timeout_s']).""" + extra = plan["wireup_timeout_s"] if (li == 0 and rep == 0) else 0 + return plan["chunk_timeout_s"] + extra # --------------------------------------------------------------------------- # @@ -946,7 +957,7 @@ def leader_reply(obj): leader_reply(("welcome", {"fingerprint": plan["fingerprint"]})) for li, req_len, rep, chunk in _schedule(plan): - arm(f"gen_{peer_idx} len={req_len} rep={rep}") + arm(f"gen_{peer_idx} len={req_len} rep={rep}", seconds=chunk_timeout_s(plan, li, rep)) msg = leader_recv() if msg[0] == "abort": raise _PeerAbort(f"gen_{peer_idx} aborted: {msg[1]}") @@ -1024,15 +1035,20 @@ def gen_run_peer(runner, peer_idx, arm, disarm): raise _TransferError(f"ctx_{peer_idx} aborted handshake: {reply[1]}") if reply[0] != "welcome": raise _TransferError(f"unexpected handshake reply from ctx_{peer_idx}: {reply[:1]}") - # Established sessions are dedicated: chunk replies are prompt. + # Established sessions are dedicated: chunk replies are prompt. The + # ZMQ timeout is only a backstop under the per-chunk alarm, so it + # includes the first-rep wire-up allowance unconditionally. if runner.is_leader: import zmq as zmq_mod - sock.setsockopt(zmq_mod.RCVTIMEO, (plan["chunk_timeout_s"] + 30) * 1000) + sock.setsockopt( + zmq_mod.RCVTIMEO, + (plan["chunk_timeout_s"] + plan["wireup_timeout_s"] + 30) * 1000, + ) case_ok = {} for li, req_len, rep, chunk in _schedule(plan): - arm(f"ctx_{peer_idx} len={req_len} rep={rep}") + arm(f"ctx_{peer_idx} len={req_len} rep={rep}", seconds=chunk_timeout_s(plan, li, rep)) reply = runner._leader_send_recv( sock, ("go", {"li": li, "rep": rep, "chunk": chunk[0]}), key ) diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py index 17f30babddf9..66f788a1f92d 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_config.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py @@ -338,3 +338,40 @@ def test_write_addr_replaces_stale_file(self, tmp_path, monkeypatch): rp.write_addr(p, {"host": "new", "port": 2, "key": "bb"}) got = rp.wait_for_addr(p, timeout_s=2) assert got["host"] == "new" and got["job"] == "22222" + + +def test_wireup_timeout_derivation(): + plan = pcfg.resolve_plan(_disagg_yaml()) # ctx dep4 -> gen dep16 + assert plan["wireup_timeout_s"] == min(1800, 150 * 16) + plan = pcfg.resolve_plan( + _disagg_yaml(gen_extra={"tensor_parallel_size": 4}) + ) + assert plan["wireup_timeout_s"] == 600 + plan = pcfg.resolve_plan(_disagg_yaml(cache_transceiver_precheck={"wireup_timeout_s": 42})) + assert plan["wireup_timeout_s"] == 42 + + +def test_rid_tags_dense_within_session(): + """The C++ notification tag is rid & 0xFFF: rids must be dense within a + (ctx, gen) session so tags cannot alias across reps/lengths.""" + import run_precheck as rp + + plan = pcfg.resolve_plan(_disagg_yaml()) # n_pairs=16 + total_reps = plan["warmup_requests"] + plan["num_requests"] + n_pairs = plan["n_pairs"] + + def session_rids(ctx_idx, gen_idx): + out = [] + for li in range(2): + for rep in range(total_reps): + for pair in range(n_pairs): + seq = (li * total_reps + rep) * n_pairs + pair + out.append(rp.make_rid(ctx_idx, gen_idx, 2, seq)) + return out + + a = session_rids(0, 0) + b = session_rids(1, 0) + assert len(set(a)) == len(a) and len(set(b)) == len(b) + assert not (set(a) & set(b)) # globally unique across sessions + tags = [r & 0xFFF for r in a] + assert len(set(tags)) == len(tags) # no tag aliasing within a session From 60bf16ebc9d50158ae97d1e3f26fb782fb2aa96f Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Mon, 20 Jul 2026 07:32:04 -0700 Subject: [PATCH 10/21] [None][test] Simplify the precheck: one owner per policy, leaner hot path Quality pass over the precheck stack (no behavior change; all topologies re-validated on hardware afterwards): - precheck_config now owns every launch-facing policy once: precheck_prefix_lines() builds the gate export lines (enable / kill-switch semantics, step-timeout default) for BOTH jenkins/scripts/perf/submit.py and the local flow, which previously carried a drifting ~40-line copy each; wireup_timeout_s()/ default_step_timeout_s() replace the formula that existed in three files. - run_precheck: reuse the maintained converters instead of local re-implementations (DisaggregatedParams.get_context_phase_params for the wire decode, llm_utils._resolve_kv_cache_manager_v2_auto for the V2 'auto' rule, the real pydantic KvCacheConfig instead of a 25-field stand-in dataclass); one exception->verdict mapper replaces three hand-rolled per-peer blocks; fill/verify expand the pattern on device (was transferring heads x the unique bytes), compute shapes without a throwaway gather, skip byte-verification for warmup reps (their result is discarded), and V2 frees sync the stream once per chunk instead of per request; assorted small cleanups (dead state, redundant conditionals, duplicate loop tests, stale-file pre-remove that os.replace + job stamps already cover). - gate script: shared ct_step_passed/ct_step_excerpt keep the console summary and the junit xml from drifting; one ct_launch_step function replaces the twin gen/ctx srun loops. - per-side plan fields in resolve_plan are built by one loop; side_plan drops the unused peer_parallel; tests import run_precheck at module level (it is stdlib-only by design). Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../disaggregated/slurm_ct_precheck_gate.sh | 81 ++++--- jenkins/scripts/perf/local/submit.py | 63 +++-- jenkins/scripts/perf/submit.py | 62 +++-- .../precheck_config.py | 86 +++++-- .../run_precheck.py | 217 ++++++++---------- .../test_cache_transceiver_precheck_config.py | 23 +- 6 files changed, 252 insertions(+), 280 deletions(-) diff --git a/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh b/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh index 86cdbba90c61..a5192eae7e7d 100644 --- a/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh +++ b/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh @@ -25,6 +25,21 @@ ct_first_errors() { "$1" 2>/dev/null || true } +# Shared verdict predicate + failing-log excerpt, so the console summary and +# the junit xml cannot drift apart. +ct_step_passed() { + local statusFile="$precheckDir/status/$1.status" + [ -f "$statusFile" ] && grep -q "^PASS" "$statusFile" +} + +ct_step_excerpt() { + local stepLog="$precheckDir/logs/$1.log" + echo "First error lines (line-numbered):" + ct_first_errors "$stepLog" + echo "Log tail ($stepLog):" + tail -n 60 "$stepLog" 2>/dev/null || true +} + # Console summary for a failed precheck: per-instance verdicts, first error # lines + tail of each failing step log, and UCX red-flag lines. # Uses: precheckDir, precheckNames. @@ -37,16 +52,9 @@ ct_print_failure_summary() { echo "" echo "Failing step logs:" for k in "${!precheckNames[@]}"; do - statusFile="$precheckDir/status/${precheckNames[$k]}.status" - if [ -f "$statusFile" ] && grep -q "^PASS" "$statusFile"; then - continue - fi - stepLog="$precheckDir/logs/${precheckNames[$k]}.log" - echo "----- ${precheckNames[$k]} ($stepLog) -----" - echo "First error lines (line-numbered):" - ct_first_errors "$stepLog" - echo "Tail:" - tail -n 60 "$stepLog" 2>/dev/null || true + ct_step_passed "${precheckNames[$k]}" && continue + echo "----- ${precheckNames[$k]} -----" + ct_step_excerpt "${precheckNames[$k]}" done echo "" echo "UCX red flags (host-staged tcp fallback / UCX errors), if any:" @@ -65,20 +73,17 @@ ct_write_junit_xml() { suiteName="$(printf '%s' "${stageName:-${SLURM_JOB_NAME:-disagg_perf_sanity}}" | ct_xml_escape)" local junitFailures=0 local junitCases="" - local name statusFile stepLog verdict detail + local name verdict detail for k in "${!precheckNames[@]}"; do name="${precheckNames[$k]}" - statusFile="$precheckDir/status/$name.status" - if [ -f "$statusFile" ] && grep -q "^PASS" "$statusFile"; then + if ct_step_passed "$name"; then junitCases+=""$'\n' continue fi junitFailures=$((junitFailures + 1)) - verdict="$( (head -n 1 "$statusFile" 2>/dev/null \ + verdict="$( (head -n 1 "$precheckDir/status/$name.status" 2>/dev/null \ || echo "NO_STATUS: step died before writing a verdict (see log)") | ct_xml_escape)" - stepLog="$precheckDir/logs/$name.log" - detail="$( { echo "First error lines:"; ct_first_errors "$stepLog"; echo ""; \ - echo "Log tail ($stepLog):"; tail -n 60 "$stepLog" 2>/dev/null; } | ct_xml_escape)" + detail="$(ct_step_excerpt "$name" | ct_xml_escape)" junitCases+="" junitCases+="$detail"$'\n' done @@ -115,36 +120,30 @@ run_cache_transceiver_precheck() { "$precheckDir"/status/*.json 2>/dev/null || true precheckPids=() precheckNames=() - local i gen_world_size ctx_world_size - for i in $(seq 0 $((numGenServers - 1))); do - gen_world_size=$((nodesPerGenServer * gpusPerNodePerGenServer)) - export DISAGG_SERVING_TYPE="GEN_PRECHECK_$i" - export pytestCommand="$pytestCommandGENPrecheck --server-idx $i" + # ct_launch_step + ct_launch_step() { + local role=$1 i=$2 nodes=$3 gpusPerNode=$4 nodeList=$5 pytestCmd=$6 + export DISAGG_SERVING_TYPE="${role^^}_PRECHECK_$i" + export pytestCommand="$pytestCmd --server-idx $i" timeout -k 60 "${ctPrecheckTimeout:-900}" \ srun "${srunArgs[@]}" --mpi=pmix --kill-on-bad-exit=1 \ - -N $nodesPerGenServer \ - -w "${genNodeLists[$i]}" \ - --ntasks=$gen_world_size \ - --ntasks-per-node=$gpusPerNodePerGenServer \ - bash $precheckRunScript &> "$precheckDir/logs/gen_$i.log" & + -N "$nodes" \ + -w "$nodeList" \ + --ntasks=$((nodes * gpusPerNode)) \ + --ntasks-per-node="$gpusPerNode" \ + bash $precheckRunScript &> "$precheckDir/logs/${role}_$i.log" & precheckPids+=($!) - precheckNames+=("gen_$i") + precheckNames+=("${role}_$i") sleep 5 # Wait for pyxis container namespace initialization to avoid race condition + } + local i + for i in $(seq 0 $((numGenServers - 1))); do + ct_launch_step gen "$i" "$nodesPerGenServer" "$gpusPerNodePerGenServer" \ + "${genNodeLists[$i]}" "$pytestCommandGENPrecheck" done for i in $(seq 0 $((numCtxServers - 1))); do - ctx_world_size=$((nodesPerCtxServer * gpusPerNodePerCtxServer)) - export DISAGG_SERVING_TYPE="CTX_PRECHECK_$i" - export pytestCommand="$pytestCommandCTXPrecheck --server-idx $i" - timeout -k 60 "${ctPrecheckTimeout:-900}" \ - srun "${srunArgs[@]}" --mpi=pmix --kill-on-bad-exit=1 \ - -N $nodesPerCtxServer \ - -w "${ctxNodeLists[$i]}" \ - --ntasks=$ctx_world_size \ - --ntasks-per-node=$gpusPerNodePerCtxServer \ - bash $precheckRunScript &> "$precheckDir/logs/ctx_$i.log" & - precheckPids+=($!) - precheckNames+=("ctx_$i") - sleep 5 # Wait for pyxis container namespace initialization to avoid race condition + ct_launch_step ctx "$i" "$nodesPerCtxServer" "$gpusPerNodePerCtxServer" \ + "${ctxNodeLists[$i]}" "$pytestCommandCTXPrecheck" done local precheckFailed=0 k rc diff --git a/jenkins/scripts/perf/local/submit.py b/jenkins/scripts/perf/local/submit.py index 8f99af0b9f2e..f603b0d02ea2 100755 --- a/jenkins/scripts/perf/local/submit.py +++ b/jenkins/scripts/perf/local/submit.py @@ -6,8 +6,20 @@ import re import shutil import subprocess +import sys from datetime import datetime + +def _import_precheck_config(llm_src): + """Import the pure-stdlib precheck config module from the repo tree + (single owner of the gate's enable policy and timeout formulas).""" + path = os.path.join(llm_src, "tests", "scripts", "perf-sanity", "cache_transceiver_precheck") + if path not in sys.path: + sys.path.insert(0, path) + import precheck_config + + return precheck_config + import yaml AGG_CONFIG_FOLDER = os.environ.get("AGG_CONFIG_FOLDER", "tests/scripts/perf-sanity/aggregated") @@ -910,43 +922,22 @@ def main(): ] ) - # Cache-transceiver network precheck (mirrors jenkins/scripts/perf/ - # submit.py): reuses the exact ucx_tls_cmd + worker env strings of the - # real worker steps. On by default; disabled per-yaml - # (cache_transceiver_precheck.enabled: false); TRTLLM_DISAGG_CT_PRECHECK - # =0/1 at generation time overrides the yaml either way. - precheck_cfg = config.get("cache_transceiver_precheck", {}) or {} - precheck_env = os.environ.get("TRTLLM_DISAGG_CT_PRECHECK") - if precheck_env is not None: - precheck_enabled = precheck_env == "1" - else: - precheck_enabled = bool(precheck_cfg.get("enabled", True)) - precheck_mode = "gen_only" if benchmark_mode == "gen_only" else "e2e" - precheck_cmd = ( - "python3 $llmSrcNode/tests/scripts/perf-sanity/cache_transceiver_precheck/" - "run_precheck.py --config $configYamlPath " - "--work-dir $testOutputDir/cache_transceiver_precheck " - f"--benchmark-mode {precheck_mode} --llm-src $llmSrcNode" - ) - # The external step timeout must cover the driver's first-rep NIXL - # wire-up allowance (precheck_config: min(1800, 150 * max world)). - max_world = max( - hardware_config.get("gpus_per_ctx_server", 0) or 0, - hardware_config.get("gpus_per_gen_server", 0) or 0, - ) - default_step_timeout = 900 + min(1800, 150 * max_world) + # Cache-transceiver network precheck (same wiring as jenkins/scripts/ + # perf/submit.py): reuses the exact ucx_tls_cmd + worker env strings + # of the real worker steps; enable policy and timeouts come from + # precheck_config (single owner). + pcfg = _import_precheck_config(llm_src) script_prefix_lines.extend( - [ - f"export ctPrecheckEnabled={1 if precheck_enabled else 0}", - f"export ctPrecheckTimeout=" - f"{int(precheck_cfg.get('step_timeout_s', default_step_timeout))}", - "export precheckRunScript=$llmSrcNode/jenkins/scripts/perf/" - "disaggregated/slurm_precheck_run.sh", - f'export pytestCommandCTXPrecheck="{ucx_tls_cmd} $CTX_WORKER_ENV_VARS' - f' {precheck_cmd} --role ctx"', - f'export pytestCommandGENPrecheck="{ucx_tls_cmd} $GEN_WORKER_ENV_VARS' - f' {precheck_cmd} --role gen"', - ] + pcfg.precheck_prefix_lines( + config, + benchmark_mode, + config_path_expr="$configYamlPath", + ucx_tls_cmd=ucx_tls_cmd, + max_world=max( + hardware_config.get("gpus_per_ctx_server", 0) or 0, + hardware_config.get("gpus_per_gen_server", 0) or 0, + ), + ) ) # Add srun args for disagg diff --git a/jenkins/scripts/perf/submit.py b/jenkins/scripts/perf/submit.py index 81aeca1ed029..83c7afa0da3d 100755 --- a/jenkins/scripts/perf/submit.py +++ b/jenkins/scripts/perf/submit.py @@ -21,9 +21,21 @@ import math import os import re +import sys import yaml + +def _import_precheck_config(llm_src): + """Import the pure-stdlib precheck config module from the repo tree + (single owner of the gate's enable policy and timeout formulas).""" + path = os.path.join(llm_src, "tests", "scripts", "perf-sanity", "cache_transceiver_precheck") + if path not in sys.path: + sys.path.insert(0, path) + import precheck_config + + return precheck_config + AGG_CONFIG_FOLDER = "tests/scripts/perf-sanity/aggregated" DISAGG_CONFIG_FOLDER = "tests/scripts/perf-sanity/disaggregated" @@ -631,43 +643,21 @@ def main(): # servers with the same instance topology, and reuses the exact # $ucx_tls_cmd / $CTX_WORKER_ENV_VARS / $GEN_WORKER_ENV_VARS strings # of the worker steps so the UCX environment matches by construction. - # On by default. Disabled per-yaml (cache_transceiver_precheck. - # enabled: false); TRTLLM_DISAGG_CT_PRECHECK=0/1 at generation time - # overrides the yaml either way (global kill switch). - precheck_cfg = config.get("cache_transceiver_precheck", {}) or {} - precheck_env = os.environ.get("TRTLLM_DISAGG_CT_PRECHECK") - if precheck_env is not None: - precheck_enabled = precheck_env == "1" - else: - precheck_enabled = bool(precheck_cfg.get("enabled", True)) - config_yaml_node = f"$llmSrcNode/{os.path.relpath(config_yaml, args.llm_src)}" - precheck_cmd = ( - "python3 $llmSrcNode/tests/scripts/perf-sanity/cache_transceiver_precheck/" - f"run_precheck.py --config {config_yaml_node} " - "--work-dir $testOutputDir/cache_transceiver_precheck " - f"--benchmark-mode {benchmark_mode} --llm-src $llmSrcNode" - ) - # The external step timeout must cover the driver's first-rep NIXL - # wire-up allowance (precheck_config: min(1800, 150 * max world)). - max_world = max( - hardware_config["gpus_per_ctx_server"], hardware_config["gpus_per_gen_server"] - ) - default_step_timeout = 900 + min(1800, 150 * max_world) + # Enable/kill-switch policy and timeouts live in precheck_config + # (single owner, shared with the local flow). + pcfg = _import_precheck_config(args.llm_src) script_prefix_lines.extend( - [ - f"export ctPrecheckEnabled={1 if precheck_enabled else 0}", - f"export ctPrecheckTimeout=" - f"{int(precheck_cfg.get('step_timeout_s', default_step_timeout))}", - # Suite name for the synthetic junit xml the launch script - # writes when the precheck fails. - f'export stageName="{args.stage_name}"', - "export precheckRunScript=$llmSrcNode/jenkins/scripts/perf/" - "disaggregated/slurm_precheck_run.sh", - f'export pytestCommandCTXPrecheck="{ucx_tls_cmd} $CTX_WORKER_ENV_VARS' - f' {precheck_cmd} --role ctx"', - f'export pytestCommandGENPrecheck="{ucx_tls_cmd} $GEN_WORKER_ENV_VARS' - f' {precheck_cmd} --role gen"', - ] + pcfg.precheck_prefix_lines( + config, + benchmark_mode, + config_path_expr=f"$llmSrcNode/{os.path.relpath(config_yaml, args.llm_src)}", + ucx_tls_cmd=ucx_tls_cmd, + max_world=max( + hardware_config["gpus_per_ctx_server"], + hardware_config["gpus_per_gen_server"], + ), + stage_name=args.stage_name, + ) ) srun_args_lines.extend( [ diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py index d732caf28b07..3b9e26e15117 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py @@ -104,6 +104,59 @@ def _spec_nextn(side): return 0 +def wireup_timeout_s(max_world): + """First-rep NIXL agent wire-up allowance (see PRECHECK_DEFAULTS).""" + return min(1800, 150 * int(max_world)) + + +def default_step_timeout_s(max_world): + """External (srun-level) timeout covering one precheck instance, + including the first-rep wire-up. Imported by the launch tooling + (jenkins/scripts/perf/{,local/}submit.py) so the outer timeout can + never drift below the driver's internal budget.""" + return 900 + wireup_timeout_s(max_world) + + +def precheck_prefix_lines(cfg, benchmark_mode, config_path_expr, ucx_tls_cmd, max_world, + stage_name=""): + """Launch-script export lines wiring the precheck gate. + + Single owner of the enable/kill-switch policy, the step-timeout default, + and the export names the gate consumes — shared by + jenkins/scripts/perf/submit.py and jenkins/scripts/perf/local/submit.py. + `config_path_expr` and the env-var references are launch-script-side + expressions ($llmSrcNode etc.), expanded at sbatch runtime. + """ + knobs = cfg.get("cache_transceiver_precheck", {}) or {} + env = os.environ.get("TRTLLM_DISAGG_CT_PRECHECK") + # On by default; yaml opts out per test; the env var (when set) overrides + # the yaml either way (global kill switch). + enabled = env == "1" if env is not None else bool(knobs.get("enabled", True)) + cmd = ( + "python3 $llmSrcNode/tests/scripts/perf-sanity/cache_transceiver_precheck/" + f"run_precheck.py --config {config_path_expr} " + "--work-dir $testOutputDir/cache_transceiver_precheck " + f"--benchmark-mode {benchmark_mode} --llm-src $llmSrcNode" + ) + lines = [ + f"export ctPrecheckEnabled={int(enabled)}", + # The external srun timeout must cover the driver's first-rep NIXL + # wire-up allowance; the default derives from the same formula the + # driver budgets with (default_step_timeout_s). + f"export ctPrecheckTimeout=" + f"{int(knobs.get('step_timeout_s', default_step_timeout_s(max_world)))}", + "export precheckRunScript=$llmSrcNode/jenkins/scripts/perf/" + "disaggregated/slurm_precheck_run.sh", + f'export pytestCommandCTXPrecheck="{ucx_tls_cmd} $CTX_WORKER_ENV_VARS {cmd} --role ctx"', + f'export pytestCommandGENPrecheck="{ucx_tls_cmd} $GEN_WORKER_ENV_VARS {cmd} --role gen"', + ] + if stage_name: + # Suite name for the synthetic junit xml the gate writes on failure + # (absent -> the gate falls back to $SLURM_JOB_NAME). + lines.append(f'export stageName="{stage_name}"') + return lines + + def resolve_plan(cfg, benchmark_mode="e2e"): """Build the shared precheck plan both roles must agree on. @@ -145,9 +198,6 @@ def resolve_plan(cfg, benchmark_mode="e2e"): f"{ctx_xcvr.get('backend')} vs {gen_xcvr.get('backend')}" ) - ctx_nextn = _spec_nextn(ctx_side) - gen_nextn = _spec_nextn(gen_side) - knobs = dict(PRECHECK_DEFAULTS) knobs.update(cfg.get("cache_transceiver_precheck", {}) or {}) @@ -179,22 +229,6 @@ def resolve_plan(cfg, benchmark_mode="e2e"): "gpus_per_node": gpus_per_node, "ctx": ctx, "gen": gen, - "ctx_cache_transceiver_config": ctx_xcvr, - "gen_cache_transceiver_config": gen_xcvr, - "ctx_num_nextn_predict_layers": ctx_nextn, - "gen_num_nextn_predict_layers": gen_nextn, - "ctx_kv_dtype": str((ctx_side.get("kv_cache_config") or {}).get("dtype", "auto")), - "gen_kv_dtype": str((gen_side.get("kv_cache_config") or {}).get("dtype", "auto")), - # Tri-state, matching KvCacheConfig's pydantic default: explicit - # True/False from the yaml wins; absent means "auto", which the - # driver resolves against the model class's get_model_defaults() at - # runtime — exactly like serving (_resolve_kv_cache_manager_v2_auto). - "ctx_use_kv_cache_manager_v2": (ctx_side.get("kv_cache_config") or {}).get( - "use_kv_cache_manager_v2", "auto" - ), - "gen_use_kv_cache_manager_v2": (gen_side.get("kv_cache_config") or {}).get( - "use_kv_cache_manager_v2", "auto" - ), "tokens_per_block": tokens_per_block, "request_lengths": req_lens, "num_requests": int(knobs["num_requests"]), @@ -205,11 +239,21 @@ def resolve_plan(cfg, benchmark_mode="e2e"): "wireup_timeout_s": int( knobs["wireup_timeout_s"] if knobs["wireup_timeout_s"] is not None - else min(1800, 150 * max(ctx["world_size"], gen["world_size"])) + else wireup_timeout_s(max(ctx["world_size"], gen["world_size"])) ), "rendezvous_timeout_s": int(knobs["rendezvous_timeout_s"]), "verify_data": bool(knobs["verify_data"]), } + for role, side, xcvr in (("ctx", ctx_side, ctx_xcvr), ("gen", gen_side, gen_xcvr)): + kv_cfg = side.get("kv_cache_config") or {} + plan[f"{role}_cache_transceiver_config"] = xcvr + plan[f"{role}_num_nextn_predict_layers"] = _spec_nextn(side) + plan[f"{role}_kv_dtype"] = str(kv_cfg.get("dtype", "auto")) + # Tri-state, matching KvCacheConfig's pydantic default: explicit + # True/False from the yaml wins; absent means "auto", which the + # driver resolves against the model class's get_model_defaults() at + # runtime — exactly like serving (_resolve_kv_cache_manager_v2_auto). + plan[f"{role}_use_kv_cache_manager_v2"] = kv_cfg.get("use_kv_cache_manager_v2", "auto") plan["fingerprint"] = plan_fingerprint(plan) return plan @@ -237,11 +281,9 @@ def plan_fingerprint(plan): def side_plan(plan, role): """Per-role view: this role's parallelism + transceiver/kv config.""" - other = "gen" if role == "ctx" else "ctx" return { "role": role, "parallel": plan[role], - "peer_parallel": plan[other], "cache_transceiver_config": plan[f"{role}_cache_transceiver_config"], "kv_dtype": plan[f"{role}_kv_dtype"], "num_nextn_predict_layers": plan[f"{role}_num_nextn_predict_layers"], diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index 2c82ed7e3c15..7c043cb722ed 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -54,7 +54,6 @@ import argparse import base64 -import dataclasses import hashlib import hmac import json @@ -145,37 +144,39 @@ def params_to_wire(p): def params_from_wire(d): - """Inverse of params_to_wire (ctor per disaggregated_params.py).""" - import tensorrt_llm.bindings.executor as tllme - - return tllme.ContextPhaseParams( - list(d["first_gen_tokens"]), - int(d["req_id"]), - base64.b64decode(d["opaque_state"]), - d["draft_tokens"], - d["ctx_dp_rank"], - d["ctx_info_endpoint"], - ) + """Inverse of params_to_wire, via the maintained DisaggregatedParams + converter (keeps us off the raw nanobind ctor signature).""" + from tensorrt_llm import DisaggregatedParams + + return DisaggregatedParams( + ctx_request_id=int(d["req_id"]), + first_gen_tokens=list(d["first_gen_tokens"]), + opaque_state=base64.b64decode(d["opaque_state"]), + draft_tokens=d["draft_tokens"], + ctx_dp_rank=d["ctx_dp_rank"], + ctx_info_endpoint=d["ctx_info_endpoint"], + ).get_context_phase_params() # --------------------------------------------------------------------------- # # KV fill / verify (heavy imports stay inside functions: --dry-run and unit # tests must work without torch / tensorrt_llm installed) # --------------------------------------------------------------------------- # -def _pattern_like(view, seed): - """Deterministic tensor matching `view`, constant along the head axis. +def _pattern_like(shape, dtype, device, seed): + """Deterministic tensor of `shape`, constant along the head axis. - `view` is an HND block slice: [nblocks, kv_factor, heads, tokens, dim]. + `shape` is an HND block slice: [nblocks, kv_factor, heads, tokens, dim]. Generated on CPU (bit-identical across nodes), head axis generated as 1 - and expanded, so ctx/gen sides with different local head counts (TP + and expanded ON DEVICE (transferring heads x the unique data would be + pure waste), so ctx/gen sides with different local head counts (TP resharding) or different local layer sets (PP re-splitting) still agree. """ import torch - nb, kv, heads, tok, dim = view.shape + nb, kv, heads, tok, dim = shape g = torch.Generator(device="cpu").manual_seed(int(seed)) rnd = torch.rand((nb, kv, 1, tok, dim), dtype=torch.float32, generator=g) - return rnd.to(view.dtype).expand(nb, kv, heads, tok, dim) + return rnd.to(dtype).to(device).expand(nb, kv, heads, tok, dim) def _request_block_views(kvm, rid): @@ -191,8 +192,8 @@ def _request_block_views(kvm, rid): def fill_request(kvm, rid): for global_layer, buf, valid in _request_block_views(kvm, rid): - view = buf[valid] - buf[valid] = _pattern_like(view, seed_for(rid, global_layer)).to(view.device) + shape = (len(valid), *buf.shape[1:]) + buf[valid] = _pattern_like(shape, buf.dtype, buf.device, seed_for(rid, global_layer)) def verify_request(kvm, rid): @@ -201,49 +202,14 @@ def verify_request(kvm, rid): for global_layer, buf, valid in _request_block_views(kvm, rid): recv = buf[valid] - exp = _pattern_like(recv, seed_for(rid, global_layer)).to(recv.device) - if not torch.equal(recv.float(), exp.float()): - bad = (recv.float() != exp.float()).sum().item() + exp = _pattern_like(recv.shape, recv.dtype, recv.device, seed_for(rid, global_layer)) + recv_f, exp_f = recv.float(), exp.float() # fp8 lacks direct compare ops + if not torch.equal(recv_f, exp_f): + bad = (recv_f != exp_f).sum().item() return False, f"layer={global_layer} mismatched_elements={bad}/{recv.numel()}" return True, "" -@dataclasses.dataclass -class _KvCacheConfigV2: - """KvCacheConfig stand-in for KVCacheManagerV2. - - Mirrors examples/disaggregated/slurm/cache_transceiver_test (and the - reference dataclass in - tests/unittest/disaggregated/test_cache_transceiver_single_process.py): - KVCacheManagerV2 reads these fields off the config object directly, so - every attribute it accesses must exist here. - """ - - max_tokens: "int | None" = None - enable_block_reuse: bool = False - max_attention_window: "list | None" = None - sink_token_length: "int | None" = None - free_gpu_memory_fraction: "float | None" = None - host_cache_size: "int | None" = None - disk_cache_size: "int | None" = None - disk_cache_path: "str | None" = None - onboard_blocks: bool = True - cross_kv_cache_fraction: "float | None" = None - secondary_offload_min_priority: "int | None" = None - event_buffer_max_size: int = 0 - kv_cache_event_hash_algo: str = "auto" - max_gpu_total_bytes: "int | None" = None - enable_partial_reuse: bool = False - copy_on_partial_reuse: bool = False - dtype: str = "auto" - pool_ratio: "list | None" = None - avg_seq_len: "int | None" = None - block_reuse_policy: str = "all_reusable" - enable_swa_scratch_reuse: bool = False - disk_prefetch_num_reqs: int = 4 - max_util_for_resume: float = 0.95 - - def _lookup_model_cls(model_dir): """Model class from config.json architectures, like serving's automodel path.""" try: @@ -287,8 +253,19 @@ def resolve_model_prefs(model_dir, side, cache_cfg): except Exception as e: # noqa: BLE001 - model hooks may need llm_args print(f"[precheck] WARNING: get_model_defaults failed ({e!r}); " f"assuming V1", flush=True) - kv_defaults = defaults.get("kv_cache_config") or {} - use_v2 = kv_defaults.get("use_kv_cache_manager_v2", False) is True + try: + # The REAL serving resolver, via the same shim pattern as the + # runtime resolution below -- one owner for the 'auto' semantics. + from tensorrt_llm.llmapi.llm_utils import _resolve_kv_cache_manager_v2_auto + + shim = types.SimpleNamespace( + kv_cache_config=types.SimpleNamespace(use_kv_cache_manager_v2="auto") + ) + use_v2 = bool(_resolve_kv_cache_manager_v2_auto(shim, defaults)) + except Exception as e: # noqa: BLE001 - fall back like a missing model + print(f"[precheck] WARNING: V2 'auto' resolution failed ({e!r}); " + f"assuming V1", flush=True) + use_v2 = False else: use_v2 = bool(setting) @@ -353,13 +330,18 @@ def build_kv_cache_manager(kv_shape, plan, side, mapping, max_req_len, use_v2): ) if use_v2: from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 + from tensorrt_llm.llmapi.llm_args import KvCacheConfig - # is_disagg=True doubles the IndexMapper capacity so in-flight - # transfers (TRANS_IN_PROGRESS) can hold slots, like real serving. + # The REAL pydantic KvCacheConfig (what serving passes): partial reuse + # is pinned off because its pydantic default is True and block reuse + # is off here. is_disagg=True doubles the IndexMapper capacity so + # in-flight transfers (TRANS_IN_PROGRESS) can hold slots. return KVCacheManagerV2( - _KvCacheConfigV2( + KvCacheConfig( max_tokens=max_tokens, enable_block_reuse=False, + enable_partial_reuse=False, + copy_on_partial_reuse=False, max_attention_window=[padded_len], ), cache_type, @@ -430,11 +412,9 @@ def add_sequence(kvm, req, prompt_len, use_v2): def free_sequence(kvm, req, use_v2): if use_v2: - import torch - # free_resources() closes the kv_cache AND releases the IndexMapper - # slot (closing the cache alone leaks slots). - torch.cuda.current_stream().synchronize() + # slot (closing the cache alone leaks slots). Callers synchronize the + # stream once per batch before freeing (see _free_all). kvm.free_resources(req) return kvm.impl.remove_sequence(req.py_request_id, req, True) @@ -461,8 +441,6 @@ def _wait_gen_complete(xcvr, req, runtime, llm_request_state): next_report = t0 + 15.0 while req.state not in terminal: xcvr.check_gen_transfer_status(1) - if req.state in terminal: - break now = time.monotonic() if now >= next_report: print( @@ -494,14 +472,10 @@ def run_token(): def write_addr(path, payload): - """Atomically publish an addr file (clearing any stale one first). It - carries the session HMAC key, so restrict it to the owning user before - it becomes visible.""" + """Atomically publish an addr file (os.replace overwrites stale ones; + readers reject wrong-job stamps). It carries the session HMAC key, so + restrict it to the owning user before it becomes visible.""" os.makedirs(os.path.dirname(path), exist_ok=True) - try: - os.remove(path) # stale file from a previous run in a reused work dir - except FileNotFoundError: - pass payload = dict(payload, job=run_token()) tmp = f"{path}.tmp.{os.getpid()}" with open(tmp, "w") as f: @@ -619,7 +593,9 @@ def parse_bandwidth_gbps(csv_dir, rank, tag="recv"): with open(path) as f: rows = list(csv_mod.DictReader(f)) col = next((c for c in (rows[0] or {}) if "Bandwidth" in c), None) - vals = [float(r[col]) / 8.0 for r in rows if col and r.get(col)] + if col is None: + return None + vals = [float(r[col]) / 8.0 for r in rows if r.get(col)] return statistics.median(vals) if vals else None except (OSError, ValueError, IndexError, StopIteration): return None @@ -634,7 +610,6 @@ class PrecheckRunner: def __init__(self, args, plan, side, comm): from mpi4py import MPI # noqa: F401 - ensures MPI initialized - self.args = args self.plan = plan self.side = side self.role = side["role"] @@ -759,13 +734,15 @@ def ctx_run_chunk(self, peer_idx, li, req_len, rep, chunk): # Params for pair k come from its owning dp rank at pp stage 0 with # attention DP; without DP every rank sends the same request, and the # instance leader's params are the ones the real server would return. - contrib = {} - if local_err is None: - if self.side["parallel"]["enable_attention_dp"]: - if self.mapping.pp_rank == 0: - contrib = {p: r.context_phase_params for p, r in reqs.items()} - elif self.is_leader: - contrib = {p: r.context_phase_params for p, r in reqs.items()} + if self.side["parallel"]["enable_attention_dp"]: + contributes = self.mapping.pp_rank == 0 + else: + contributes = self.is_leader + contrib = ( + {p: r.context_phase_params for p, r in reqs.items()} + if local_err is None and contributes + else {} + ) gathered = self.comm.gather(contrib, root=0) params_by_pair = {} if self.is_leader and reason is None: @@ -807,7 +784,11 @@ def ctx_finish_chunk(self, reqs): raise _TransferError(f"ctx transfer failed: {reason}") def gen_run_chunk(self, peer_idx, li, req_len, rep, chunk, params_by_pair): - """Receive + verify owned pairs. Returns (ok, mismatch_detail).""" + """Receive + verify owned pairs. Returns (ok, mismatch_detail). + + Warmup reps skip the (CPU-heavy) byte verification: their result is + discarded by the caller either way, transfer errors still raise. + """ import torch import tensorrt_llm @@ -849,7 +830,7 @@ def gen_run_chunk(self, peer_idx, li, req_len, rep, chunk, params_by_pair): ] if bad: local_err = _TransferError(f"gen DISAGG_TRANS_ERROR on pairs {bad}") - elif self.plan["verify_data"]: + elif self.plan["verify_data"] and rep >= self.plan["warmup_requests"]: for pair, req in reqs.items(): ok, detail = verify_request(self.kvm, req.py_request_id) if not ok: @@ -866,6 +847,10 @@ def gen_run_chunk(self, peer_idx, li, req_len, rep, chunk, params_by_pair): return (not mismatches, "; ".join(mismatches[:4])) def _free_all(self, reqs): + if reqs and self.use_v2: + import torch + + torch.cuda.current_stream().synchronize() # V2 frees need quiesced stream for req in reqs.values(): try: free_sequence(self.kvm, req, self.use_v2) @@ -1039,10 +1024,9 @@ def gen_run_peer(runner, peer_idx, arm, disarm): # ZMQ timeout is only a backstop under the per-chunk alarm, so it # includes the first-rep wire-up allowance unconditionally. if runner.is_leader: - import zmq as zmq_mod - + zmq, _ = runner._zmq() sock.setsockopt( - zmq_mod.RCVTIMEO, + zmq.RCVTIMEO, (plan["chunk_timeout_s"] + plan["wireup_timeout_s"] + 30) * 1000, ) @@ -1122,7 +1106,7 @@ def main(argv=None): plan, side = load_plan(args) if args.dry_run: - print(json.dumps({k: v for k, v in plan.items()}, indent=2, default=str)) + print(json.dumps(plan, indent=2, default=str)) return 0 if plan.get("skip"): print(f"[precheck] SKIP: {plan['skip_reason']}", flush=True) @@ -1240,6 +1224,20 @@ def disarm(): flush=True, ) + def record_peer_failure(peer, exc): + """Map a per-peer exception to its verdict (isolation: other peers + keep running).""" + disarm() + if isinstance(exc, _Timeout): + runner.recorder.record( + peer, 0, "TIMEOUT", + f"exceeded the budget during {current_cell['what']}", + ) + elif isinstance(exc, _PeerAbort): + runner.recorder.record(peer, 0, "TRANSFER_ERROR", str(exc)) + else: + runner.recorder.record(peer, 0, "TRANSFER_ERROR", repr(exc)) + # --- sessions ------------------------------------------------------------- num_peers = side["num_peers"] try: @@ -1268,44 +1266,23 @@ def disarm(): try: ctx_serve_peer( runner, - socks.get(gj) if runner.is_leader else None, + socks.get(gj), gj, arm, disarm, keys.get(gj), ) runner.recorder.record(f"gen_{gj}", 0, "PASS", "served all transfers") - except _Timeout: - disarm() - runner.recorder.record( - f"gen_{gj}", - 0, - "TIMEOUT", - f"exceeded {plan['chunk_timeout_s']}s during {current_cell['what']}", - ) - except _PeerAbort as e: - disarm() - runner.recorder.record(f"gen_{gj}", 0, "TRANSFER_ERROR", str(e)) except Exception as e: # noqa: BLE001 - per-peer isolation - disarm() - runner.recorder.record(f"gen_{gj}", 0, "TRANSFER_ERROR", repr(e)) + record_peer_failure(f"gen_{gj}", e) else: open_sessions = [] for ci in range(num_peers): try: sock, sess_key = gen_run_peer(runner, ci, arm, disarm) open_sessions.append((ci, sock, sess_key)) - except _Timeout: - disarm() - runner.recorder.record( - f"ctx_{ci}", - 0, - "TIMEOUT", - f"exceeded {plan['chunk_timeout_s']}s during {current_cell['what']}", - ) except Exception as e: # noqa: BLE001 - per-peer isolation - disarm() - runner.recorder.record(f"ctx_{ci}", 0, "TRANSFER_ERROR", repr(e)) + record_peer_failure(f"ctx_{ci}", e) # Deferred "done": released only after EVERY ctx peer's schedule # finished, so all ctx instances stay alive for the whole # precheck (real-serving lifecycle -- no dead-agent connections @@ -1313,14 +1290,8 @@ def disarm(): for ci, sock, sess_key in open_sessions: try: gen_release_peer(runner, ci, sock, sess_key, arm, disarm) - except _Timeout: - disarm() - runner.recorder.record( - f"ctx_{ci}", 0, "TIMEOUT", "releasing session (deferred done)" - ) except Exception as e: # noqa: BLE001 - best-effort release - disarm() - runner.recorder.record(f"ctx_{ci}", 0, "TRANSFER_ERROR", repr(e)) + record_peer_failure(f"ctx_{ci}", e) finally: disarm() hang_detector.stop() diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py index 66f788a1f92d..6b25e0508d28 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_config.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py @@ -36,6 +36,7 @@ sys.path.insert(0, os.path.abspath(_PRECHECK_DIR)) import precheck_config as pcfg # noqa: E402 +import run_precheck as rp # noqa: E402 (stdlib-only at import time) def _disagg_yaml(ctx_extra=None, gen_extra=None, **overrides): @@ -212,7 +213,6 @@ def test_side_plan_views(): ctx_view = pcfg.side_plan(plan, "ctx") gen_view = pcfg.side_plan(plan, "gen") assert ctx_view["parallel"]["world_size"] == 4 - assert ctx_view["peer_parallel"]["world_size"] == 16 assert ctx_view["num_peers"] == 1 and gen_view["num_peers"] == 1 assert gen_view["cache_transceiver_config"]["max_tokens_in_buffer"] == 16384 @@ -220,19 +220,12 @@ def test_side_plan_views(): class TestControlWireFormat: """run_precheck's HMAC-JSON control frames (importable without torch).""" - def _rp(self): - import run_precheck as rp - - return rp - def test_roundtrip(self): - rp = self._rp() key = b"\x01" * 32 msg = ["go", {"li": 0, "rep": 1, "chunk": 2}] assert rp.unpack_msg(rp.pack_msg(msg, key), key) == msg def test_tampered_frame_rejected(self): - rp = self._rp() key = b"\x01" * 32 raw = rp.pack_msg(["hello", {}], key) bad = raw[:-1] + bytes([raw[-1] ^ 0xFF]) @@ -240,18 +233,15 @@ def test_tampered_frame_rejected(self): rp.unpack_msg(bad, key) def test_wrong_key_rejected(self): - rp = self._rp() raw = rp.pack_msg(["hello", {}], b"\x01" * 32) with pytest.raises(rp._TransferError): rp.unpack_msg(raw, b"\x02" * 32) def test_short_frame_rejected(self): - rp = self._rp() with pytest.raises(rp._TransferError): rp.unpack_msg(b"tiny", b"\x01" * 32) def test_addr_file_owner_only(self, tmp_path): - rp = self._rp() path = str(tmp_path / "rendezvous" / "ctx0_gen0.addr") rp.write_addr(path, {"host": "h", "port": 1, "key": "aa"}) assert (os.stat(path).st_mode & 0o777) == 0o600 @@ -299,13 +289,7 @@ def test_model_kv_shape_vocab_size(tmp_path): class TestRendezvousStaleness: """wait_for_addr must skip addr files stamped by a previous run.""" - def _rp(self): - import run_precheck as rp - - return rp - def test_same_job_accepted(self, tmp_path, monkeypatch): - rp = self._rp() monkeypatch.setenv("SLURM_JOB_ID", "12345") p = str(tmp_path / "rendezvous" / "ctx0_gen0.addr") rp.write_addr(p, {"host": "h", "port": 1, "key": "aa"}) @@ -313,7 +297,6 @@ def test_same_job_accepted(self, tmp_path, monkeypatch): assert got["job"] == "12345" and got["port"] == 1 def test_stale_job_skipped_until_timeout(self, tmp_path, monkeypatch): - rp = self._rp() monkeypatch.setenv("SLURM_JOB_ID", "11111") p = str(tmp_path / "rendezvous" / "ctx0_gen0.addr") rp.write_addr(p, {"host": "h", "port": 1, "key": "aa"}) # stamped 11111 @@ -322,7 +305,6 @@ def test_stale_job_skipped_until_timeout(self, tmp_path, monkeypatch): rp.wait_for_addr(p, timeout_s=2) def test_no_job_id_accepts_any(self, tmp_path, monkeypatch): - rp = self._rp() monkeypatch.setenv("SLURM_JOB_ID", "11111") p = str(tmp_path / "rendezvous" / "ctx0_gen0.addr") rp.write_addr(p, {"host": "h", "port": 1, "key": "aa"}) @@ -330,7 +312,6 @@ def test_no_job_id_accepts_any(self, tmp_path, monkeypatch): assert rp.wait_for_addr(p, timeout_s=2)["port"] == 1 def test_write_addr_replaces_stale_file(self, tmp_path, monkeypatch): - rp = self._rp() monkeypatch.setenv("SLURM_JOB_ID", "11111") p = str(tmp_path / "rendezvous" / "ctx0_gen0.addr") rp.write_addr(p, {"host": "old", "port": 1, "key": "aa"}) @@ -354,8 +335,6 @@ def test_wireup_timeout_derivation(): def test_rid_tags_dense_within_session(): """The C++ notification tag is rid & 0xFFF: rids must be dense within a (ctx, gen) session so tags cannot alias across reps/lengths.""" - import run_precheck as rp - plan = pcfg.resolve_plan(_disagg_yaml()) # n_pairs=16 total_reps = plan["warmup_requests"] + plan["num_requests"] n_pairs = plan["n_pairs"] From 5c5f778a2bb75a8ac8431eb36015f3ad37ed58bb Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Mon, 20 Jul 2026 17:41:32 -0700 Subject: [PATCH 11/21] [None][test] Readability pass on the precheck driver - main() is now a short narrative (parse -> env -> topology check -> watchdog -> setup gate -> sessions -> verdict): the watchdog wiring moved to _install_watchdog(), the exception->verdict mapping to _make_peer_failure_recorder(), and the two session drivers to _serve_gen_peers() / _drive_ctx_peers() with role-level docstrings. - The leader-to-leader session protocol now has an ASCII sequence diagram at the section header (hello/welcome, go/params, deferred done/bye). - 'chunk' is defined everywhere a reader first meets it (module vocabulary block, pcfg.chunks(), the max_concurrent_pairs knob): a chunk is the batch of TRANSFER PAIRS in flight at once, NOT token/data chunking -- the term collides with chunked prefill and confused a reviewer. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> [None][test] CPU-level coverage for multi-peer precheck orchestration Runs the REAL 2-ctx x 1-gen session protocol in-process (real PrecheckRunner, ZMQ loopback sockets, HMAC frames, StatusRecorder, rendezvous files; only the GPU transfer methods stubbed): full-pass asserts every ctx serves the complete schedule and receives the deferred done, and the failure-isolation case asserts one aborting ctx yields TRANSFER_ERROR on the gen while the healthy peer still passes and is released. This is the multi-instance logic the hardware 'B' topology exercises, now guarded without GPUs. Addr files publish loopback (SLURMD_NODENAME=127.0.0.1) since CI/sandbox hosts may not resolve their own hostname. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../precheck_config.py | 13 +- .../run_precheck.py | 254 +++++++++++------- .../test_cache_transceiver_precheck_config.py | 140 ++++++++++ 3 files changed, 302 insertions(+), 105 deletions(-) diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py index 3b9e26e15117..9885007a296d 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py @@ -40,7 +40,10 @@ # Measured requests per (peer, request_length) after warmup. "num_requests": 2, "warmup_requests": 1, - # Upper bound on concurrently in-flight transfer pairs per chunk. + # Upper bound on concurrently in-flight transfer pairs. The n_pairs + # dp-rank pairings are processed in batches ("chunks") of at most this + # many -- see chunks() for the term's exact meaning (it is NOT token + # chunking). "max_concurrent_pairs": 8, # signal.alarm / hang-detector budget for one chunk of transfers. "chunk_timeout_s": 180, @@ -319,7 +322,13 @@ def max_owned_per_chunk(plan, role): def chunks(plan): - """Pair indices grouped into concurrency-bounded chunks.""" + """Pair indices grouped into concurrency-bounded chunks. + + A "chunk" is the batch of TRANSFER PAIRS in flight at once (at most + max_concurrent_pairs of the n_pairs dp-rank pairings) -- NOT token/data + chunking. It bounds the synthetic KV pool size and gives the per-chunk + alarm a precise target while still exercising concurrent transfers. + """ pairs = list(range(plan["n_pairs"])) size = plan["chunk_size"] return [pairs[i : i + size] for i in range(0, len(pairs), size)] diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index 7c043cb722ed..75c17934e76d 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -25,6 +25,12 @@ before any model is loaded, so network/UCX misconfiguration is caught in minutes instead of after a full model bring-up. +Vocabulary: a PAIR is one dp-rank-to-dp-rank transfer pairing (n_pairs = +max of the two sides' dp sizes, so every dp rank is exercised); a REP is one +repetition of all pairs (1 warmup + num_requests measured); a CHUNK is the +batch of pairs in flight at once (at most max_concurrent_pairs) -- NOT +token/data chunking. + Asymmetric parallelism (e.g. ctx dep4 -> gen dep16, ctx pp8 -> gen tp32) is supported: the fill pattern is seeded per (request, GLOBAL layer) and is constant along the KV-head axis, so any TP resharding or PP re-splitting on @@ -912,6 +918,25 @@ def chunk_timeout_s(plan, li, rep): # --------------------------------------------------------------------------- # # ctx / gen session loops +# +# Per (ctx, gen) pair, the leaders speak a lockstep REQ/REP protocol (all +# frames HMAC-JSON; KV bytes themselves go through the transceiver under +# test, never over ZMQ): +# +# gen leader ctx leader +# | -- hello {fingerprint} --------------> | yaml mismatch -> abort +# | <---------------- welcome ------------ | +# | -- go {li, rep, chunk} --------------> | every rank posts its sends +# | <----- params {pair: ctx_phase} ------ | (from the owning dp ranks) +# | ...KV transfer + byte verification.. | +# | (repeat per schedule entry) | +# | -- done (deferred: after ALL peers) -> | ctx exits only now +# | <------------------ bye -------------- | +# +# A "chunk" here is a batch of TRANSFER PAIRS in flight at once (at most +# max_concurrent_pairs of the n_pairs dp-rank pairings) -- NOT token/data +# chunking. It bounds the tiny synthetic KV pool and gives the per-chunk +# alarm a precise target, while still exercising concurrent transfers. # --------------------------------------------------------------------------- # def ctx_serve_peer(runner, sock, peer_idx, arm, disarm, key): """Serve one gen peer's full schedule on a dedicated REP socket.""" @@ -1101,6 +1126,127 @@ def load_plan(args): return plan, role_side +def _install_watchdog(runner, plan, rank): + """Two-layer stall protection around every phase of the run. + + signal.alarm catches Python-level stalls; HangDetector catches + GIL-released native hangs (dumps stacks, records TIMEOUT, SIGKILLs so + `srun --kill-on-bad-exit` tears the step down). The external `timeout` + around the srun is the guaranteed backstop for GIL-held hangs. + + Returns (arm, disarm, stop, current_cell): per-phase alarm control, the + final shutdown, and the mutable "what phase are we in" marker used by + failure messages. + """ + from tensorrt_llm._torch.pyexecutor.hang_detector import HangDetector + + signal.signal(signal.SIGALRM, _alarm_handler) + current_cell = {"what": "startup"} + + def _on_hang(): + runner.recorder.record("-", 0, "TIMEOUT", f"hang detected during {current_cell['what']}") + runner.recorder.finalize() + sys.stderr.write( + f"[precheck {runner.role}_{runner.server_idx} r{rank}] WATCHDOG_KILL " + f"{current_cell['what']}\n" + ) + sys.stderr.flush() + os.kill(os.getpid(), signal.SIGKILL) + + # The detector must outlast the LONGEST legitimate wait (peer handshakes + # are serialized across sessions); per-cell alarms are the tighter bound + # for actual transfer work. + hang_detector = HangDetector( + timeout=hello_timeout_s(plan, runner.side["num_peers"]) + plan["chunk_timeout_s"] + 60, + on_detected=_on_hang, + ) + hang_detector.start() + + def arm(what, seconds=None): + current_cell["what"] = what + signal.alarm(seconds or plan["chunk_timeout_s"]) + hang_detector.checkpoint() + + def disarm(): + signal.alarm(0) + hang_detector.cancel_task() + + def stop(): + disarm() + hang_detector.stop() + + return arm, disarm, stop, current_cell + + +def _make_peer_failure_recorder(runner, disarm, current_cell): + """Exception -> verdict mapping shared by all per-peer loops (isolation: + one bad peer never stops the others).""" + + def record_peer_failure(peer, exc): + disarm() + if isinstance(exc, _Timeout): + runner.recorder.record( + peer, 0, "TIMEOUT", f"exceeded the budget during {current_cell['what']}" + ) + elif isinstance(exc, _PeerAbort): + runner.recorder.record(peer, 0, "TRANSFER_ERROR", str(exc)) + else: + runner.recorder.record(peer, 0, "TRANSFER_ERROR", repr(exc)) + + return record_peer_failure + + +def _serve_gen_peers(runner, plan, arm, disarm, record_peer_failure): + """ctx role: bind one dedicated REP socket per gen peer (avoids REQ + interleaving across sessions on a shared socket), publish the addr files, + then serve each peer's full schedule. Each session gets a fresh HMAC key, + shared only through the work-dir addr file (0600).""" + num_peers = runner.side["num_peers"] + socks, keys = {}, {} + if runner.is_leader: + zmq, zctx = runner._zmq() + host = os.environ.get("SLURMD_NODENAME") or socket.gethostname() + for gj in range(num_peers): + s = zctx.socket(zmq.REP) + s.setsockopt(zmq.LINGER, 0) + # Generous: gen peers are serialized across ctx servers. + s.setsockopt(zmq.RCVTIMEO, hello_timeout_s(plan, num_peers) * 1000) + port = s.bind_to_random_port("tcp://*") + keys[gj] = secrets.token_bytes(32) + write_addr( + addr_path(runner.work_dir, runner.server_idx, gj), + {"host": host, "port": port, "key": keys[gj].hex()}, + ) + socks[gj] = s + for gj in range(num_peers): + try: + ctx_serve_peer(runner, socks.get(gj), gj, arm, disarm, keys.get(gj)) + runner.recorder.record(f"gen_{gj}", 0, "PASS", "served all transfers") + except Exception as e: # noqa: BLE001 - per-peer isolation + record_peer_failure(f"gen_{gj}", e) + + +def _drive_ctx_peers(runner, arm, disarm, record_peer_failure): + """gen role: run every ctx peer's schedule, then release all sessions. + + The release ("done") is deferred until EVERY peer's schedule finished, so + all ctx instances stay alive for the whole precheck -- matching real + serving, where no transceiver ever holds connections to a dead agent + while transfers are still running.""" + open_sessions = [] + for ci in range(runner.side["num_peers"]): + try: + sock, sess_key = gen_run_peer(runner, ci, arm, disarm) + open_sessions.append((ci, sock, sess_key)) + except Exception as e: # noqa: BLE001 - per-peer isolation + record_peer_failure(f"ctx_{ci}", e) + for ci, sock, sess_key in open_sessions: + try: + gen_release_peer(runner, ci, sock, sess_key, arm, disarm) + except Exception as e: # noqa: BLE001 - best-effort release + record_peer_failure(f"ctx_{ci}", e) + + def main(argv=None): args = parse_args(argv) plan, side = load_plan(args) @@ -1158,42 +1304,7 @@ def main(argv=None): flush=True, ) - # --- watchdog: signal.alarm for Python-level stalls + HangDetector for - # GIL-released native hangs (dumps stacks, records TIMEOUT, SIGKILLs so - # `srun --kill-on-bad-exit` tears the step down; the external `timeout` - # around the srun is the guaranteed backstop for GIL-held hangs). - signal.signal(signal.SIGALRM, _alarm_handler) - from tensorrt_llm._torch.pyexecutor.hang_detector import HangDetector - - current_cell = {"what": "startup"} - - def _on_hang(): - runner.recorder.record("-", 0, "TIMEOUT", f"hang detected during {current_cell['what']}") - runner.recorder.finalize() - sys.stderr.write( - f"[precheck {args.role}_{args.server_idx} r{rank}] WATCHDOG_KILL " - f"{current_cell['what']}\n" - ) - sys.stderr.flush() - os.kill(os.getpid(), signal.SIGKILL) - - # The detector must outlast the LONGEST legitimate wait (peer handshakes - # are serialized across sessions); per-cell alarms below are the tighter - # bound for actual transfer work. - hang_detector = HangDetector( - timeout=hello_timeout_s(plan, side["num_peers"]) + plan["chunk_timeout_s"] + 60, - on_detected=_on_hang, - ) - hang_detector.start() - - def arm(what, seconds=None): - current_cell["what"] = what - signal.alarm(seconds or plan["chunk_timeout_s"]) - hang_detector.checkpoint() - - def disarm(): - signal.alarm(0) - hang_detector.cancel_task() + arm, disarm, stop_watchdog, current_cell = _install_watchdog(runner, plan, rank) # --- setup: KV pool + transceiver (same config as the real test) --------- setup_err = None @@ -1224,77 +1335,14 @@ def disarm(): flush=True, ) - def record_peer_failure(peer, exc): - """Map a per-peer exception to its verdict (isolation: other peers - keep running).""" - disarm() - if isinstance(exc, _Timeout): - runner.recorder.record( - peer, 0, "TIMEOUT", - f"exceeded the budget during {current_cell['what']}", - ) - elif isinstance(exc, _PeerAbort): - runner.recorder.record(peer, 0, "TRANSFER_ERROR", str(exc)) - else: - runner.recorder.record(peer, 0, "TRANSFER_ERROR", repr(exc)) - - # --- sessions ------------------------------------------------------------- - num_peers = side["num_peers"] + record_peer_failure = _make_peer_failure_recorder(runner, disarm, current_cell) try: if args.role == "ctx": - # One dedicated REP socket per gen peer (avoids REQ interleaving - # across sessions on a shared socket), published via addr files. - # Each session gets a fresh HMAC key, shared only through the - # work-dir addr file (0600). - socks, keys = {}, {} - if runner.is_leader: - zmq, zctx = runner._zmq() - host = os.environ.get("SLURMD_NODENAME") or socket.gethostname() - for gj in range(num_peers): - s = zctx.socket(zmq.REP) - s.setsockopt(zmq.LINGER, 0) - # Generous: gen peers are serialized across ctx servers. - s.setsockopt(zmq.RCVTIMEO, hello_timeout_s(plan, num_peers) * 1000) - port = s.bind_to_random_port("tcp://*") - keys[gj] = secrets.token_bytes(32) - write_addr( - addr_path(runner.work_dir, args.server_idx, gj), - {"host": host, "port": port, "key": keys[gj].hex()}, - ) - socks[gj] = s - for gj in range(num_peers): - try: - ctx_serve_peer( - runner, - socks.get(gj), - gj, - arm, - disarm, - keys.get(gj), - ) - runner.recorder.record(f"gen_{gj}", 0, "PASS", "served all transfers") - except Exception as e: # noqa: BLE001 - per-peer isolation - record_peer_failure(f"gen_{gj}", e) + _serve_gen_peers(runner, plan, arm, disarm, record_peer_failure) else: - open_sessions = [] - for ci in range(num_peers): - try: - sock, sess_key = gen_run_peer(runner, ci, arm, disarm) - open_sessions.append((ci, sock, sess_key)) - except Exception as e: # noqa: BLE001 - per-peer isolation - record_peer_failure(f"ctx_{ci}", e) - # Deferred "done": released only after EVERY ctx peer's schedule - # finished, so all ctx instances stay alive for the whole - # precheck (real-serving lifecycle -- no dead-agent connections - # in anyone's transceiver while transfers are still running). - for ci, sock, sess_key in open_sessions: - try: - gen_release_peer(runner, ci, sock, sess_key, arm, disarm) - except Exception as e: # noqa: BLE001 - best-effort release - record_peer_failure(f"ctx_{ci}", e) + _drive_ctx_peers(runner, arm, disarm, record_peer_failure) finally: - disarm() - hang_detector.stop() + stop_watchdog() # --- teardown + result ------------------------------------------------------ bw = None diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py index 6b25e0508d28..616dd7213341 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_config.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py @@ -354,3 +354,143 @@ def session_rids(ctx_idx, gen_idx): assert not (set(a) & set(b)) # globally unique across sessions tags = [r & 0xFFF for r in a] assert len(set(tags)) == len(tags) # no tag aliasing within a session + + +class TestMultiPeerOrchestration: + """CPU-only end-to-end run of the 2-ctx x 1-gen session protocol. + + Exercises the exact multi-instance logic of the hardware "B" topology: + real ZMQ sockets + HMAC frames + StatusRecorder + rendezvous files via + the real PrecheckRunner/_serve_gen_peers/_drive_ctx_peers, with only the + GPU transfer methods stubbed out. + """ + + class _FakeComm: + def Get_rank(self): + return 0 + + def Get_size(self): + return 1 + + def bcast(self, obj, root=0): + return obj + + def gather(self, obj, root=0): + return [obj] + + def allgather(self, obj): + return [obj] + + class _FakeParams: + first_gen_tokens = [0] + req_id = 1 + opaque_state = b"op" + draft_tokens = None + ctx_dp_rank = 0 + disagg_info_endpoint = None + + def _mk_runner(self, role, server_idx, plan, work_dir, monkeypatch, fail_ctx=False): + import sys + import types + + # PrecheckRunner.__init__ imports mpi4py only to ensure MPI init. + monkeypatch.setitem(sys.modules, "mpi4py", types.SimpleNamespace(MPI=None)) + # The gen side converts wire params through tensorrt_llm bindings; + # identity is fine here (params_to_wire is covered separately). + monkeypatch.setattr(rp, "params_from_wire", lambda d: d) + + args = types.SimpleNamespace(server_idx=server_idx, work_dir=work_dir) + side = pcfg.side_plan(plan, role) + runner = rp.PrecheckRunner(args, plan, side, self._FakeComm()) + + calls = {"chunks": 0} + + def ctx_run_chunk(peer_idx, li, req_len, rep, chunk): + if fail_ctx: + raise rp._TransferError("injected ctx failure") + calls["chunks"] += 1 + return {p: self._FakeParams() for p in chunk}, {} + + runner.ctx_run_chunk = ctx_run_chunk + runner.ctx_finish_chunk = lambda reqs: None + runner.gen_run_chunk = ( + lambda peer_idx, li, req_len, rep, chunk, params: (True, "") + ) + runner._calls = calls + return runner + + def _run(self, tmp_path, monkeypatch, fail_ctx1=False): + import threading + + monkeypatch.setenv("SLURM_JOB_ID", "777") + # Publish loopback in the addr files: the real node hostname may not + # resolve in sandboxed/CI environments, and everything is one process. + monkeypatch.setenv("SLURMD_NODENAME", "127.0.0.1") + cfg = _disagg_yaml( + hardware={"gpus_per_node": 4, "num_ctx_servers": 2, "num_gen_servers": 1}, + cache_transceiver_precheck={ + "request_lengths": [32], + "num_requests": 1, + "warmup_requests": 1, + "rendezvous_timeout_s": 30, + "chunk_timeout_s": 30, + "wireup_timeout_s": 0, + }, + ) + plan = pcfg.resolve_plan(cfg) + work = str(tmp_path) + noop = lambda *a, **k: None # noqa: E731 - signal.alarm needs main thread + + gen = self._mk_runner("gen", 0, plan, work, monkeypatch) + ctxs = [ + self._mk_runner("ctx", i, plan, work, monkeypatch, fail_ctx=(fail_ctx1 and i == 1)) + for i in range(2) + ] + + failures = [] + + def rec(peer, exc): + failures.append((peer, type(exc).__name__)) + + threads = [ + threading.Thread( + target=rp._serve_gen_peers, args=(c, plan, noop, noop, rec), daemon=True + ) + for c in ctxs + ] + for t in threads: + t.start() + rp._drive_ctx_peers( + gen, noop, noop, rp._make_peer_failure_recorder(gen, noop, {"what": "test"}) + ) + for t in threads: + t.join(timeout=60) + assert not t.is_alive(), "ctx serve thread wedged" + return plan, gen, ctxs, failures + + def test_two_ctx_full_pass(self, tmp_path, monkeypatch): + plan, gen, ctxs, failures = self._run(tmp_path, monkeypatch) + assert not failures + # gen recorded a PASS per (peer, req_len) + assert {(c["peer"], c["status"]) for c in gen.recorder.cases} == { + ("ctx_0", "PASS"), + ("ctx_1", "PASS"), + } + # every ctx served the full schedule (reps x chunks) and got its + # deferred done (PASS recorded only after done/bye completes) + total_chunks = len(pcfg.chunks(plan)) * ( + plan["warmup_requests"] + plan["num_requests"] + ) + for c in ctxs: + assert c._calls["chunks"] == total_chunks + assert [x["status"] for x in c.recorder.cases] == ["PASS"] + + def test_ctx_failure_isolated(self, tmp_path, monkeypatch): + plan, gen, ctxs, failures = self._run(tmp_path, monkeypatch, fail_ctx1=True) + # gen side: healthy peer unaffected, failing peer gets a clear verdict + by_peer = {c["peer"]: c["status"] for c in gen.recorder.cases} + assert by_peer == {"ctx_0": "PASS", "ctx_1": "TRANSFER_ERROR"} + # ctx_1's own serve loop surfaced the failure (its peer is gen_0) + assert ("gen_0", "_TransferError") in failures + # ctx_0 served its full schedule and got the deferred done + assert [c["status"] for c in ctxs[0].recorder.cases] == ["PASS"] From 760d8e1f5d40058c9e79ddfea4ffcd62a0c19cd8 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Mon, 20 Jul 2026 22:23:04 -0700 Subject: [PATCH 12/21] [None][test] Stop capturing NIXL_* env in precheck verdicts The env snapshot in status/*.json exists to capture variables that influence transport behavior (UCX_/TRTLLM_/TLLM_). The only NIXL_* variable ever observed is NIXL_VERSION, a leftover marker from the NGC base image's bundled NIXL rather than the library TRT-LLM links from /opt/nvidia/nvda_nixl (it read '1.0.0' in containers shipping NIXL v1.3.0), so it misstated the transport version in every verdict. Drop the NIXL_ prefix from the snapshot entirely. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../cache_transceiver_precheck/run_precheck.py | 6 +++++- .../others/test_cache_transceiver_precheck_config.py | 12 ++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index 75c17934e76d..a55ee02add65 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -533,7 +533,11 @@ def __init__(self, work_dir, role, server_idx, is_leader): os.makedirs(status_dir, exist_ok=True) self.json_path = os.path.join(status_dir, f"{role}_{server_idx}.json") self.text_path = os.path.join(status_dir, f"{role}_{server_idx}.status") - env_keys = ("UCX_", "NIXL_", "TRTLLM_", "TLLM_") + # NIXL_* is deliberately absent: the only such variable seen in + # practice is NIXL_VERSION, a stale marker from the NGC base image's + # bundled NIXL (not the library TRT-LLM links from + # /opt/nvidia/nvda_nixl), which misstates the transport version. + env_keys = ("UCX_", "TRTLLM_", "TLLM_") self.env = {k: v for k, v in sorted(os.environ.items()) if k.startswith(env_keys)} def record(self, peer, req_len, status, reason=""): diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py index 616dd7213341..df338b434749 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_config.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py @@ -494,3 +494,15 @@ def test_ctx_failure_isolated(self, tmp_path, monkeypatch): assert ("gen_0", "_TransferError") in failures # ctx_0 served its full schedule and got the deferred done assert [c["status"] for c in ctxs[0].recorder.cases] == ["PASS"] + + +def test_status_env_snapshot_excludes_nixl(tmp_path, monkeypatch): + """NIXL_* is not captured: the only such variable seen in practice is + NIXL_VERSION, a stale NGC-base-image marker that misstates the version of + the actually-linked library.""" + monkeypatch.setenv("NIXL_VERSION", "1.0.0") + monkeypatch.setenv("NIXL_PLUGIN_DIR", "/opt/x") + monkeypatch.setenv("UCX_TLS", "rc,cuda_copy") + rec = rp.StatusRecorder(str(tmp_path), "gen", 0, is_leader=True) + assert not any(k.startswith("NIXL_") for k in rec.env) + assert rec.env["UCX_TLS"] == "rc,cuda_copy" # behavioral vars still captured From de67ddd4738563ee96db874862e4923e2c1a0029 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Mon, 20 Jul 2026 23:58:26 -0700 Subject: [PATCH 13/21] [None][test] Run the network check on unmodelable models via a simplified pool; add Python-transceiver bandwidth Two follow-ups so the precheck covers more real configs: 1. Sparse-attention models (DeepSeek V4 / DSA: an indexer INDEX_KEY pool separate from the main K/V pool) cannot be expressed as the precheck's single (layers, kv_heads, head_dim) pool. Since the precheck is a NETWORK/transceiver availability check -- not a KV-correctness test -- it must still RUN, not skip. model_kv_shape now returns a simple MLA-flavored stand-in pool for these models (real layer count, one latent head, is_mla=True), so the transceiver is set up with SELFKONLY + AttentionType.MLA exactly like the real sparse model and the transfer exercises the same UCX/NIXL path -- only the per-token bytes are generic. The KV manager version (V1/V2) and transceiver runtime still resolve against the model + yaml (V4 -> V2 + Python), so the check stays consistent with what the model actually runs. A SIMPLIFIED note is printed and recorded. 2. The Python transceiver records bandwidth through perf_logger.py (KVSendTask throughput_mbs on the ctx side) rather than the C++ TRTLLM_KVCACHE_TIME_OUTPUT_PATH CSV, so V2xV2 verdicts showed no bandwidth. setup() now also enables TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO + TLLM_KV_TRANSFER_PERF_LOG_FILE, and the result step medians throughput from the ctx-side perf CSVs on the Python path (gen-side recv CSVs on the C++ path, as before). CPU unit tests cover the simplified MLA stand-in and the Python CSV bandwidth parser. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../precheck_config.py | 40 +++++++++++- .../run_precheck.py | 47 +++++++++++++- .../test_cache_transceiver_precheck_config.py | 63 +++++++++++++++++++ 3 files changed, 147 insertions(+), 3 deletions(-) diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py index 9885007a296d..a32965e9fdf7 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py @@ -374,13 +374,33 @@ def resolve_model_dir(cfg, llm_src=None, llm_models_root=None): return None +def unmodelable_kv_reason(hf_cfg): + """Why the precheck cannot faithfully model this model's KV layout, or None. + + The precheck allocates a SINGLE KV pool from (layers, kv_heads, head_dim). + Sparse-attention models (DeepSeek V4 / DSA: an indexer with its own + INDEX_KEY pool at a different stride than the main K/V pool) cannot be + represented that way. The precheck is a NETWORK/transceiver availability + check, not a KV-correctness test, so rather than skip these (the check + must still run) it falls back to a simple generic pool -- the transfer + exercises the exact same UCX/NIXL path, topology, and transceiver config, + just with generic bytes. + """ + if any(k in hf_cfg for k in ("index_topk", "index_n_heads", "index_head_dim")): + return "sparse-attention model (DSA/indexer): using a generic KV pool -- this checks the network/transceiver path, not V4's exact KV layout" + return None + + def model_kv_shape(model_dir): """KV cache shape (per-token layout) from the model's config.json. Handles MLA checkpoints (kv_lora_rank present: one latent 'head' of kv_lora_rank + qk_rope_head_dim) and GQA/MHA. `num_layers` excludes the MTP nextn layers -- those are added by the KV cache manager via - spec_config, mirroring real serving. + spec_config, mirroring real serving. A model whose KV layout the precheck + can't model (see unmodelable_kv_reason) returns the generic + FALLBACK_KV_SHAPE plus a ``simplified`` note, so the network check still + runs with a stand-in pool instead of skipping. """ if not model_dir: return dict(FALLBACK_KV_SHAPE) @@ -393,6 +413,24 @@ def model_kv_shape(model_dir): hf_cfg = hf_cfg["text_config"] num_layers = hf_cfg.get("num_hidden_layers") + + reason = unmodelable_kv_reason(hf_cfg) + if reason: + # Simple MLA-flavored stand-in pool: real layer count, one latent + # head, is_mla=True so the transceiver is set up with SELFKONLY + + # AttentionType.MLA exactly like the real sparse model -- only the + # per-token bytes are generic. Enough to check the network path. + latent = int(hf_cfg.get("kv_lora_rank") or 0) + int(hf_cfg.get("qk_rope_head_dim") or 0) + return { + "num_layers": int(num_layers) if num_layers else FALLBACK_KV_SHAPE["num_layers"], + "num_kv_heads": 1, + "head_dim": latent or int(hf_cfg.get("head_dim") or 0) or 576, + "is_mla": True, + "vocab_size": int(hf_cfg.get("vocab_size") or 0) or None, + "source": "config.json (simplified MLA stand-in)", + "simplified": reason, + } + if num_layers is None: return dict(FALLBACK_KV_SHAPE) diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index a55ee02add65..eee55323c3ba 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -611,6 +611,30 @@ def parse_bandwidth_gbps(csv_dir, rank, tag="recv"): return None +def parse_python_bandwidth_gbps(csv_dir): + """Median KV-send throughput in GB/s from the Python transceiver's perf + CSVs (perf_logger.py). Enabled via TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO=1 + + TLLM_KV_TRANSFER_PERF_LOG_FILE=, which writes + __.csv. throughput_mbs (MB/s) is on the SENDER + (ctx) side, task_type=KVSendTask; receiver rows have no throughput. + Best-effort -- returns None when no perf CSV exists. + """ + import csv as csv_mod + import glob + import statistics + + vals = [] + for path in glob.glob(os.path.join(csv_dir, "perf_*_*.csv")): + try: + with open(path) as f: + for r in csv_mod.DictReader(f): + if r.get("task_type") == "KVSendTask" and r.get("throughput_mbs"): + vals.append(float(r["throughput_mbs"]) / 1024.0) # MB/s -> GB/s + except (OSError, ValueError): + continue + return statistics.median(vals) if vals else None + + # --------------------------------------------------------------------------- # # Chunk execution # --------------------------------------------------------------------------- # @@ -672,14 +696,26 @@ def setup(self, kv_shape, max_req_len): enable_attention_dp=par["enable_attention_dp"], ) os.makedirs(self.csv_dir, exist_ok=True) + # C++ transceiver bandwidth CSV (per-rank recv). os.environ["TRTLLM_KVCACHE_TIME_OUTPUT_PATH"] = self.csv_dir + # Python transceiver bandwidth CSV (perf_logger.py): writes + # {base}_{instance}_{rank}.csv with KVSendTask throughput on the ctx + # side. Harmless to set unconditionally -- the C++ path ignores it. + os.environ["TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO"] = "1" + os.environ["TLLM_KV_TRANSFER_PERF_LOG_FILE"] = os.path.join(self.csv_dir, "perf") # Built VERBATIM from the disagg yaml's cache_transceiver_config so # backend/max_tokens_in_buffer/timeouts match the real test exactly. cache_cfg = CacheTransceiverConfig(**self.side["cache_transceiver_config"]) # Yaml-absent settings resolve against the model's preferences, like - # serving does (kv manager version + transceiver runtime). + # serving does (kv manager version + transceiver runtime) -- this holds + # even for the simplified stand-in pool: only the KV SHAPE is generic; + # the V1/V2 manager version and the transceiver runtime must still + # match what the real model runs (e.g. V4 -> V2 + Python). self.use_v2 = resolve_model_prefs(self.plan.get("_model_dir"), self.side, cache_cfg) + if kv_shape.get("simplified") and self.is_leader: + print(f"[precheck {self.role}_{self.server_idx}] SIMPLIFIED: " + f"{kv_shape['simplified']}", flush=True) # KVCacheManagerV2 only works with the Python transceiver (see # cache_transceiver_test/report.py); reject the pairing up front with # a clear INIT_ERROR instead of a C++ binding type error. @@ -1349,8 +1385,15 @@ def main(argv=None): stop_watchdog() # --- teardown + result ------------------------------------------------------ + # Bandwidth lives on different sides per transceiver: C++ records it on the + # receiver (gen recv CSVs), the Python transceiver on the sender (ctx perf + # CSVs). csv_dir is shared across the instance's ranks, so for the Python + # path the leader alone medians over all ranks' perf files. bw = None - if args.role == "gen": + if runner.runtime == "PYTHON": + if args.role == "ctx" and runner.is_leader: + bw = parse_python_bandwidth_gbps(runner.csv_dir) + elif args.role == "gen": local_bw = parse_bandwidth_gbps(runner.csv_dir, rank) bws = [b for b in comm.gather(local_bw, root=0) or [] if b] if runner.is_leader and bws: diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py index df338b434749..e2c55d7a8216 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_config.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py @@ -506,3 +506,66 @@ def test_status_env_snapshot_excludes_nixl(tmp_path, monkeypatch): rec = rp.StatusRecorder(str(tmp_path), "gen", 0, is_leader=True) assert not any(k.startswith("NIXL_") for k in rec.env) assert rec.env["UCX_TLS"] == "rc,cuda_copy" # behavioral vars still captured + + +def test_sparse_attention_model_uses_simplified_mla_pool(tmp_path): + """DeepSeek V4 / DSA can't be modeled as a single KV pool. The precheck is + a NETWORK check, so it falls back to a simple MLA-flavored stand-in pool + (real layer count, one latent head, is_mla=True) and still runs -- never + skips.""" + d = tmp_path / "v4" + d.mkdir() + (d / "config.json").write_text( + json.dumps( + { + "architectures": ["DeepseekV4ForCausalLM"], + "num_hidden_layers": 43, + "num_attention_heads": 64, + "num_key_value_heads": 1, + "head_dim": 512, + "index_head_dim": 128, + "index_n_heads": 64, + "index_topk": 512, + "sliding_window": 128, + "vocab_size": 129280, + } + ) + ) + shape = pcfg.model_kv_shape(str(d)) + assert shape.get("simplified") and "sparse" in shape["simplified"] + assert shape["is_mla"] is True and shape["num_kv_heads"] == 1 + assert shape["num_layers"] == 43 # real layer count preserved + assert shape["head_dim"] == 512 and shape["vocab_size"] == 129280 + # a plain MLA model is modeled normally (no simplified marker) + m = tmp_path / "mla" + m.mkdir() + (m / "config.json").write_text( + json.dumps({"num_hidden_layers": 4, "kv_lora_rank": 512, "qk_rope_head_dim": 64}) + ) + assert not pcfg.model_kv_shape(str(m)).get("simplified") + + +def test_python_transceiver_bandwidth_csv(tmp_path): + """Bandwidth from the Python transceiver's perf_logger CSVs: median over + KVSendTask throughput_mbs (MB/s -> GB/s), receiver rows ignored.""" + header = ( + "timestamp,task_type,unique_rid,peer_rank,transfer_size_bytes," + "avg_segment_size_bytes,transfer_entry_count,prepare_args_latency_ms," + "queue_latency_ms,transfer_latency_ms,task_latency_ms,throughput_mbs" + ) + # two ctx ranks, each its own perf file; KVSendTask rows carry throughput + (tmp_path / "perf_abc_0.csv").write_text( + header + "\n" + "t,KVSendTask,1,0,1000,,,0,0,0,0,102400.00\n" # 100 GB/s + "t,AuxSendTask,1,0,10,,,0,0,0,0,1.00\n" # tiny metadata, ignored + ) + (tmp_path / "perf_abc_1.csv").write_text( + header + "\n" + "t,KVSendTask,2,1,1000,,,0,0,0,0,204800.00\n" # 200 GB/s + ) + # a receiver file (no throughput) must not contribute + (tmp_path / "perf_def_8.csv").write_text(header + "\nt,KVRecvTask,3,0,,,,,,,5.0,\n") + bw = rp.parse_python_bandwidth_gbps(str(tmp_path)) + assert bw == 150.0 # median(100, 200) + # no perf files -> None + assert rp.parse_python_bandwidth_gbps(str(tmp_path / "empty")) is None From 0c3b7d93a55c5c5dccc0dad26cfb461ce7893680 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Tue, 21 Jul 2026 00:22:23 -0700 Subject: [PATCH 14/21] [None][test] Rename the precheck's 'chunk' concept to 'wave' 'chunk' collided with chunked prefill (splitting a request's tokens), which is an unrelated concept -- reviewers kept reading the precheck's concurrency batch as token chunking. Rename the concept throughout to 'wave' (a wave of transfer pairs in flight at once): waves()/wave_size/ wave_timeout_s (yaml knob), ctx_run_wave/gen_run_wave/ctx_finish_wave, and all comments/docstrings. The max_concurrent_pairs knob keeps its name (already unambiguous). The docstring now states the collision explicitly so the term doesn't drift back. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../cache_transceiver_precheck/README.md | 4 +- .../precheck_config.py | 38 ++++----- .../run_precheck.py | 79 ++++++++++--------- .../test_cache_transceiver_precheck_config.py | 42 +++++----- 4 files changed, 82 insertions(+), 81 deletions(-) diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md index 4ebfdd10656b..715c66b18cfd 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md @@ -36,7 +36,7 @@ layers) is read from the real model's `config.json` under `$LLM_MODELS_ROOT`. # optional overrides (defaults in precheck_config.PRECHECK_DEFAULTS): # request_lengths: [1024, 8192] # num_requests: 2 - # chunk_timeout_s: 180 + # wave_timeout_s: 180 # wireup_timeout_s: 1800 # first-rep NIXL agent wire-up allowance # step_timeout_s: 2700 # external srun timeout (default derives from topology) ``` @@ -52,7 +52,7 @@ The first rep of the schedule (the warmup rep) additionally budgets path pays a one-time serialized `fetchRemoteMD` metadata exchange per (receiver rank, ctx rank) agent pair, and cold cross-rack fetches were measured at 100-170s each — real serving absorbs this as slow first requests, -so the precheck does too. Later reps run under the tight `chunk_timeout_s`, +so the precheck does too. Later reps run under the tight `wave_timeout_s`, which is what actually catches hangs. Set `PRECHECK_DEBUG=1` in the worker env to raise the C++/Python transceiver log levels when debugging a stall. diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py index a32965e9fdf7..5ef6455d09f4 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py @@ -41,12 +41,12 @@ "num_requests": 2, "warmup_requests": 1, # Upper bound on concurrently in-flight transfer pairs. The n_pairs - # dp-rank pairings are processed in batches ("chunks") of at most this - # many -- see chunks() for the term's exact meaning (it is NOT token + # dp-rank pairings are processed in batches ("waves") of at most this + # many -- see waves() for the term's exact meaning (it is NOT token # chunking). "max_concurrent_pairs": 8, - # signal.alarm / hang-detector budget for one chunk of transfers. - "chunk_timeout_s": 180, + # signal.alarm / hang-detector budget for one wave of transfers. + "wave_timeout_s": 180, # Extra budget for the FIRST rep of the schedule, which pays the one-time # NIXL agent metadata wire-up (fetchRemoteMD over the mgmt network): with # MLA + ctx pp>1 every receiving rank must connect to every ctx rank, and @@ -223,7 +223,7 @@ def resolve_plan(cfg, benchmark_mode="e2e"): # cover every dp rank on both sides; without DP a single request already # involves every rank of the instance (KV is sharded across TP/PP). n_pairs = max(ctx["dp_size"], gen["dp_size"], 1) - chunk_size = max(1, min(n_pairs, int(knobs["max_concurrent_pairs"]))) + wave_size = max(1, min(n_pairs, int(knobs["max_concurrent_pairs"]))) plan = { "skip": False, @@ -237,8 +237,8 @@ def resolve_plan(cfg, benchmark_mode="e2e"): "num_requests": int(knobs["num_requests"]), "warmup_requests": int(knobs["warmup_requests"]), "n_pairs": n_pairs, - "chunk_size": chunk_size, - "chunk_timeout_s": int(knobs["chunk_timeout_s"]), + "wave_size": wave_size, + "wave_timeout_s": int(knobs["wave_timeout_s"]), "wireup_timeout_s": int( knobs["wireup_timeout_s"] if knobs["wireup_timeout_s"] is not None @@ -277,7 +277,7 @@ def plan_fingerprint(plan): "num_requests", "warmup_requests", "n_pairs", - "chunk_size", + "wave_size", ) return json.dumps({k: plan[k] for k in keys}, sort_keys=True) @@ -308,29 +308,29 @@ def pair_participates(plan, role, tp_rank, pair_idx): return tp_rank == pair_idx % side["dp_size"] -def owned_pairs(plan, role, tp_rank, chunk_pairs): - return [k for k in chunk_pairs if pair_participates(plan, role, tp_rank, k)] +def owned_pairs(plan, role, tp_rank, wave_pairs): + return [k for k in wave_pairs if pair_participates(plan, role, tp_rank, k)] -def max_owned_per_chunk(plan, role): +def max_owned_per_wave(plan, role): """Max concurrently owned pairs per rank (KV pool sizing).""" side = plan[role] - chunk = plan["chunk_size"] + wave = plan["wave_size"] if not side["enable_attention_dp"]: - return chunk - return (chunk + side["dp_size"] - 1) // side["dp_size"] + return wave + return (wave + side["dp_size"] - 1) // side["dp_size"] -def chunks(plan): - """Pair indices grouped into concurrency-bounded chunks. +def waves(plan): + """Pair indices grouped into concurrency-bounded waves. - A "chunk" is the batch of TRANSFER PAIRS in flight at once (at most + A "wave" is the batch of TRANSFER PAIRS in flight at once (at most max_concurrent_pairs of the n_pairs dp-rank pairings) -- NOT token/data - chunking. It bounds the synthetic KV pool size and gives the per-chunk + chunking. It bounds the synthetic KV pool size and gives the per-wave alarm a precise target while still exercising concurrent transfers. """ pairs = list(range(plan["n_pairs"])) - size = plan["chunk_size"] + size = plan["wave_size"] return [pairs[i : i + size] for i in range(0, len(pairs), size)] diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index eee55323c3ba..0ae93d9e0e11 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -27,9 +27,10 @@ Vocabulary: a PAIR is one dp-rank-to-dp-rank transfer pairing (n_pairs = max of the two sides' dp sizes, so every dp rank is exercised); a REP is one -repetition of all pairs (1 warmup + num_requests measured); a CHUNK is the -batch of pairs in flight at once (at most max_concurrent_pairs) -- NOT -token/data chunking. +repetition of all pairs (1 warmup + num_requests measured); a WAVE is the +batch of pairs in flight at once (at most max_concurrent_pairs). "Wave" is +used instead of "chunk" on purpose -- it would otherwise collide with +chunked prefill (splitting a request's tokens), which is unrelated. Asymmetric parallelism (e.g. ctx dep4 -> gen dep16, ctx pp8 -> gen tp32) is supported: the fill pattern is seeded per (request, GLOBAL layer) and is @@ -317,7 +318,7 @@ def build_kv_cache_manager(kv_shape, plan, side, mapping, max_req_len, use_v2): tpb = plan["tokens_per_block"] padded_len = ((max_req_len + tpb - 1) // tpb) * tpb - owned = pcfg.max_owned_per_chunk(plan, side["role"]) + owned = pcfg.max_owned_per_wave(plan, side["role"]) max_tokens = owned * padded_len + 2 * tpb # concurrent pairs + headroom # Real MLA serving uses SELFKONLY (kv_factor=1: one latent plane, no V) — @@ -636,7 +637,7 @@ def parse_python_bandwidth_gbps(csv_dir): # --------------------------------------------------------------------------- # -# Chunk execution +# Wave execution # --------------------------------------------------------------------------- # class PrecheckRunner: """One MPI world = one ctx or gen server instance.""" @@ -741,7 +742,7 @@ def setup(self, kv_shape, max_req_len): # preference on this path -> C++), so read the effective runtime back. self.runtime = cache_cfg.transceiver_runtime or "CPP" - # ---- per-chunk transfer logic ------------------------------------------- + # ---- per-wave transfer logic ------------------------------------------- def _pair_rid(self, peer_idx, li, rep, pair): ctx_idx = self.server_idx if self.is_ctx else peer_idx gen_idx = peer_idx if self.is_ctx else self.server_idx @@ -749,10 +750,10 @@ def _pair_rid(self, peer_idx, li, rep, pair): seq = (li * total_reps + rep) * self.plan["n_pairs"] + pair return make_rid(ctx_idx, gen_idx, self.plan["num_ctx_servers"], seq) - def _owned(self, chunk): - return pcfg.owned_pairs(self.plan, self.role, self.mapping.tp_rank, chunk) + def _owned(self, wave): + return pcfg.owned_pairs(self.plan, self.role, self.mapping.tp_rank, wave) - def ctx_run_chunk(self, peer_idx, li, req_len, rep, chunk): + def ctx_run_wave(self, peer_idx, li, req_len, rep, wave): """Fill + send owned pairs. Returns {pair: context_phase_params} on the leader (params from each @@ -760,7 +761,7 @@ def ctx_run_chunk(self, peer_idx, li, req_len, rep, chunk): """ import tensorrt_llm - owned = self._owned(chunk) + owned = self._owned(wave) reqs, local_err = {}, None try: for pair in owned: @@ -794,7 +795,7 @@ def ctx_run_chunk(self, peer_idx, li, req_len, rep, chunk): if self.is_leader and reason is None: for d in gathered: params_by_pair.update(d or {}) - missing = [p for p in chunk if p not in params_by_pair] + missing = [p for p in wave if p not in params_by_pair] if missing: reason = f"missing context_phase_params for pairs {missing}" @@ -803,8 +804,8 @@ def ctx_run_chunk(self, peer_idx, li, req_len, rep, chunk): raise _TransferError(f"ctx send setup failed: {reason}") return params_by_pair, reqs - def ctx_finish_chunk(self, reqs): - """Wait for all in-flight sends of this chunk, then free.""" + def ctx_finish_wave(self, reqs): + """Wait for all in-flight sends of this wave, then free.""" import tensorrt_llm t0 = time.monotonic() @@ -822,14 +823,14 @@ def ctx_finish_chunk(self, reqs): self._free_all(reqs) states = {p: str(r.state) for p, r in reqs.items()} tensorrt_llm.logger.info( - f"[ctx{self.server_idx} r{self.rank}] chunk sends finished in " + f"[ctx{self.server_idx} r{self.rank}] wave sends finished in " f"{time.monotonic() - t0:.1f}s states={states}" ) reason = self._consensus_error(local_err) if reason is not None: raise _TransferError(f"ctx transfer failed: {reason}") - def gen_run_chunk(self, peer_idx, li, req_len, rep, chunk, params_by_pair): + def gen_run_wave(self, peer_idx, li, req_len, rep, wave, params_by_pair): """Receive + verify owned pairs. Returns (ok, mismatch_detail). Warmup reps skip the (CPU-heavy) byte verification: their result is @@ -839,7 +840,7 @@ def gen_run_chunk(self, peer_idx, li, req_len, rep, chunk, params_by_pair): import tensorrt_llm - owned = self._owned(chunk) + owned = self._owned(wave) reqs, local_err = {}, None try: for pair in owned: @@ -867,7 +868,7 @@ def gen_run_chunk(self, peer_idx, li, req_len, rep, chunk, params_by_pair): _wait_gen_complete(self.xcvr, req, self.runtime, self.llm_request_state) if reqs: tensorrt_llm.logger.info( - f"[gen{self.server_idx} r{self.rank}] chunk recvs finished in " + f"[gen{self.server_idx} r{self.rank}] wave recvs finished in " f"{time.monotonic() - t0:.1f}s" ) torch.cuda.synchronize() # receive may land on a side stream @@ -928,13 +929,13 @@ def _leader_send_recv(self, sock, obj, key): def _schedule(plan): - """Deterministic (li, req_len, rep, chunk) schedule both sides iterate.""" + """Deterministic (li, req_len, rep, wave) schedule both sides iterate.""" out = [] total_reps = plan["warmup_requests"] + plan["num_requests"] for li, req_len in enumerate(plan["request_lengths"]): for rep in range(total_reps): - for chunk in pcfg.chunks(plan): - out.append((li, req_len, rep, chunk)) + for wave in pcfg.waves(plan): + out.append((li, req_len, rep, wave)) return out @@ -949,11 +950,11 @@ def hello_timeout_s(plan, num_peers): return plan["rendezvous_timeout_s"] + num_peers * (300 + plan["wireup_timeout_s"]) -def chunk_timeout_s(plan, li, rep): - """Per-chunk budget: the first rep additionally pays the one-time NIXL +def wave_timeout_s(plan, li, rep): + """Per-wave budget: the first rep additionally pays the one-time NIXL agent wire-up (see PRECHECK_DEFAULTS['wireup_timeout_s']).""" extra = plan["wireup_timeout_s"] if (li == 0 and rep == 0) else 0 - return plan["chunk_timeout_s"] + extra + return plan["wave_timeout_s"] + extra # --------------------------------------------------------------------------- # @@ -966,16 +967,16 @@ def chunk_timeout_s(plan, li, rep): # gen leader ctx leader # | -- hello {fingerprint} --------------> | yaml mismatch -> abort # | <---------------- welcome ------------ | -# | -- go {li, rep, chunk} --------------> | every rank posts its sends +# | -- go {li, rep, wave} --------------> | every rank posts its sends # | <----- params {pair: ctx_phase} ------ | (from the owning dp ranks) # | ...KV transfer + byte verification.. | # | (repeat per schedule entry) | # | -- done (deferred: after ALL peers) -> | ctx exits only now # | <------------------ bye -------------- | # -# A "chunk" here is a batch of TRANSFER PAIRS in flight at once (at most +# A "wave" here is a batch of TRANSFER PAIRS in flight at once (at most # max_concurrent_pairs of the n_pairs dp-rank pairings) -- NOT token/data -# chunking. It bounds the tiny synthetic KV pool and gives the per-chunk +# chunking. It bounds the tiny synthetic KV pool and gives the per-wave # alarm a precise target, while still exercising concurrent transfers. # --------------------------------------------------------------------------- # def ctx_serve_peer(runner, sock, peer_idx, arm, disarm, key): @@ -1006,8 +1007,8 @@ def leader_reply(obj): raise _TransferError(f"handshake with gen_{peer_idx} failed: {msg[:1]}") leader_reply(("welcome", {"fingerprint": plan["fingerprint"]})) - for li, req_len, rep, chunk in _schedule(plan): - arm(f"gen_{peer_idx} len={req_len} rep={rep}", seconds=chunk_timeout_s(plan, li, rep)) + for li, req_len, rep, wave in _schedule(plan): + arm(f"gen_{peer_idx} len={req_len} rep={rep}", seconds=wave_timeout_s(plan, li, rep)) msg = leader_recv() if msg[0] == "abort": raise _PeerAbort(f"gen_{peer_idx} aborted: {msg[1]}") @@ -1016,13 +1017,13 @@ def leader_reply(obj): f"schedule desync with gen_{peer_idx}: expected li={li} rep={rep}, got {msg}" ) try: - params_by_pair, reqs = runner.ctx_run_chunk(peer_idx, li, req_len, rep, chunk) + params_by_pair, reqs = runner.ctx_run_wave(peer_idx, li, req_len, rep, wave) except _TransferError as e: leader_reply(("abort", str(e))) raise # JSON object keys are strings; the gen side converts back to int. leader_reply(("params", {str(p): params_to_wire(v) for p, v in params_by_pair.items()})) - runner.ctx_finish_chunk(reqs) + runner.ctx_finish_wave(reqs) # The gen defers "done" until it has finished the schedules of ALL its # ctx peers, so every ctx instance stays alive for the whole precheck -- @@ -1085,26 +1086,26 @@ def gen_run_peer(runner, peer_idx, arm, disarm): raise _TransferError(f"ctx_{peer_idx} aborted handshake: {reply[1]}") if reply[0] != "welcome": raise _TransferError(f"unexpected handshake reply from ctx_{peer_idx}: {reply[:1]}") - # Established sessions are dedicated: chunk replies are prompt. The - # ZMQ timeout is only a backstop under the per-chunk alarm, so it + # Established sessions are dedicated: wave replies are prompt. The + # ZMQ timeout is only a backstop under the per-wave alarm, so it # includes the first-rep wire-up allowance unconditionally. if runner.is_leader: zmq, _ = runner._zmq() sock.setsockopt( zmq.RCVTIMEO, - (plan["chunk_timeout_s"] + plan["wireup_timeout_s"] + 30) * 1000, + (plan["wave_timeout_s"] + plan["wireup_timeout_s"] + 30) * 1000, ) case_ok = {} - for li, req_len, rep, chunk in _schedule(plan): - arm(f"ctx_{peer_idx} len={req_len} rep={rep}", seconds=chunk_timeout_s(plan, li, rep)) + for li, req_len, rep, wave in _schedule(plan): + arm(f"ctx_{peer_idx} len={req_len} rep={rep}", seconds=wave_timeout_s(plan, li, rep)) reply = runner._leader_send_recv( - sock, ("go", {"li": li, "rep": rep, "chunk": chunk[0]}), key + sock, ("go", {"li": li, "rep": rep, "wave": wave[0]}), key ) if reply[0] == "abort": raise _TransferError(f"ctx_{peer_idx} aborted: {reply[1]}") params_by_pair = {int(p): params_from_wire(v) for p, v in reply[1].items()} - ok, detail = runner.gen_run_chunk(peer_idx, li, req_len, rep, chunk, params_by_pair) + ok, detail = runner.gen_run_wave(peer_idx, li, req_len, rep, wave, params_by_pair) if rep >= plan["warmup_requests"]: prev_ok, prev_detail = case_ok.get(req_len, (True, "")) case_ok[req_len] = (prev_ok and ok, prev_detail or detail) @@ -1197,14 +1198,14 @@ def _on_hang(): # are serialized across sessions); per-cell alarms are the tighter bound # for actual transfer work. hang_detector = HangDetector( - timeout=hello_timeout_s(plan, runner.side["num_peers"]) + plan["chunk_timeout_s"] + 60, + timeout=hello_timeout_s(plan, runner.side["num_peers"]) + plan["wave_timeout_s"] + 60, on_detected=_on_hang, ) hang_detector.start() def arm(what, seconds=None): current_cell["what"] = what - signal.alarm(seconds or plan["chunk_timeout_s"]) + signal.alarm(seconds or plan["wave_timeout_s"]) hang_detector.checkpoint() def disarm(): diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py index e2c55d7a8216..296b8f1948a5 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_config.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py @@ -85,7 +85,7 @@ def test_resolve_plan_adp_asymmetric(): assert plan["gen"]["world_size"] == 16 and plan["gen"]["dp_size"] == 16 # Cover every gen dp rank. assert plan["n_pairs"] == 16 - assert plan["chunk_size"] == 8 + assert plan["wave_size"] == 8 assert plan["request_lengths"] == [1024, 8192] assert plan["ctx_num_nextn_predict_layers"] == 1 assert plan["ctx_cache_transceiver_config"]["backend"] == "NIXL" @@ -128,7 +128,7 @@ def test_backend_mismatch_raises(): pcfg.resolve_plan(cfg) -def test_pair_participation_and_chunks(): +def test_pair_participation_and_waves(): plan = pcfg.resolve_plan(_disagg_yaml()) # ADP ctx (dp4): pair k belongs to tp_rank k % 4. assert pcfg.pair_participates(plan, "ctx", 1, 5) @@ -136,12 +136,12 @@ def test_pair_participation_and_chunks(): # ADP gen (dp16): 1:1. assert pcfg.pair_participates(plan, "gen", 5, 5) assert not pcfg.pair_participates(plan, "gen", 4, 5) - assert pcfg.chunks(plan) == [list(range(8)), list(range(8, 16))] - # ctx rank owns 2 pairs per chunk of 8; gen rank owns at most 1. - assert pcfg.max_owned_per_chunk(plan, "ctx") == 2 - assert pcfg.max_owned_per_chunk(plan, "gen") == 1 + assert pcfg.waves(plan) == [list(range(8)), list(range(8, 16))] + # ctx rank owns 2 pairs per wave of 8; gen rank owns at most 1. + assert pcfg.max_owned_per_wave(plan, "ctx") == 2 + assert pcfg.max_owned_per_wave(plan, "gen") == 1 - # Non-ADP side participates everywhere and owns the whole chunk. + # Non-ADP side participates everywhere and owns the whole wave. plan_pp = pcfg.resolve_plan( _disagg_yaml( ctx_extra={ @@ -153,7 +153,7 @@ def test_pair_participation_and_chunks(): ) assert plan_pp["ctx"]["dp_size"] == 1 and plan_pp["n_pairs"] == 16 assert pcfg.pair_participates(plan_pp, "ctx", 0, 11) - assert pcfg.max_owned_per_chunk(plan_pp, "ctx") == plan_pp["chunk_size"] + assert pcfg.max_owned_per_wave(plan_pp, "ctx") == plan_pp["wave_size"] def test_fingerprint_role_agnostic(): @@ -222,7 +222,7 @@ class TestControlWireFormat: def test_roundtrip(self): key = b"\x01" * 32 - msg = ["go", {"li": 0, "rep": 1, "chunk": 2}] + msg = ["go", {"li": 0, "rep": 1, "wave": 2}] assert rp.unpack_msg(rp.pack_msg(msg, key), key) == msg def test_tampered_frame_rejected(self): @@ -403,18 +403,18 @@ def _mk_runner(self, role, server_idx, plan, work_dir, monkeypatch, fail_ctx=Fal side = pcfg.side_plan(plan, role) runner = rp.PrecheckRunner(args, plan, side, self._FakeComm()) - calls = {"chunks": 0} + calls = {"waves": 0} - def ctx_run_chunk(peer_idx, li, req_len, rep, chunk): + def ctx_run_wave(peer_idx, li, req_len, rep, wave): if fail_ctx: raise rp._TransferError("injected ctx failure") - calls["chunks"] += 1 - return {p: self._FakeParams() for p in chunk}, {} + calls["waves"] += 1 + return {p: self._FakeParams() for p in wave}, {} - runner.ctx_run_chunk = ctx_run_chunk - runner.ctx_finish_chunk = lambda reqs: None - runner.gen_run_chunk = ( - lambda peer_idx, li, req_len, rep, chunk, params: (True, "") + runner.ctx_run_wave = ctx_run_wave + runner.ctx_finish_wave = lambda reqs: None + runner.gen_run_wave = ( + lambda peer_idx, li, req_len, rep, wave, params: (True, "") ) runner._calls = calls return runner @@ -433,7 +433,7 @@ def _run(self, tmp_path, monkeypatch, fail_ctx1=False): "num_requests": 1, "warmup_requests": 1, "rendezvous_timeout_s": 30, - "chunk_timeout_s": 30, + "wave_timeout_s": 30, "wireup_timeout_s": 0, }, ) @@ -476,13 +476,13 @@ def test_two_ctx_full_pass(self, tmp_path, monkeypatch): ("ctx_0", "PASS"), ("ctx_1", "PASS"), } - # every ctx served the full schedule (reps x chunks) and got its + # every ctx served the full schedule (reps x waves) and got its # deferred done (PASS recorded only after done/bye completes) - total_chunks = len(pcfg.chunks(plan)) * ( + total_waves = len(pcfg.waves(plan)) * ( plan["warmup_requests"] + plan["num_requests"] ) for c in ctxs: - assert c._calls["chunks"] == total_chunks + assert c._calls["waves"] == total_waves assert [x["status"] for x in c.recorder.cases] == ["PASS"] def test_ctx_failure_isolated(self, tmp_path, monkeypatch): From 88e134b4a72f079359d3bf1d1ebe3f77ca7e251e Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Wed, 22 Jul 2026 01:19:04 -0700 Subject: [PATCH 15/21] [None][test] Fail-fast the cache-transceiver precheck and fix review findings Fail-fast: once any (ctx, gen) pair fails, stop testing the rest instead of running every pair against a fabric already known bad. A work-dir abort flag (job-id stamped, like the addr files, so a reused work dir does not carry a stale flag) is dropped on the first failure; the gen driver skips remaining peers and tells each to abort over the control channel, so the ctx side tears down promptly rather than waiting out its handshake alarm. Sessions that already passed still get their deferred "done". Skipped peers record a non-failing SKIP, which does not pollute the overall verdict. Review fixes: - run_precheck: broadcast the leader-only "missing context_phase_params" verdict so every rank raises together; previously only the leader raised while the others entered the next collective, deadlocking the step until the watchdog SIGKILLed it (misreported as TIMEOUT). - precheck_config: parse the TRTLLM_DISAGG_CT_PRECHECK kill switch as a boolean (1/0/true/false/on/off/yes/no) and reject ambiguous values, so a force-enable like =true is no longer silently read as "off". - precheck_config: add gate_library_content() as the single owner of loading and splicing the gate shell library; both submit.py call it (removing the drifted duplicate), and it falls back to the in-repo copy when the draft lives outside the repo instead of failing with FileNotFoundError. - gate: also clear csv/ on requeue so stale per-run bandwidth CSVs are not median'd across runs; print SKIPPED (not an empty "PASSED") when the run wrote no verdicts. Tests: cover fail-fast skip/abort, the abort-flag staleness rule, the missing-params broadcast, the kill-switch parsing, and the gate library fallback. Register the (pure-CPU) test file in l0_a10.yml. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../disaggregated/slurm_ct_precheck_gate.sh | 23 ++- jenkins/scripts/perf/local/submit.py | 8 +- jenkins/scripts/perf/submit.py | 10 +- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../precheck_config.py | 46 ++++- .../run_precheck.py | 155 ++++++++++++++--- .../test_cache_transceiver_precheck_config.py | 161 +++++++++++++++++- 7 files changed, 355 insertions(+), 49 deletions(-) diff --git a/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh b/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh index a5192eae7e7d..9a888ced5b94 100644 --- a/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh +++ b/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh @@ -112,12 +112,15 @@ run_cache_transceiver_precheck() { precheckDir="$testOutputDir/cache_transceiver_precheck" mkdir -p "$precheckDir/logs" # A reused work dir (Slurm requeue reruns this batch script with the same - # directories) may hold a previous run's rendezvous/status files: stale - # addr files would point gen leaders at dead ports, and stale status - # files would pollute the verdict aggregation below. The driver also - # job-id-stamps addr files as a second line of defense. + # directories) may hold a previous run's rendezvous/status/csv files: stale + # addr files would point gen leaders at dead ports, stale status files would + # pollute the verdict aggregation below, and stale bandwidth CSVs (the + # Python transceiver's perf__.csv are per-run and appended) + # would make parse_python_bandwidth_gbps median over two runs' samples. The + # driver also job-id-stamps addr files as a second line of defense. rm -f "$precheckDir"/rendezvous/*.addr "$precheckDir"/status/*.status \ "$precheckDir"/status/*.json 2>/dev/null || true + rm -rf "$precheckDir"/csv 2>/dev/null || true precheckPids=() precheckNames=() # ct_launch_step @@ -162,6 +165,14 @@ run_cache_transceiver_precheck() { ct_write_junit_xml cleanup_on_failure "Cache transceiver precheck failed. See summary above and $precheckDir" fi - echo "Cache transceiver precheck PASSED:" - cat "$precheckDir"/status/*.status 2>/dev/null || true + # No status files means every step took the skip path (e.g. the yaml has no + # cache_transceiver_config.backend): the run wrote no verdicts, so report it + # as SKIPPED rather than an empty "PASSED" that reads like real validation. + if ls "$precheckDir"/status/*.status >/dev/null 2>&1; then + echo "Cache transceiver precheck PASSED:" + cat "$precheckDir"/status/*.status 2>/dev/null || true + else + echo "Cache transceiver precheck SKIPPED: not applicable for this config" \ + "(no verdicts written; see $precheckDir/logs)" + fi } diff --git a/jenkins/scripts/perf/local/submit.py b/jenkins/scripts/perf/local/submit.py index f603b0d02ea2..f1b7ab75de15 100755 --- a/jenkins/scripts/perf/local/submit.py +++ b/jenkins/scripts/perf/local/submit.py @@ -1036,13 +1036,11 @@ def main(): draft_launch_lines = remove_whitespace_lines(draft_launch_lines) draft_launch_content = "\n".join(draft_launch_lines) - # The disagg draft calls run_cache_transceiver_precheck; splice in the - # gate function library (sibling of the draft) ahead of it. + # The disagg draft calls run_cache_transceiver_precheck; splice in the gate + # function library ahead of it (single owner: precheck_config). gate_content = "" if runtime_mode == "disaggregated": - gate_sh = os.path.join(os.path.dirname(draft_launch_sh), "slurm_ct_precheck_gate.sh") - with open(gate_sh, "r") as f: - gate_content = "\n".join(remove_whitespace_lines(f.read().split("\n"))) + "\n\n" + gate_content = pcfg.gate_library_content(draft_launch_sh, llm_src) # Combine and write launch script script_prefix = "\n".join(script_prefix_lines) diff --git a/jenkins/scripts/perf/submit.py b/jenkins/scripts/perf/submit.py index 83c7afa0da3d..abb9e42a808c 100755 --- a/jenkins/scripts/perf/submit.py +++ b/jenkins/scripts/perf/submit.py @@ -678,15 +678,11 @@ def main(): draft_launch_lines = remove_whitespace_lines(draft_launch_content.split("\n")) draft_launch_content = "\n".join(draft_launch_lines) - # The disagg draft calls run_cache_transceiver_precheck; splice in the - # gate function library (sibling of the draft) ahead of it. + # The disagg draft calls run_cache_transceiver_precheck; splice in the gate + # function library ahead of it (single owner: precheck_config). gate_content = "" if runtime_mode == "disaggregated": - gate_sh = os.path.join( - os.path.dirname(args.draft_launch_sh), "slurm_ct_precheck_gate.sh" - ) - with open(gate_sh, "r") as f: - gate_content = "\n".join(remove_whitespace_lines(f.read().split("\n"))) + "\n" + gate_content = pcfg.gate_library_content(args.draft_launch_sh, args.llm_src) with open(args.launch_sh, "w") as f: f.write(f"{script_prefix}\n{srun_args}\n{gate_content}{draft_launch_content}") diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index b5d9a4ba7c18..c0ce6be1abf3 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -55,6 +55,7 @@ l0_a10: - unittest/inputs/test_multimodal.py - unittest/inputs/test_multimodal_input_processor.py - unittest/inputs/test_video_decode.py + - unittest/others/test_cache_transceiver_precheck_config.py - unittest/others/test_convert_utils.py - unittest/others/test_lora_manager.py - unittest/others/test_lora_module_count.py diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py index 5ef6455d09f4..229637f1ca70 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py @@ -131,10 +131,24 @@ def precheck_prefix_lines(cfg, benchmark_mode, config_path_expr, ucx_tls_cmd, ma expressions ($llmSrcNode etc.), expanded at sbatch runtime. """ knobs = cfg.get("cache_transceiver_precheck", {}) or {} - env = os.environ.get("TRTLLM_DISAGG_CT_PRECHECK") # On by default; yaml opts out per test; the env var (when set) overrides - # the yaml either way (global kill switch). - enabled = env == "1" if env is not None else bool(knobs.get("enabled", True)) + # the yaml either way (global kill switch). Parse the usual boolean spellings + # so a well-meant TRTLLM_DISAGG_CT_PRECHECK=true force-enable is not silently + # read as "off"; reject anything ambiguous instead of guessing. + env = os.environ.get("TRTLLM_DISAGG_CT_PRECHECK") + if env is not None: + val = env.strip().lower() + if val in ("1", "true", "on", "yes"): + enabled = True + elif val in ("0", "false", "off", "no"): + enabled = False + else: + raise ValueError( + "TRTLLM_DISAGG_CT_PRECHECK must be a boolean " + f"(1/0/true/false/on/off/yes/no), got {env!r}" + ) + else: + enabled = bool(knobs.get("enabled", True)) cmd = ( "python3 $llmSrcNode/tests/scripts/perf-sanity/cache_transceiver_precheck/" f"run_precheck.py --config {config_path_expr} " @@ -160,6 +174,32 @@ def precheck_prefix_lines(cfg, benchmark_mode, config_path_expr, ucx_tls_cmd, ma return lines +def gate_library_content(draft_launch_sh, llm_src): + """The precheck gate shell library (defines run_cache_transceiver_precheck), + spliced ahead of the disagg draft by both submit.py. Single owner of the + load/splice, so the two launch generators can't drift. + + Located next to the draft by default; falls back to the in-repo copy when + the draft lives outside the repo (a custom --draft-launch-sh), so launch + generation never dies with FileNotFoundError. Whitespace-only lines are + dropped to match the rest of the assembled launch script.""" + candidates = [ + os.path.join(os.path.dirname(draft_launch_sh), "slurm_ct_precheck_gate.sh"), + os.path.join( + llm_src, "jenkins", "scripts", "perf", "disaggregated", "slurm_ct_precheck_gate.sh" + ), + ] + for path in candidates: + if os.path.exists(path): + with open(path) as f: + lines = [ln.strip() for ln in f.read().split("\n") if ln.strip()] + return "\n".join(lines) + "\n" + raise FileNotFoundError( + "precheck gate library slurm_ct_precheck_gate.sh not found; looked in: " + + ", ".join(candidates) + ) + + def resolve_plan(cfg, benchmark_mode="e2e"): """Build the shared precheck plan both roles must agree on. diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index 0ae93d9e0e11..8b3d6a01f980 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -514,6 +514,46 @@ def wait_for_addr(path, timeout_s): raise _Timeout(f"rendezvous file {path} not published within {timeout_s}s") +def abort_flag_path(work_dir): + return os.path.join(work_dir, "precheck.abort") + + +def raise_abort_flag(work_dir, reason): + """Fail-fast signal, shared across instances through the work dir: the + first peer failure (in ANY ctx/gen instance) drops this file so the others + stop starting new sessions instead of each re-discovering the dead fabric + on its own. Best-effort and write-once. Stamped with the job id so a stale + flag in a reused work dir (Slurm requeue) is ignored, like addr files.""" + path = abort_flag_path(work_dir) + if abort_flag_reason(work_dir) is not None: + return + try: + tmp = f"{path}.tmp.{os.getpid()}" + with open(tmp, "w") as f: + json.dump( + {"reason": (reason or "peer failure").splitlines()[0][:400], "job": run_token()}, + f, + ) + os.replace(tmp, path) + except OSError: + pass # best-effort: a missed flag only costs the usual per-peer timeout + + +def abort_flag_reason(work_dir): + """The reason recorded by the first failing peer of THIS run, or None. A + flag stamped with another job id is stale (reused work dir) and ignored.""" + try: + with open(abort_flag_path(work_dir)) as f: + payload = json.load(f) + except (OSError, json.JSONDecodeError): + return None + expect_job = run_token() + stamped = payload.get("job", "") + if expect_job and stamped and stamped != expect_job: + return None + return payload.get("reason") or "peer failure" + + # --------------------------------------------------------------------------- # # Status recording # --------------------------------------------------------------------------- # @@ -792,12 +832,19 @@ def ctx_run_wave(self, peer_idx, li, req_len, rep, wave): ) gathered = self.comm.gather(contrib, root=0) params_by_pair = {} - if self.is_leader and reason is None: + if self.is_leader: for d in gathered: params_by_pair.update(d or {}) - missing = [p for p in wave if p not in params_by_pair] - if missing: - reason = f"missing context_phase_params for pairs {missing}" + if reason is None: + missing = [p for p in wave if p not in params_by_pair] + if missing: + reason = f"missing context_phase_params for pairs {missing}" + # The missing-params check runs only on the leader (only it holds the + # gathered params). Broadcast the verdict so EVERY rank raises together: + # otherwise the leader raises here while the other ranks return and enter + # ctx_finish_wave's collective, the collective sequence diverges, and the + # step deadlocks until the watchdog SIGKILLs it (misreported as TIMEOUT). + reason = self.comm.bcast(reason, root=0) if reason is not None: self._free_all(reqs) @@ -1011,6 +1058,9 @@ def leader_reply(obj): arm(f"gen_{peer_idx} len={req_len} rep={rep}", seconds=wave_timeout_s(plan, li, rep)) msg = leader_recv() if msg[0] == "abort": + # Ack so the gen's REQ send/recv completes (fail-fast teardown + # sends this in place of the schedule; see gen_abort_peer). + leader_reply(("aborted", {})) raise _PeerAbort(f"gen_{peer_idx} aborted: {msg[1]}") if msg[0] != "go" or (msg[1]["li"], msg[1]["rep"]) != (li, rep): raise _TransferError( @@ -1039,20 +1089,19 @@ def leader_reply(obj): disarm() -def gen_run_peer(runner, peer_idx, arm, disarm): - """Run the full schedule against ctx server `peer_idx`. +def _gen_open_session(runner, peer_idx, arm): + """Rendezvous with ctx server `peer_idx` and complete the hello/welcome + handshake; return (sock, key) on an OPEN REQ socket. - Returns (sock, key) with the session STILL OPEN on success -- the caller - sends the deferred "done" only after every ctx peer's schedule finished, - keeping all ctx instances alive for the whole precheck (real-serving - lifecycle; see ctx_serve_peer). On failure the socket is closed here. + Shared by gen_run_peer (runs the schedule) and gen_abort_peer (sends an + early abort for fail-fast). Rendezvous reaches instance-wide consensus via + bcast: if only the leader raised, the other ranks would deadlock in the + next bcast. The socket is closed here if the handshake itself fails. """ plan = runner.plan comm = runner.comm hello_s = hello_timeout_s(plan, runner.side["num_peers"]) - # Rendezvous with instance-wide consensus: if only the leader raised here, - # the other ranks would deadlock in the next bcast. sock, key, err = None, None, None arm(f"rendezvous ctx_{peer_idx}", seconds=hello_s) if runner.is_leader: @@ -1086,6 +1135,24 @@ def gen_run_peer(runner, peer_idx, arm, disarm): raise _TransferError(f"ctx_{peer_idx} aborted handshake: {reply[1]}") if reply[0] != "welcome": raise _TransferError(f"unexpected handshake reply from ctx_{peer_idx}: {reply[:1]}") + return sock, key + except BaseException: + if sock is not None: + sock.close(linger=0) + raise + + +def gen_run_peer(runner, peer_idx, arm, disarm): + """Run the full schedule against ctx server `peer_idx`. + + Returns (sock, key) with the session STILL OPEN on success -- the caller + sends the deferred "done" only after every ctx peer's schedule finished, + keeping all ctx instances alive for the whole precheck (real-serving + lifecycle; see ctx_serve_peer). On failure the socket is closed here. + """ + plan = runner.plan + sock, key = _gen_open_session(runner, peer_idx, arm) + try: # Established sessions are dedicated: wave replies are prompt. The # ZMQ timeout is only a backstop under the per-wave alarm, so it # includes the first-rep wire-up allowance unconditionally. @@ -1125,6 +1192,21 @@ def gen_run_peer(runner, peer_idx, arm, disarm): raise +def gen_abort_peer(runner, peer_idx, reason, arm, disarm): + """Fail-fast teardown of a not-yet-run ctx peer: open the session and tell + it to abort (it is blocked awaiting our hello), so it stops promptly + instead of waiting out the handshake alarm. Best-effort; always closes.""" + sock = None + try: + sock, key = _gen_open_session(runner, peer_idx, arm) + arm(f"abort ctx_{peer_idx}", seconds=hello_timeout_s(runner.plan, runner.side["num_peers"])) + runner._leader_send_recv(sock, ("abort", f"peer fail-fast: {reason}"), key) + disarm() + finally: + if sock is not None: + sock.close(linger=0) + + def gen_release_peer(runner, peer_idx, sock, key, arm, disarm): """Deferred session teardown: send "done" and close (best-effort).""" try: @@ -1220,19 +1302,21 @@ def stop(): def _make_peer_failure_recorder(runner, disarm, current_cell): - """Exception -> verdict mapping shared by all per-peer loops (isolation: - one bad peer never stops the others).""" + """Exception -> verdict mapping shared by all per-peer loops. Recording a + failure also drops the fail-fast flag so the remaining peers (here and in + the other instances) are skipped instead of tested against a fabric + already known bad -- see _drive_ctx_peers / raise_abort_flag.""" def record_peer_failure(peer, exc): disarm() if isinstance(exc, _Timeout): - runner.recorder.record( - peer, 0, "TIMEOUT", f"exceeded the budget during {current_cell['what']}" - ) + status, reason = "TIMEOUT", f"exceeded the budget during {current_cell['what']}" elif isinstance(exc, _PeerAbort): - runner.recorder.record(peer, 0, "TRANSFER_ERROR", str(exc)) + status, reason = "TRANSFER_ERROR", str(exc) else: - runner.recorder.record(peer, 0, "TRANSFER_ERROR", repr(exc)) + status, reason = "TRANSFER_ERROR", repr(exc) + runner.recorder.record(peer, 0, status, reason) + raise_abort_flag(runner.work_dir, f"{peer} {status}: {reason}") return record_peer_failure @@ -1263,6 +1347,15 @@ def _serve_gen_peers(runner, plan, arm, disarm, record_peer_failure): try: ctx_serve_peer(runner, socks.get(gj), gj, arm, disarm, keys.get(gj)) runner.recorder.record(f"gen_{gj}", 0, "PASS", "served all transfers") + except _PeerAbort as e: + # A gen driver that failed elsewhere aborts our session as part of + # fail-fast: record a (non-failing) SKIP, not our own failure -- + # the real failure is recorded by whoever hit it. Absent the flag, + # a genuine peer abort is still a real failure. + if abort_flag_reason(runner.work_dir) is not None: + runner.recorder.record(f"gen_{gj}", 0, "SKIP", f"aborted by fail-fast: {e}") + else: + record_peer_failure(f"gen_{gj}", e) except Exception as e: # noqa: BLE001 - per-peer isolation record_peer_failure(f"gen_{gj}", e) @@ -1270,16 +1363,30 @@ def _serve_gen_peers(runner, plan, arm, disarm, record_peer_failure): def _drive_ctx_peers(runner, arm, disarm, record_peer_failure): """gen role: run every ctx peer's schedule, then release all sessions. - The release ("done") is deferred until EVERY peer's schedule finished, so - all ctx instances stay alive for the whole precheck -- matching real - serving, where no transceiver ever holds connections to a dead agent - while transfers are still running.""" + Fail-fast: once any pair has failed (this instance or another -- signalled + through the work-dir abort flag), the remaining ctx peers are not tested; + each is told to abort so it tears down promptly rather than waiting out its + handshake alarm. Sessions that already succeeded still get their deferred + "done" (below) so those ctx instances shut down cleanly. + + The release ("done") is deferred until EVERY driven peer's schedule + finished, so those ctx instances stay alive for the whole precheck -- + matching real serving, where no transceiver ever holds connections to a + dead agent while transfers are still running.""" open_sessions = [] for ci in range(runner.side["num_peers"]): + reason = abort_flag_reason(runner.work_dir) + if reason is not None: + try: + gen_abort_peer(runner, ci, reason, arm, disarm) + except Exception: # noqa: BLE001 - best-effort teardown + pass + runner.recorder.record(f"ctx_{ci}", 0, "SKIP", f"fail-fast: {reason}") + continue try: sock, sess_key = gen_run_peer(runner, ci, arm, disarm) open_sessions.append((ci, sock, sess_key)) - except Exception as e: # noqa: BLE001 - per-peer isolation + except Exception as e: # noqa: BLE001 - failure sets the fail-fast flag record_peer_failure(f"ctx_{ci}", e) for ci, sock, sess_key in open_sessions: try: diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py index 296b8f1948a5..2d63de210458 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_config.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py @@ -332,6 +332,56 @@ def test_wireup_timeout_derivation(): assert plan["wireup_timeout_s"] == 42 +def _enabled_line(cfg): + lines = pcfg.precheck_prefix_lines(cfg, "e2e", "$c", "unset &&", max_world=8) + return next(x for x in lines if x.startswith("export ctPrecheckEnabled")) + + +def test_precheck_env_kill_switch_truthy(monkeypatch): + """The TRTLLM_DISAGG_CT_PRECHECK global kill switch parses the usual boolean + spellings (so a force-enable like =true is not silently read as "off") and + rejects anything ambiguous instead of guessing.""" + cfg = {"cache_transceiver_precheck": {"enabled": True}} + monkeypatch.delenv("TRTLLM_DISAGG_CT_PRECHECK", raising=False) + assert _enabled_line(cfg).endswith("=1") # yaml default + for v in ("1", "true", "on", "YES", " True "): + monkeypatch.setenv("TRTLLM_DISAGG_CT_PRECHECK", v) + assert _enabled_line(cfg).endswith("=1"), v + for v in ("0", "false", "off", "no"): + monkeypatch.setenv("TRTLLM_DISAGG_CT_PRECHECK", v) + assert _enabled_line(cfg).endswith("=0"), v + # env overrides yaml either way (kill switch): yaml opt-out but env force-on + monkeypatch.setenv("TRTLLM_DISAGG_CT_PRECHECK", "true") + assert _enabled_line({"cache_transceiver_precheck": {"enabled": False}}).endswith("=1") + monkeypatch.setenv("TRTLLM_DISAGG_CT_PRECHECK", "maybe") + with pytest.raises(ValueError): + _enabled_line(cfg) + + +def test_gate_library_content(tmp_path): + """The gate library loads from next to the draft, falls back to the in-repo + copy for an external draft, strips blank lines, and errors if truly absent.""" + dd = tmp_path / "disaggregated" + dd.mkdir() + (dd / "slurm_ct_precheck_gate.sh").write_text( + "run_cache_transceiver_precheck() { :; }\n\n\nrun_cache_transceiver_precheck\n" + ) + draft = str(dd / "slurm_launch_draft.sh") + got = pcfg.gate_library_content(draft, str(tmp_path)) + assert "run_cache_transceiver_precheck()" in got + assert got.endswith("\n") and "\n\n" not in got # blank lines stripped + + # external draft -> fall back to /jenkins/.../slurm_ct_precheck_gate.sh + repo = tmp_path / "repo" + gate2 = repo / "jenkins" / "scripts" / "perf" / "disaggregated" / "slurm_ct_precheck_gate.sh" + gate2.parent.mkdir(parents=True) + gate2.write_text("echo hi\n") + assert pcfg.gate_library_content("/nowhere/draft.sh", str(repo)) == "echo hi\n" + + with pytest.raises(FileNotFoundError): + pcfg.gate_library_content("/nowhere/draft.sh", str(tmp_path / "empty")) + + def test_rid_tags_dense_within_session(): """The C++ notification tag is rid & 0xFFF: rids must be dense within a (ctx, gen) session so tags cannot alias across reps/lengths.""" @@ -419,7 +469,7 @@ def ctx_run_wave(peer_idx, li, req_len, rep, wave): runner._calls = calls return runner - def _run(self, tmp_path, monkeypatch, fail_ctx1=False): + def _run(self, tmp_path, monkeypatch, fail_ctx_idx=None): import threading monkeypatch.setenv("SLURM_JOB_ID", "777") @@ -443,7 +493,10 @@ def _run(self, tmp_path, monkeypatch, fail_ctx1=False): gen = self._mk_runner("gen", 0, plan, work, monkeypatch) ctxs = [ - self._mk_runner("ctx", i, plan, work, monkeypatch, fail_ctx=(fail_ctx1 and i == 1)) + self._mk_runner( + "ctx", i, plan, work, monkeypatch, + fail_ctx=(fail_ctx_idx is not None and i == fail_ctx_idx), + ) for i in range(2) ] @@ -485,8 +538,10 @@ def test_two_ctx_full_pass(self, tmp_path, monkeypatch): assert c._calls["waves"] == total_waves assert [x["status"] for x in c.recorder.cases] == ["PASS"] - def test_ctx_failure_isolated(self, tmp_path, monkeypatch): - plan, gen, ctxs, failures = self._run(tmp_path, monkeypatch, fail_ctx1=True) + def test_ctx_failure_last_peer(self, tmp_path, monkeypatch): + # The failing pair is driven LAST: the earlier healthy peer already + # completed, so there is nothing left to fail-fast/skip. + plan, gen, ctxs, failures = self._run(tmp_path, monkeypatch, fail_ctx_idx=1) # gen side: healthy peer unaffected, failing peer gets a clear verdict by_peer = {c["peer"]: c["status"] for c in gen.recorder.cases} assert by_peer == {"ctx_0": "PASS", "ctx_1": "TRANSFER_ERROR"} @@ -495,6 +550,104 @@ def test_ctx_failure_isolated(self, tmp_path, monkeypatch): # ctx_0 served its full schedule and got the deferred done assert [c["status"] for c in ctxs[0].recorder.cases] == ["PASS"] + def test_fail_fast_skips_remaining(self, tmp_path, monkeypatch): + # The FIRST-driven pair fails: the remaining pair must be skipped + # (not tested against a fabric already known bad), and told to abort + # so it tears down promptly instead of waiting out its handshake alarm. + plan, gen, ctxs, failures = self._run(tmp_path, monkeypatch, fail_ctx_idx=0) + by_peer = {c["peer"]: c["status"] for c in gen.recorder.cases} + assert by_peer == {"ctx_0": "TRANSFER_ERROR", "ctx_1": "SKIP"} + # ctx_1 never ran a single transfer wave: fail-fast reached it first. + assert ctxs[1]._calls["waves"] == 0 + # ctx_1 recorded a non-failing SKIP (its driver aborted the session). + assert [c["status"] for c in ctxs[1].recorder.cases] == ["SKIP"] + # the shared, job-stamped abort flag was dropped. + assert rp.abort_flag_reason(str(tmp_path)) is not None + # SKIP does not count toward the overall verdict; only ctx_0 failed. + assert [c["peer"] for c in gen.recorder.failed_cases()] == ["ctx_0"] + + def test_abort_flag_stale_job_ignored(self, tmp_path, monkeypatch): + # A flag left by a previous run (different SLURM_JOB_ID) in a reused + # work dir must not fail-fast a fresh run -- same staleness rule as + # addr files. + monkeypatch.setenv("SLURM_JOB_ID", "111") + rp.raise_abort_flag(str(tmp_path), "old run failure") + assert rp.abort_flag_reason(str(tmp_path)) == "old run failure" + monkeypatch.setenv("SLURM_JOB_ID", "222") + assert rp.abort_flag_reason(str(tmp_path)) is None + + +def test_ctx_run_wave_missing_params_broadcast(tmp_path, monkeypatch): + """#4 regression: the "missing context_phase_params" verdict is computed + only on the instance leader (only it holds the gathered params), so it must + be broadcast -- otherwise a NON-leader rank keeps reason=None, returns, and + enters the next collective while the leader raises, deadlocking the step + until the watchdog SIGKILLs it (misreported as TIMEOUT). + + Run the real ctx_run_wave on a non-leader rank: with the leader's verdict + delivered via bcast the rank must raise; with a clean (None) broadcast it + must return normally. + """ + import sys + import types + + monkeypatch.setitem(sys.modules, "mpi4py", types.SimpleNamespace(MPI=None)) + # ctx_run_wave imports tensorrt_llm only for logger.info, never reached with + # no owned pairs; a stub keeps the test pure-CPU. + monkeypatch.setitem( + sys.modules, + "tensorrt_llm", + types.SimpleNamespace(logger=types.SimpleNamespace(info=lambda *a, **k: None)), + ) + + class _NonLeaderComm: + def __init__(self, bcast_ret): + self._bcast_ret = bcast_ret + + def Get_rank(self): + return 1 # non-leader (leader is rank 0) + + def Get_size(self): + return 2 + + def allgather(self, obj): + return ["", ""] # no local send error on any rank + + def gather(self, obj, root=0): + return None # only the leader receives the gathered params + + def bcast(self, obj, root=0): + return self._bcast_ret # the leader's verdict reaching this rank + + plan = pcfg.resolve_plan( + _disagg_yaml( + cache_transceiver_precheck={ + "request_lengths": [32], + "num_requests": 1, + "warmup_requests": 1, + "wireup_timeout_s": 0, + } + ) + ) + side = pcfg.side_plan(plan, "ctx") + args = types.SimpleNamespace(server_idx=0, work_dir=str(tmp_path)) + + def _mk(bcast_ret): + r = rp.PrecheckRunner(args, plan, side, _NonLeaderComm(bcast_ret)) + r.mapping = types.SimpleNamespace(pp_rank=1, tp_rank=1) # unused non-leader + r._owned = lambda wave: [] # skip the GPU send path; verdict arrives via bcast + return r + + # leader broadcast a missing-params verdict -> the non-leader raises it too + with pytest.raises(rp._TransferError, match="missing context_phase_params"): + _mk("missing context_phase_params for pairs [0]").ctx_run_wave( + peer_idx=0, li=0, req_len=32, rep=0, wave=[0] + ) + + # leader broadcast None (all good) -> the non-leader returns cleanly + params, reqs = _mk(None).ctx_run_wave(peer_idx=0, li=0, req_len=32, rep=0, wave=[0]) + assert params == {} and reqs == {} + def test_status_env_snapshot_excludes_nixl(tmp_path, monkeypatch): """NIXL_* is not captured: the only such variable seen in practice is From a1dd9c91fe065c558bb5bbcc235f13812469923d Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Wed, 22 Jul 2026 08:53:12 +0000 Subject: [PATCH 16/21] [None][test] Fix pre-commit lint findings across the precheck code Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- jenkins/scripts/perf/local/submit.py | 9 +- jenkins/scripts/perf/submit.py | 7 +- .../precheck_config.py | 26 ++-- .../run_precheck.py | 111 +++++++++++------- .../test_cache_transceiver_precheck_config.py | 85 ++++++++------ 5 files changed, 148 insertions(+), 90 deletions(-) diff --git a/jenkins/scripts/perf/local/submit.py b/jenkins/scripts/perf/local/submit.py index f1b7ab75de15..538af4bc5280 100755 --- a/jenkins/scripts/perf/local/submit.py +++ b/jenkins/scripts/perf/local/submit.py @@ -9,10 +9,14 @@ import sys from datetime import datetime +import yaml + def _import_precheck_config(llm_src): - """Import the pure-stdlib precheck config module from the repo tree - (single owner of the gate's enable policy and timeout formulas).""" + """Import the pure-stdlib precheck config module from the repo tree. + + It is the single owner of the gate's enable policy and timeout formulas. + """ path = os.path.join(llm_src, "tests", "scripts", "perf-sanity", "cache_transceiver_precheck") if path not in sys.path: sys.path.insert(0, path) @@ -20,7 +24,6 @@ def _import_precheck_config(llm_src): return precheck_config -import yaml AGG_CONFIG_FOLDER = os.environ.get("AGG_CONFIG_FOLDER", "tests/scripts/perf-sanity/aggregated") DISAGG_CONFIG_FOLDER = os.environ.get( diff --git a/jenkins/scripts/perf/submit.py b/jenkins/scripts/perf/submit.py index abb9e42a808c..69f69bac27ea 100755 --- a/jenkins/scripts/perf/submit.py +++ b/jenkins/scripts/perf/submit.py @@ -27,8 +27,10 @@ def _import_precheck_config(llm_src): - """Import the pure-stdlib precheck config module from the repo tree - (single owner of the gate's enable policy and timeout formulas).""" + """Import the pure-stdlib precheck config module from the repo tree. + + It is the single owner of the gate's enable policy and timeout formulas. + """ path = os.path.join(llm_src, "tests", "scripts", "perf-sanity", "cache_transceiver_precheck") if path not in sys.path: sys.path.insert(0, path) @@ -36,6 +38,7 @@ def _import_precheck_config(llm_src): return precheck_config + AGG_CONFIG_FOLDER = "tests/scripts/perf-sanity/aggregated" DISAGG_CONFIG_FOLDER = "tests/scripts/perf-sanity/disaggregated" diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py index 229637f1ca70..c02e24a20fbc 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py @@ -113,15 +113,18 @@ def wireup_timeout_s(max_world): def default_step_timeout_s(max_world): - """External (srun-level) timeout covering one precheck instance, - including the first-rep wire-up. Imported by the launch tooling + """External (srun-level) timeout covering one precheck instance. + + Includes the first-rep wire-up. Imported by the launch tooling (jenkins/scripts/perf/{,local/}submit.py) so the outer timeout can - never drift below the driver's internal budget.""" + never drift below the driver's internal budget. + """ return 900 + wireup_timeout_s(max_world) -def precheck_prefix_lines(cfg, benchmark_mode, config_path_expr, ucx_tls_cmd, max_world, - stage_name=""): +def precheck_prefix_lines( + cfg, benchmark_mode, config_path_expr, ucx_tls_cmd, max_world, stage_name="" +): """Launch-script export lines wiring the precheck gate. Single owner of the enable/kill-switch policy, the step-timeout default, @@ -175,14 +178,16 @@ def precheck_prefix_lines(cfg, benchmark_mode, config_path_expr, ucx_tls_cmd, ma def gate_library_content(draft_launch_sh, llm_src): - """The precheck gate shell library (defines run_cache_transceiver_precheck), - spliced ahead of the disagg draft by both submit.py. Single owner of the + """The precheck gate shell library (defines run_cache_transceiver_precheck). + + Spliced ahead of the disagg draft by both submit.py. Single owner of the load/splice, so the two launch generators can't drift. Located next to the draft by default; falls back to the in-repo copy when the draft lives outside the repo (a custom --draft-launch-sh), so launch generation never dies with FileNotFoundError. Whitespace-only lines are - dropped to match the rest of the assembled launch script.""" + dropped to match the rest of the assembled launch script. + """ candidates = [ os.path.join(os.path.dirname(draft_launch_sh), "slurm_ct_precheck_gate.sh"), os.path.join( @@ -427,7 +432,10 @@ def unmodelable_kv_reason(hf_cfg): just with generic bytes. """ if any(k in hf_cfg for k in ("index_topk", "index_n_heads", "index_head_dim")): - return "sparse-attention model (DSA/indexer): using a generic KV pool -- this checks the network/transceiver path, not V4's exact KV layout" + return ( + "sparse-attention model (DSA/indexer): using a generic KV pool -- this checks " + "the network/transceiver path, not V4's exact KV layout" + ) return None diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index 8b3d6a01f980..9dad2594b9e1 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -151,8 +151,11 @@ def params_to_wire(p): def params_from_wire(d): - """Inverse of params_to_wire, via the maintained DisaggregatedParams - converter (keeps us off the raw nanobind ctor signature).""" + """Inverse of params_to_wire. + + Uses the maintained DisaggregatedParams converter (keeps us off the raw + nanobind ctor signature). + """ from tensorrt_llm import DisaggregatedParams return DisaggregatedParams( @@ -229,7 +232,6 @@ def _lookup_model_cls(model_dir): if not archs: return None, hf_view import tensorrt_llm._torch.models # noqa: F401 - populates the registry - from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPING return MODEL_CLASS_MAPPING.get(archs[0]), hf_view @@ -258,8 +260,10 @@ def resolve_model_prefs(model_dir, side, cache_cfg): try: defaults = model_cls.get_model_defaults(None) or {} except Exception as e: # noqa: BLE001 - model hooks may need llm_args - print(f"[precheck] WARNING: get_model_defaults failed ({e!r}); " - f"assuming V1", flush=True) + print( + f"[precheck] WARNING: get_model_defaults failed ({e!r}); assuming V1", + flush=True, + ) try: # The REAL serving resolver, via the same shim pattern as the # runtime resolution below -- one owner for the 'auto' semantics. @@ -270,8 +274,9 @@ def resolve_model_prefs(model_dir, side, cache_cfg): ) use_v2 = bool(_resolve_kv_cache_manager_v2_auto(shim, defaults)) except Exception as e: # noqa: BLE001 - fall back like a missing model - print(f"[precheck] WARNING: V2 'auto' resolution failed ({e!r}); " - f"assuming V1", flush=True) + print( + f"[precheck] WARNING: V2 'auto' resolution failed ({e!r}); assuming V1", flush=True + ) use_v2 = False else: use_v2 = bool(setting) @@ -479,9 +484,12 @@ def run_token(): def write_addr(path, payload): - """Atomically publish an addr file (os.replace overwrites stale ones; - readers reject wrong-job stamps). It carries the session HMAC key, so - restrict it to the owning user before it becomes visible.""" + """Atomically publish an addr file. + + os.replace overwrites stale ones; readers reject wrong-job stamps. It + carries the session HMAC key, so restrict it to the owning user before it + becomes visible. + """ os.makedirs(os.path.dirname(path), exist_ok=True) payload = dict(payload, job=run_token()) tmp = f"{path}.tmp.{os.getpid()}" @@ -494,8 +502,11 @@ def write_addr(path, payload): def wait_for_addr(path, timeout_s): - """Wait for THIS run's addr file; files stamped with another run's job id - are treated as stale and skipped (keep polling).""" + """Wait for THIS run's addr file. + + Files stamped with another run's job id are treated as stale and skipped + (keep polling). + """ expect_job = run_token() deadline = time.monotonic() + timeout_s while time.monotonic() < deadline: @@ -519,11 +530,14 @@ def abort_flag_path(work_dir): def raise_abort_flag(work_dir, reason): - """Fail-fast signal, shared across instances through the work dir: the - first peer failure (in ANY ctx/gen instance) drops this file so the others - stop starting new sessions instead of each re-discovering the dead fabric - on its own. Best-effort and write-once. Stamped with the job id so a stale - flag in a reused work dir (Slurm requeue) is ignored, like addr files.""" + """Fail-fast signal, shared across instances through the work dir. + + The first peer failure (in ANY ctx/gen instance) drops this file so the + others stop starting new sessions instead of each re-discovering the dead + fabric on its own. Best-effort and write-once. Stamped with the job id so + a stale flag in a reused work dir (Slurm requeue) is ignored, like addr + files. + """ path = abort_flag_path(work_dir) if abort_flag_reason(work_dir) is not None: return @@ -540,8 +554,10 @@ def raise_abort_flag(work_dir, reason): def abort_flag_reason(work_dir): - """The reason recorded by the first failing peer of THIS run, or None. A - flag stamped with another job id is stale (reused work dir) and ignored.""" + """The reason recorded by the first failing peer of THIS run, or None. + + A flag stamped with another job id is stale (reused work dir) and ignored. + """ try: with open(abort_flag_path(work_dir)) as f: payload = json.load(f) @@ -653,8 +669,9 @@ def parse_bandwidth_gbps(csv_dir, rank, tag="recv"): def parse_python_bandwidth_gbps(csv_dir): - """Median KV-send throughput in GB/s from the Python transceiver's perf - CSVs (perf_logger.py). Enabled via TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO=1 + """Median KV-send throughput in GB/s from the Python transceiver's perf CSVs. + + Written by perf_logger.py; enabled via TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO=1 + TLLM_KV_TRANSFER_PERF_LOG_FILE=, which writes __.csv. throughput_mbs (MB/s) is on the SENDER (ctx) side, task_type=KVSendTask; receiver rows have no throughput. @@ -755,8 +772,10 @@ def setup(self, kv_shape, max_req_len): # match what the real model runs (e.g. V4 -> V2 + Python). self.use_v2 = resolve_model_prefs(self.plan.get("_model_dir"), self.side, cache_cfg) if kv_shape.get("simplified") and self.is_leader: - print(f"[precheck {self.role}_{self.server_idx}] SIMPLIFIED: " - f"{kv_shape['simplified']}", flush=True) + print( + f"[precheck {self.role}_{self.server_idx}] SIMPLIFIED: {kv_shape['simplified']}", + flush=True, + ) # KVCacheManagerV2 only works with the Python transceiver (see # cache_transceiver_test/report.py); reject the pairing up front with # a clear INIT_ERROR instead of a C++ binding type error. @@ -998,8 +1017,11 @@ def hello_timeout_s(plan, num_peers): def wave_timeout_s(plan, li, rep): - """Per-wave budget: the first rep additionally pays the one-time NIXL - agent wire-up (see PRECHECK_DEFAULTS['wireup_timeout_s']).""" + """Per-wave budget. + + The first rep additionally pays the one-time NIXL agent wire-up (see + PRECHECK_DEFAULTS['wireup_timeout_s']). + """ extra = plan["wireup_timeout_s"] if (li == 0 and rep == 0) else 0 return plan["wave_timeout_s"] + extra @@ -1090,10 +1112,9 @@ def leader_reply(obj): def _gen_open_session(runner, peer_idx, arm): - """Rendezvous with ctx server `peer_idx` and complete the hello/welcome - handshake; return (sock, key) on an OPEN REQ socket. + """Rendezvous with ctx server `peer_idx`; return (sock, key) on an OPEN REQ socket. - Shared by gen_run_peer (runs the schedule) and gen_abort_peer (sends an + Completes the hello/welcome handshake. Shared by gen_run_peer (runs the schedule) and gen_abort_peer (sends an early abort for fail-fast). Rendezvous reaches instance-wide consensus via bcast: if only the leader raised, the other ranks would deadlock in the next bcast. The socket is closed here if the handshake itself fails. @@ -1193,9 +1214,12 @@ def gen_run_peer(runner, peer_idx, arm, disarm): def gen_abort_peer(runner, peer_idx, reason, arm, disarm): - """Fail-fast teardown of a not-yet-run ctx peer: open the session and tell - it to abort (it is blocked awaiting our hello), so it stops promptly - instead of waiting out the handshake alarm. Best-effort; always closes.""" + """Fail-fast teardown of a not-yet-run ctx peer. + + Open the session and tell it to abort (it is blocked awaiting our hello), + so it stops promptly instead of waiting out the handshake alarm. + Best-effort; always closes. + """ sock = None try: sock, key = _gen_open_session(runner, peer_idx, arm) @@ -1302,10 +1326,12 @@ def stop(): def _make_peer_failure_recorder(runner, disarm, current_cell): - """Exception -> verdict mapping shared by all per-peer loops. Recording a - failure also drops the fail-fast flag so the remaining peers (here and in - the other instances) are skipped instead of tested against a fabric - already known bad -- see _drive_ctx_peers / raise_abort_flag.""" + """Exception -> verdict mapping shared by all per-peer loops. + + Recording a failure also drops the fail-fast flag so the remaining peers + (here and in the other instances) are skipped instead of tested against a + fabric already known bad -- see _drive_ctx_peers / raise_abort_flag. + """ def record_peer_failure(peer, exc): disarm() @@ -1322,10 +1348,12 @@ def record_peer_failure(peer, exc): def _serve_gen_peers(runner, plan, arm, disarm, record_peer_failure): - """ctx role: bind one dedicated REP socket per gen peer (avoids REQ - interleaving across sessions on a shared socket), publish the addr files, - then serve each peer's full schedule. Each session gets a fresh HMAC key, - shared only through the work-dir addr file (0600).""" + """Ctx role: bind per-peer REP sockets, publish addrs, serve each schedule. + + One dedicated REP socket per gen peer avoids REQ interleaving across + sessions on a shared socket. Each session gets a fresh HMAC key, shared + only through the work-dir addr file (0600). + """ num_peers = runner.side["num_peers"] socks, keys = {}, {} if runner.is_leader: @@ -1361,7 +1389,7 @@ def _serve_gen_peers(runner, plan, arm, disarm, record_peer_failure): def _drive_ctx_peers(runner, arm, disarm, record_peer_failure): - """gen role: run every ctx peer's schedule, then release all sessions. + """Gen role: run every ctx peer's schedule, then release all sessions. Fail-fast: once any pair has failed (this instance or another -- signalled through the work-dir abort flag), the remaining ctx peers are not tested; @@ -1372,7 +1400,8 @@ def _drive_ctx_peers(runner, arm, disarm, record_peer_failure): The release ("done") is deferred until EVERY driven peer's schedule finished, so those ctx instances stay alive for the whole precheck -- matching real serving, where no transceiver ever holds connections to a - dead agent while transfers are still running.""" + dead agent while transfers are still running. + """ open_sessions = [] for ci in range(runner.side["num_peers"]): reason = abort_flag_reason(runner.work_dir) diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py index 2d63de210458..c2c77a349713 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_config.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py @@ -324,9 +324,7 @@ def test_write_addr_replaces_stale_file(self, tmp_path, monkeypatch): def test_wireup_timeout_derivation(): plan = pcfg.resolve_plan(_disagg_yaml()) # ctx dep4 -> gen dep16 assert plan["wireup_timeout_s"] == min(1800, 150 * 16) - plan = pcfg.resolve_plan( - _disagg_yaml(gen_extra={"tensor_parallel_size": 4}) - ) + plan = pcfg.resolve_plan(_disagg_yaml(gen_extra={"tensor_parallel_size": 4})) assert plan["wireup_timeout_s"] == 600 plan = pcfg.resolve_plan(_disagg_yaml(cache_transceiver_precheck={"wireup_timeout_s": 42})) assert plan["wireup_timeout_s"] == 42 @@ -338,9 +336,11 @@ def _enabled_line(cfg): def test_precheck_env_kill_switch_truthy(monkeypatch): - """The TRTLLM_DISAGG_CT_PRECHECK global kill switch parses the usual boolean - spellings (so a force-enable like =true is not silently read as "off") and - rejects anything ambiguous instead of guessing.""" + """The TRTLLM_DISAGG_CT_PRECHECK kill switch parses the usual boolean spellings. + + So a force-enable like =true is not silently read as "off", and anything + ambiguous is rejected instead of guessed at. + """ cfg = {"cache_transceiver_precheck": {"enabled": True}} monkeypatch.delenv("TRTLLM_DISAGG_CT_PRECHECK", raising=False) assert _enabled_line(cfg).endswith("=1") # yaml default @@ -359,8 +359,11 @@ def test_precheck_env_kill_switch_truthy(monkeypatch): def test_gate_library_content(tmp_path): - """The gate library loads from next to the draft, falls back to the in-repo - copy for an external draft, strips blank lines, and errors if truly absent.""" + """The gate library loads from next to the draft, with an in-repo fallback. + + It falls back to the in-repo copy for an external draft, strips blank + lines, and errors if truly absent. + """ dd = tmp_path / "disaggregated" dd.mkdir() (dd / "slurm_ct_precheck_gate.sh").write_text( @@ -383,8 +386,11 @@ def test_gate_library_content(tmp_path): def test_rid_tags_dense_within_session(): - """The C++ notification tag is rid & 0xFFF: rids must be dense within a - (ctx, gen) session so tags cannot alias across reps/lengths.""" + """Rids must be dense within a (ctx, gen) session. + + The C++ notification tag is rid & 0xFFF, so dense rids keep tags from + aliasing across reps/lengths. + """ plan = pcfg.resolve_plan(_disagg_yaml()) # n_pairs=16 total_reps = plan["warmup_requests"] + plan["num_requests"] n_pairs = plan["n_pairs"] @@ -463,9 +469,7 @@ def ctx_run_wave(peer_idx, li, req_len, rep, wave): runner.ctx_run_wave = ctx_run_wave runner.ctx_finish_wave = lambda reqs: None - runner.gen_run_wave = ( - lambda peer_idx, li, req_len, rep, wave, params: (True, "") - ) + runner.gen_run_wave = lambda peer_idx, li, req_len, rep, wave, params: (True, "") runner._calls = calls return runner @@ -494,7 +498,11 @@ def _run(self, tmp_path, monkeypatch, fail_ctx_idx=None): gen = self._mk_runner("gen", 0, plan, work, monkeypatch) ctxs = [ self._mk_runner( - "ctx", i, plan, work, monkeypatch, + "ctx", + i, + plan, + work, + monkeypatch, fail_ctx=(fail_ctx_idx is not None and i == fail_ctx_idx), ) for i in range(2) @@ -531,9 +539,7 @@ def test_two_ctx_full_pass(self, tmp_path, monkeypatch): } # every ctx served the full schedule (reps x waves) and got its # deferred done (PASS recorded only after done/bye completes) - total_waves = len(pcfg.waves(plan)) * ( - plan["warmup_requests"] + plan["num_requests"] - ) + total_waves = len(pcfg.waves(plan)) * (plan["warmup_requests"] + plan["num_requests"]) for c in ctxs: assert c._calls["waves"] == total_waves assert [x["status"] for x in c.recorder.cases] == ["PASS"] @@ -578,11 +584,13 @@ def test_abort_flag_stale_job_ignored(self, tmp_path, monkeypatch): def test_ctx_run_wave_missing_params_broadcast(tmp_path, monkeypatch): - """#4 regression: the "missing context_phase_params" verdict is computed - only on the instance leader (only it holds the gathered params), so it must - be broadcast -- otherwise a NON-leader rank keeps reason=None, returns, and - enters the next collective while the leader raises, deadlocking the step - until the watchdog SIGKILLs it (misreported as TIMEOUT). + """#4 regression: the "missing context_phase_params" verdict must be broadcast. + + It is computed only on the instance leader (only it holds the gathered + params) -- without the broadcast, a NON-leader rank keeps reason=None, + returns, and enters the next collective while the leader raises, + deadlocking the step until the watchdog SIGKILLs it (misreported as + TIMEOUT). Run the real ctx_run_wave on a non-leader rank: with the leader's verdict delivered via bcast the rank must raise; with a clean (None) broadcast it @@ -650,9 +658,12 @@ def _mk(bcast_ret): def test_status_env_snapshot_excludes_nixl(tmp_path, monkeypatch): - """NIXL_* is not captured: the only such variable seen in practice is - NIXL_VERSION, a stale NGC-base-image marker that misstates the version of - the actually-linked library.""" + """NIXL_* is not captured. + + The only such variable seen in practice is NIXL_VERSION, a stale + NGC-base-image marker that misstates the version of the actually-linked + library. + """ monkeypatch.setenv("NIXL_VERSION", "1.0.0") monkeypatch.setenv("NIXL_PLUGIN_DIR", "/opt/x") monkeypatch.setenv("UCX_TLS", "rc,cuda_copy") @@ -662,10 +673,12 @@ def test_status_env_snapshot_excludes_nixl(tmp_path, monkeypatch): def test_sparse_attention_model_uses_simplified_mla_pool(tmp_path): - """DeepSeek V4 / DSA can't be modeled as a single KV pool. The precheck is - a NETWORK check, so it falls back to a simple MLA-flavored stand-in pool - (real layer count, one latent head, is_mla=True) and still runs -- never - skips.""" + """DeepSeek V4 / DSA can't be modeled as a single KV pool. + + The precheck is a NETWORK check, so it falls back to a simple MLA-flavored + stand-in pool (real layer count, one latent head, is_mla=True) and still + runs -- never skips. + """ d = tmp_path / "v4" d.mkdir() (d / "config.json").write_text( @@ -699,8 +712,11 @@ def test_sparse_attention_model_uses_simplified_mla_pool(tmp_path): def test_python_transceiver_bandwidth_csv(tmp_path): - """Bandwidth from the Python transceiver's perf_logger CSVs: median over - KVSendTask throughput_mbs (MB/s -> GB/s), receiver rows ignored.""" + """Bandwidth from the Python transceiver's perf_logger CSVs. + + Median over KVSendTask throughput_mbs (MB/s -> GB/s), receiver rows + ignored. + """ header = ( "timestamp,task_type,unique_rid,peer_rank,transfer_size_bytes," "avg_segment_size_bytes,transfer_entry_count,prepare_args_latency_ms," @@ -709,12 +725,11 @@ def test_python_transceiver_bandwidth_csv(tmp_path): # two ctx ranks, each its own perf file; KVSendTask rows carry throughput (tmp_path / "perf_abc_0.csv").write_text( header + "\n" - "t,KVSendTask,1,0,1000,,,0,0,0,0,102400.00\n" # 100 GB/s - "t,AuxSendTask,1,0,10,,,0,0,0,0,1.00\n" # tiny metadata, ignored + "t,KVSendTask,1,0,1000,,,0,0,0,0,102400.00\n" # 100 GB/s + "t,AuxSendTask,1,0,10,,,0,0,0,0,1.00\n" # tiny metadata, ignored ) (tmp_path / "perf_abc_1.csv").write_text( - header + "\n" - "t,KVSendTask,2,1,1000,,,0,0,0,0,204800.00\n" # 200 GB/s + header + "\nt,KVSendTask,2,1,1000,,,0,0,0,0,204800.00\n" # 200 GB/s ) # a receiver file (no throughput) must not contribute (tmp_path / "perf_def_8.csv").write_text(header + "\nt,KVRecvTask,3,0,,,,,,,5.0,\n") From 12ffec3c16fae8a208f54115603025767b9d159c Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Thu, 23 Jul 2026 04:09:52 +0000 Subject: [PATCH 17/21] [None][infra] Print the exact test-results artifact URL on Slurm failure The failure echo used a '*' glob in place of the attempt suffix, so the printed Artifactory URL was not directly downloadable. The tracking script is regenerated per attempt with postTag in scope, and the teardown upload uses the same postTag, so interpolating it yields the exact uploaded filename. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- jenkins/L0_Test.groovy | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 087c9b99e741..5cdf61ebe6e1 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -2044,7 +2044,7 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG else echo "Pytest failed in Slurm job \$jobId" echo "Full test output (logs not shown above) is uploaded after stage teardown to:" - echo " https://urm.nvidia.com/artifactory/${UPLOAD_PATH}/test-results/results-${stageName}*.tar.gz" + echo " https://urm.nvidia.com/artifactory/${UPLOAD_PATH}/test-results/results-${stageName}${postTag}.tar.gz" fi echo "Status: \$STATUS | Exit_code \$EXIT_CODE" exit 0 From 8af0ac887d45bef359fbd58b96e25849f83b0662 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Thu, 23 Jul 2026 04:23:14 +0000 Subject: [PATCH 18/21] [None][fix] Address review: MTP depth fallback, abort-flag cleanup and collective read - precheck_config: _spec_nextn falls back to max_draft_len (the field num_nextn_predict_layers is its deprecated alias); several disagg yamls spell the MTP depth only as max_draft_len and resolved to 0 spec layers. - gate script: also remove a stale precheck.abort before launching -- a requeued job keeps its SLURM_JOB_ID, so the driver's job-id stamp cannot reject it and the rerun would be falsely fail-fast-skipped. - run_precheck: read the fail-fast flag collectively (leader reads, bcast) so ranks cannot race the flag write into divergent branches whose MPI collective sequences differ (gen_run_peer vs gen_abort_peer). Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../disaggregated/slurm_ct_precheck_gate.sh | 12 ++++---- .../precheck_config.py | 13 ++++++--- .../run_precheck.py | 21 ++++++++++++-- .../test_cache_transceiver_precheck_config.py | 29 +++++++++++++++++++ 4 files changed, 63 insertions(+), 12 deletions(-) diff --git a/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh b/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh index 9a888ced5b94..e3b02f88741d 100644 --- a/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh +++ b/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh @@ -112,14 +112,16 @@ run_cache_transceiver_precheck() { precheckDir="$testOutputDir/cache_transceiver_precheck" mkdir -p "$precheckDir/logs" # A reused work dir (Slurm requeue reruns this batch script with the same - # directories) may hold a previous run's rendezvous/status/csv files: stale - # addr files would point gen leaders at dead ports, stale status files would - # pollute the verdict aggregation below, and stale bandwidth CSVs (the - # Python transceiver's perf__.csv are per-run and appended) + # directories) may hold a previous run's rendezvous/status/csv/abort files: + # stale addr files would point gen leaders at dead ports, stale status files + # would pollute the verdict aggregation below, a stale precheck.abort would + # fail-fast-skip the whole rerun (a requeued job keeps its SLURM_JOB_ID, so + # the driver's job-id stamp cannot tell it apart), and stale bandwidth CSVs + # (the Python transceiver's perf__.csv are per-run and appended) # would make parse_python_bandwidth_gbps median over two runs' samples. The # driver also job-id-stamps addr files as a second line of defense. rm -f "$precheckDir"/rendezvous/*.addr "$precheckDir"/status/*.status \ - "$precheckDir"/status/*.json 2>/dev/null || true + "$precheckDir"/status/*.json "$precheckDir"/precheck.abort 2>/dev/null || true rm -rf "$precheckDir"/csv 2>/dev/null || true precheckPids=() precheckNames=() diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py index c02e24a20fbc..916f2c4678a3 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py @@ -101,10 +101,15 @@ def _parallel(side): def _spec_nextn(side): spec = side.get("speculative_config") or {} - decoding_type = str(spec.get("decoding_type", "")).upper() - if decoding_type == "MTP": - return int(spec.get("num_nextn_predict_layers", 0) or 0) - return 0 + if str(spec.get("decoding_type", "")).upper() != "MTP": + return 0 + # Test yamls spell the MTP depth either way: num_nextn_predict_layers is + # the deprecated alias MTPDecodingConfig remaps to max_draft_len, so both + # must resolve to the same KV-pool spec-layer count here. + n = spec.get("num_nextn_predict_layers") + if n is None: + n = spec.get("max_draft_len", 0) + return int(n or 0) def wireup_timeout_s(max_world): diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index 9dad2594b9e1..853659ab1ee1 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -1347,6 +1347,19 @@ def record_peer_failure(peer, exc): return record_peer_failure +def _consensus_abort_reason(runner): + """Instance-wide agreed view of the fail-fast flag (leader reads, bcast). + + The flag file can appear at any moment (any instance's failure drops it), + so per-rank reads can race it and disagree -- and the branches they select + (gen_run_peer vs gen_abort_peer, failure vs SKIP verdict) issue different + MPI collective sequences, deadlocking or cross-pairing the instance. Must + be called collectively by every rank. + """ + reason = abort_flag_reason(runner.work_dir) if runner.is_leader else None + return runner.comm.bcast(reason, root=0) + + def _serve_gen_peers(runner, plan, arm, disarm, record_peer_failure): """Ctx role: bind per-peer REP sockets, publish addrs, serve each schedule. @@ -1379,8 +1392,10 @@ def _serve_gen_peers(runner, plan, arm, disarm, record_peer_failure): # A gen driver that failed elsewhere aborts our session as part of # fail-fast: record a (non-failing) SKIP, not our own failure -- # the real failure is recorded by whoever hit it. Absent the flag, - # a genuine peer abort is still a real failure. - if abort_flag_reason(runner.work_dir) is not None: + # a genuine peer abort is still a real failure. The consensus read + # is collectively safe here: _PeerAbort is only raised after a + # bcast (leader_recv), so every rank reaches this handler. + if _consensus_abort_reason(runner) is not None: runner.recorder.record(f"gen_{gj}", 0, "SKIP", f"aborted by fail-fast: {e}") else: record_peer_failure(f"gen_{gj}", e) @@ -1404,7 +1419,7 @@ def _drive_ctx_peers(runner, arm, disarm, record_peer_failure): """ open_sessions = [] for ci in range(runner.side["num_peers"]): - reason = abort_flag_reason(runner.work_dir) + reason = _consensus_abort_reason(runner) if reason is not None: try: gen_abort_peer(runner, ci, reason, arm, disarm) diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py index c2c77a349713..9212f17bfc05 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_config.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py @@ -91,6 +91,35 @@ def test_resolve_plan_adp_asymmetric(): assert plan["ctx_cache_transceiver_config"]["backend"] == "NIXL" +def test_spec_nextn_max_draft_len_fallback(): + # Checked-in yamls spell the MTP depth either way (num_nextn_predict_layers + # is MTPDecodingConfig's deprecated alias of max_draft_len); both must + # resolve to the same spec-layer count. + cfg = _disagg_yaml( + ctx_extra={"speculative_config": {"decoding_type": "MTP", "max_draft_len": 3}}, + gen_extra={ + "speculative_config": { + "decoding_type": "MTP", + "num_nextn_predict_layers": 3, + "max_draft_len": 3, + } + }, + ) + plan = pcfg.resolve_plan(cfg) + assert plan["ctx_num_nextn_predict_layers"] == 3 + assert plan["gen_num_nextn_predict_layers"] == 3 + + # Non-MTP speculation never contributes MTP KV layers, even with a + # max_draft_len present. + cfg = _disagg_yaml( + ctx_extra={"speculative_config": {"decoding_type": "Eagle", "max_draft_len": 3}}, + gen_extra={"speculative_config": {"decoding_type": "Eagle", "max_draft_len": 3}}, + ) + plan = pcfg.resolve_plan(cfg) + assert plan["ctx_num_nextn_predict_layers"] == 0 + assert plan["gen_num_nextn_predict_layers"] == 0 + + def test_request_lengths_clamped_by_buffer_and_cap(): cfg = _disagg_yaml(benchmark={"mode": "e2e", "input_length": 131072}) cfg["worker_config"]["ctx"]["cache_transceiver_config"]["max_tokens_in_buffer"] = 131104 From 14d0a675e3435324d8f73d13bbc1b97dc30cbe72 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Fri, 24 Jul 2026 04:05:02 +0000 Subject: [PATCH 19/21] [None][test] Add unit, contract and single-GPU e2e coverage for the cache-transceiver precheck driver The precheck driver (run_precheck.py) had no test coverage and reached into unstable TRT-LLM internals from eight scattered call sites, so an upstream refactor would only surface as INIT_ERROR aborts on the SLURM disagg perf pipeline. - run_precheck.py: centralize every tensorrt_llm import behind a lazy load_internal_apis() (single failure point; --dry-run and pure-logic tests still work without torch installed). - test_cache_transceiver_precheck_run.py: pure-logic tests (HMAC wire format, rid/seed scheme, rendezvous/abort-flag files, StatusRecorder, bandwidth CSV parsing, schedule/timeouts) plus internal-API contract tests that pin the constructor/signature shapes the driver relies on, so internal-API drift fails in the refactorer's pre-merge CI. - test_cache_transceiver_precheck_e2e.py: drives run_precheck.py like the SLURM gate (one mpirun world per ctx/gen instance, shared work-dir, synthetic model config.json) on one GPU: symmetric 1x1, multi-instance 2x2, asymmetric tp2->tp1, and a fingerprint-mismatch fail-fast negative path asserting FAIL verdicts instead of hangs. - Wire the new tests into l0_a10 (unit/contract) and l0_h100 (e2e, next to the existing single-node transceiver harness test). Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../test_lists/test-db/l0_h100.yml | 1 + .../cache_transceiver_precheck/README.md | 12 +- .../run_precheck.py | 186 +++++--- .../test_cache_transceiver_precheck_e2e.py | 329 ++++++++++++++ .../test_cache_transceiver_precheck_run.py | 419 ++++++++++++++++++ 6 files changed, 884 insertions(+), 64 deletions(-) create mode 100644 tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py create mode 100644 tests/unittest/others/test_cache_transceiver_precheck_run.py diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index c0ce6be1abf3..aee432e33528 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -56,6 +56,7 @@ l0_a10: - unittest/inputs/test_multimodal_input_processor.py - unittest/inputs/test_video_decode.py - unittest/others/test_cache_transceiver_precheck_config.py + - unittest/others/test_cache_transceiver_precheck_run.py - unittest/others/test_convert_utils.py - unittest/others/test_lora_manager.py - unittest/others/test_lora_module_count.py diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 15048d2b8bc6..f4a3fcefef03 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -110,6 +110,7 @@ l0_h100: - unittest/disaggregated/test_cache_transceiver_single_process.py::test_cache_transceiver_v1_dsa_indexer - unittest/disaggregated/test_cache_transceiver_harness_report.py - unittest/disaggregated/test_cache_transceiver_harness.py + - unittest/disaggregated/test_cache_transceiver_precheck_e2e.py - unittest/others/test_kv_cache_transceiver.py::test_kv_cache_transceiver_single_process[PYTHON-mha-ctx_fp16_gen_fp16] - unittest/others/test_kv_cache_transceiver.py::test_kv_cache_transceiver_single_process[PYTHON-mla-ctx_fp16_gen_fp16] - unittest/llmapi/test_llm_telemetry.py diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md index 715c66b18cfd..9ec289a4f230 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md @@ -86,7 +86,17 @@ srun -N2 --ntasks=8 --mpi=pmix python3 run_precheck.py --role gen --server-idx 0 wait ``` -Unit tests (CPU-only): `pytest tests/unittest/others/test_cache_transceiver_precheck_config.py`. +Tests: + +- Config resolution (CPU-only): + `pytest tests/unittest/others/test_cache_transceiver_precheck_config.py` +- Driver logic + internal-API contract (the precheck drives unstable + TRT-LLM internals via `run_precheck.load_internal_apis()`; the contract + tests catch upstream renames in pre-merge CI): + `pytest tests/unittest/others/test_cache_transceiver_precheck_run.py` +- End-to-end on one GPU (one `mpirun` world per ctx/gen instance, real + NIXL/PYTHON transfer + verdicts, incl. multi-instance and asymmetric TP): + `pytest tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py` Lineage: adapted from `examples/disaggregated/slurm/cache_transceiver_test` (the UCX tuning harness), reduced to a single sweep and extended to diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index 853659ab1ee1..85f769167566 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -156,16 +156,99 @@ def params_from_wire(d): Uses the maintained DisaggregatedParams converter (keeps us off the raw nanobind ctor signature). """ + return ( + load_internal_apis() + .DisaggregatedParams( + ctx_request_id=int(d["req_id"]), + first_gen_tokens=list(d["first_gen_tokens"]), + opaque_state=base64.b64decode(d["opaque_state"]), + draft_tokens=d["draft_tokens"], + ctx_dp_rank=d["ctx_dp_rank"], + ctx_info_endpoint=d["ctx_info_endpoint"], + ) + .get_context_phase_params() + ) + + +# --------------------------------------------------------------------------- # +# TRT-LLM internal API surface (single owner) +# --------------------------------------------------------------------------- # +_INTERNAL_APIS = None + + +def load_internal_apis(): + """Every tensorrt_llm symbol the precheck touches, imported in ONE place. + + The precheck bypasses the serving stack and drives internal APIs directly + (_torch.pyexecutor.*, bindings.internal.*, private llm_utils resolvers), + none of which carry a stability promise. Centralizing the imports means an + upstream rename breaks exactly here, and the contract test + (tests/unittest/others/test_cache_transceiver_precheck_run.py) fails in + the refactorer's pre-merge CI instead of aborting the SLURM disagg perf + pipeline at runtime. + + Deliberately lazy (NOT at module import): --dry-run and the pure-logic + unit tests must work without torch / tensorrt_llm installed. + """ + global _INTERNAL_APIS + if _INTERNAL_APIS is not None: + return _INTERNAL_APIS + import types + + import tensorrt_llm + import tensorrt_llm._torch.models # noqa: F401 - populates the model registry + import tensorrt_llm.bindings + import tensorrt_llm.bindings.executor as trtllm_executor from tensorrt_llm import DisaggregatedParams + from tensorrt_llm._torch.distributed import Distributed + from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPING + from tensorrt_llm._torch.pyexecutor.hang_detector import HangDetector + from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 + from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import create_kv_cache_transceiver + from tensorrt_llm._torch.pyexecutor.llm_request import ( + LlmRequest, + LlmRequestState, + LlmRequestType, + ) + from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager + from tensorrt_llm.llmapi.llm_args import ( + CacheTransceiverConfig, + KvCacheConfig, + MTPDecodingConfig, + ) + from tensorrt_llm.llmapi.llm_utils import ( + _resolve_kv_cache_manager_v2_auto, + _resolve_transceiver_runtime_auto, + ) + from tensorrt_llm.mapping import Mapping + from tensorrt_llm.sampling_params import SamplingParams - return DisaggregatedParams( - ctx_request_id=int(d["req_id"]), - first_gen_tokens=list(d["first_gen_tokens"]), - opaque_state=base64.b64decode(d["opaque_state"]), - draft_tokens=d["draft_tokens"], - ctx_dp_rank=d["ctx_dp_rank"], - ctx_info_endpoint=d["ctx_info_endpoint"], - ).get_context_phase_params() + _INTERNAL_APIS = types.SimpleNamespace( + tensorrt_llm=tensorrt_llm, + DataType=tensorrt_llm.bindings.DataType, + SamplingConfigCpp=tensorrt_llm.bindings.SamplingConfig, + CacheTypeCpp=tensorrt_llm.bindings.internal.batch_manager.CacheType, + AttentionTypeCpp=tensorrt_llm.bindings.internal.batch_manager.AttentionType, + KvCacheConfigCpp=trtllm_executor.KvCacheConfig, + DisaggregatedParams=DisaggregatedParams, + Distributed=Distributed, + MODEL_CLASS_MAPPING=MODEL_CLASS_MAPPING, + HangDetector=HangDetector, + KVCacheManager=KVCacheManager, + KVCacheManagerV2=KVCacheManagerV2, + create_kv_cache_transceiver=create_kv_cache_transceiver, + LlmRequest=LlmRequest, + LlmRequestState=LlmRequestState, + LlmRequestType=LlmRequestType, + CacheTransceiverConfig=CacheTransceiverConfig, + KvCacheConfig=KvCacheConfig, + MTPDecodingConfig=MTPDecodingConfig, + resolve_kv_cache_manager_v2_auto=_resolve_kv_cache_manager_v2_auto, + resolve_transceiver_runtime_auto=_resolve_transceiver_runtime_auto, + Mapping=Mapping, + SamplingParams=SamplingParams, + ) + return _INTERNAL_APIS # --------------------------------------------------------------------------- # @@ -231,10 +314,7 @@ def _lookup_model_cls(model_dir): hf_view = type("HFConfigView", (), hf_cfg) # attribute access for the pref hook if not archs: return None, hf_view - import tensorrt_llm._torch.models # noqa: F401 - populates the registry - from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPING - - return MODEL_CLASS_MAPPING.get(archs[0]), hf_view + return load_internal_apis().MODEL_CLASS_MAPPING.get(archs[0]), hf_view def resolve_model_prefs(model_dir, side, cache_cfg): @@ -251,6 +331,7 @@ def resolve_model_prefs(model_dir, side, cache_cfg): """ import types + api = load_internal_apis() model_cls, hf_view = _lookup_model_cls(model_dir) setting = side["use_kv_cache_manager_v2"] @@ -267,12 +348,10 @@ def resolve_model_prefs(model_dir, side, cache_cfg): try: # The REAL serving resolver, via the same shim pattern as the # runtime resolution below -- one owner for the 'auto' semantics. - from tensorrt_llm.llmapi.llm_utils import _resolve_kv_cache_manager_v2_auto - shim = types.SimpleNamespace( kv_cache_config=types.SimpleNamespace(use_kv_cache_manager_v2="auto") ) - use_v2 = bool(_resolve_kv_cache_manager_v2_auto(shim, defaults)) + use_v2 = bool(api.resolve_kv_cache_manager_v2_auto(shim, defaults)) except Exception as e: # noqa: BLE001 - fall back like a missing model print( f"[precheck] WARNING: V2 'auto' resolution failed ({e!r}); assuming V1", flush=True @@ -283,10 +362,8 @@ def resolve_model_prefs(model_dir, side, cache_cfg): if getattr(cache_cfg, "transceiver_runtime", None) == "auto": try: - from tensorrt_llm.llmapi.llm_utils import _resolve_transceiver_runtime_auto - shim = types.SimpleNamespace(cache_transceiver_config=cache_cfg) - _resolve_transceiver_runtime_auto(shim, model_cls, hf_view) + api.resolve_transceiver_runtime_auto(shim, model_cls, hf_view) except Exception as e: # noqa: BLE001 - fall back to the create() default (CPP) print( f"[precheck] WARNING: transceiver_runtime 'auto' resolution failed " @@ -297,29 +374,24 @@ def resolve_model_prefs(model_dir, side, cache_cfg): def build_kv_cache_manager(kv_shape, plan, side, mapping, max_req_len, use_v2): - import tensorrt_llm.bindings - import tensorrt_llm.bindings.executor as trtllm_executor - from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager - - DataType = tensorrt_llm.bindings.DataType - CacheTypeCpp = tensorrt_llm.bindings.internal.batch_manager.CacheType + api = load_internal_apis() dtype_map = { - "fp8": DataType.FP8, - "fp16": DataType.HALF, - "half": DataType.HALF, - "bf16": DataType.BF16, + "fp8": api.DataType.FP8, + "fp16": api.DataType.HALF, + "half": api.DataType.HALF, + "bf16": api.DataType.BF16, } dtype_str = side["kv_dtype"].lower() dtype = dtype_map.get(dtype_str) if dtype is None: print(f"[precheck] kv dtype {dtype_str!r} not mapped, using BF16", flush=True) - dtype = DataType.BF16 + dtype = api.DataType.BF16 spec_config = None if side["num_nextn_predict_layers"] > 0: - from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig - - spec_config = MTPDecodingConfig(num_nextn_predict_layers=side["num_nextn_predict_layers"]) + spec_config = api.MTPDecodingConfig( + num_nextn_predict_layers=side["num_nextn_predict_layers"] + ) tpb = plan["tokens_per_block"] padded_len = ((max_req_len + tpb - 1) // tpb) * tpb @@ -328,7 +400,7 @@ def build_kv_cache_manager(kv_shape, plan, side, mapping, max_req_len, use_v2): # Real MLA serving uses SELFKONLY (kv_factor=1: one latent plane, no V) — # see _torch/pyexecutor/_util.py; SELF would double the per-token bytes. - cache_type = CacheTypeCpp.SELFKONLY if kv_shape["is_mla"] else CacheTypeCpp.SELF + cache_type = api.CacheTypeCpp.SELFKONLY if kv_shape["is_mla"] else api.CacheTypeCpp.SELF common = dict( num_layers=kv_shape["num_layers"], num_kv_heads=kv_shape["num_kv_heads"], @@ -341,15 +413,12 @@ def build_kv_cache_manager(kv_shape, plan, side, mapping, max_req_len, use_v2): spec_config=spec_config, ) if use_v2: - from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 - from tensorrt_llm.llmapi.llm_args import KvCacheConfig - # The REAL pydantic KvCacheConfig (what serving passes): partial reuse # is pinned off because its pydantic default is True and block reuse # is off here. is_disagg=True doubles the IndexMapper capacity so # in-flight transfers (TRANS_IN_PROGRESS) can hold slots. - return KVCacheManagerV2( - KvCacheConfig( + return api.KVCacheManagerV2( + api.KvCacheConfig( max_tokens=max_tokens, enable_block_reuse=False, enable_partial_reuse=False, @@ -361,8 +430,8 @@ def build_kv_cache_manager(kv_shape, plan, side, mapping, max_req_len, use_v2): is_disagg=True, **common, ) - return KVCacheManager( - trtllm_executor.KvCacheConfig(max_tokens=max_tokens, enable_block_reuse=False), + return api.KVCacheManager( + api.KvCacheConfigCpp(max_tokens=max_tokens, enable_block_reuse=False), cache_type, **common, ) @@ -370,17 +439,16 @@ def build_kv_cache_manager(kv_shape, plan, side, mapping, max_req_len, use_v2): def make_request(is_ctx, rid, req_len, runtime, ctx_params=None): """Build a ctx or gen LlmRequest (mirrors the UCX-tuning harness).""" - import tensorrt_llm.bindings - from tensorrt_llm import DisaggregatedParams - from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestType - from tensorrt_llm.sampling_params import SamplingParams + api = load_internal_apis() + LlmRequest, LlmRequestType = api.LlmRequest, api.LlmRequestType + DisaggregatedParams = api.DisaggregatedParams - sampling = SamplingParams() + sampling = api.SamplingParams() common = dict( request_id=rid, max_new_tokens=1, input_tokens=list(range(req_len)), - sampling_config=tensorrt_llm.bindings.SamplingConfig(sampling._get_sampling_config()), + sampling_config=api.SamplingConfigCpp(sampling._get_sampling_config()), is_streaming=False, ) if is_ctx: @@ -734,17 +802,11 @@ def _consensus_error(self, local_err): # ---- setup ------------------------------------------------------------- def setup(self, kv_shape, max_req_len): - import tensorrt_llm - import tensorrt_llm.bindings - from tensorrt_llm._torch.distributed import Distributed - from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import create_kv_cache_transceiver - from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState - from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig - from tensorrt_llm.mapping import Mapping - - self.llm_request_state = LlmRequestState + api = load_internal_apis() + + self.llm_request_state = api.LlmRequestState par = self.side["parallel"] - self.mapping = Mapping( + self.mapping = api.Mapping( world_size=par["world_size"], rank=self.rank, gpus_per_node=self.plan["gpus_per_node"], @@ -764,7 +826,7 @@ def setup(self, kv_shape, max_req_len): # Built VERBATIM from the disagg yaml's cache_transceiver_config so # backend/max_tokens_in_buffer/timeouts match the real test exactly. - cache_cfg = CacheTransceiverConfig(**self.side["cache_transceiver_config"]) + cache_cfg = api.CacheTransceiverConfig(**self.side["cache_transceiver_config"]) # Yaml-absent settings resolve against the model's preferences, like # serving does (kv manager version + transceiver runtime) -- this holds # even for the simplified stand-in pool: only the KV SHAPE is generic; @@ -789,10 +851,10 @@ def setup(self, kv_shape, max_req_len): self.kvm = build_kv_cache_manager( kv_shape, self.plan, self.side, self.mapping, max_req_len, self.use_v2 ) - AttentionTypeCpp = tensorrt_llm.bindings.internal.batch_manager.AttentionType + AttentionTypeCpp = api.AttentionTypeCpp attention_type = AttentionTypeCpp.MLA if kv_shape["is_mla"] else AttentionTypeCpp.DEFAULT - dist_obj = Distributed.get(self.mapping) - self.xcvr = create_kv_cache_transceiver( + dist_obj = api.Distributed.get(self.mapping) + self.xcvr = api.create_kv_cache_transceiver( self.mapping, dist_obj, self.kvm, attention_type, cache_cfg ) if self.xcvr is None: @@ -1285,8 +1347,6 @@ def _install_watchdog(runner, plan, rank): final shutdown, and the mutable "what phase are we in" marker used by failure messages. """ - from tensorrt_llm._torch.pyexecutor.hang_detector import HangDetector - signal.signal(signal.SIGALRM, _alarm_handler) current_cell = {"what": "startup"} @@ -1303,7 +1363,7 @@ def _on_hang(): # The detector must outlast the LONGEST legitimate wait (peer handshakes # are serialized across sessions); per-cell alarms are the tighter bound # for actual transfer work. - hang_detector = HangDetector( + hang_detector = load_internal_apis().HangDetector( timeout=hello_timeout_s(plan, runner.side["num_peers"]) + plan["wave_timeout_s"] + 60, on_detected=_on_hang, ) diff --git a/tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py b/tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py new file mode 100644 index 000000000000..7abde75cd12e --- /dev/null +++ b/tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py @@ -0,0 +1,329 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""End-to-end test for the disagg perf-sanity cache-transceiver PRECHECK driver. + +Drives ``run_precheck.py`` (tests/scripts/perf-sanity/cache_transceiver_precheck) +exactly like the SLURM gate does — one ``mpirun`` world per ctx/gen server +instance, a shared ``--work-dir`` for rendezvous/status, the disagg yaml as +the single config source — but on ONE node with every process sharing one +physical GPU. This is the CI net for the precheck's NON-network failure +modes: internal TRT-LLM API drift, KV pool construction, transceiver +setup/transfer/verification, multi-instance pairing, and failure verdicts. + +Requires: 1 GPU, mpirun, mpi4py, tensorrt_llm. Modeled on +tests/unittest/disaggregated/test_cache_transceiver_harness.py (same +single-node NIXL/PYTHON transceiver combination that stage already runs). +""" + +import json +import os +import shutil +import signal +import subprocess +import sys +import time + +import pytest +import yaml + +# Same single-GPU environment pinning as kv_transfer_harness.py: +# - One NIXL worker thread per agent: the default (8) causes heavy +# contention when many agents share a single GPU. +# - ``^ib,gdr_copy`` disables InfiniBand and GDR copy, which are +# unavailable (and flaky) on single-node loopback. +# Set here (not just in the child env) so any in-process tensorrt_llm +# import sees them too; the mpirun children inherit via os.environ.copy(). +os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" +os.environ["UCX_TLS"] = "^ib,gdr_copy" + +PRECHECK_DIR = os.path.normpath( + os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + os.pardir, + "tests", + "scripts", + "perf-sanity", + "cache_transceiver_precheck", + ) +) +DRIVER_SCRIPT = os.path.join(PRECHECK_DIR, "run_precheck.py") +_LOG_TAIL_CHARS = 16 * 1024 +_PROCESS_TIMEOUT_SECONDS = 240 +_TERMINATE_GRACE_SECONDS = 5 + +TINY_MODEL_CONFIG = { + # Llama-shaped so model_kv_shape() exercises the real config.json path + # (2 layers x 1 kv head x 64 dim: a few hundred KB of KV per process). + "architectures": ["LlamaForCausalLM"], + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 64, + "hidden_size": 256, + "vocab_size": 32000, +} + + +def _find_mpirun(): + path = shutil.which("mpirun") + if path is None: + pytest.skip("mpirun not found on PATH") + return path + + +def _log_tail(path): + try: + with open(path, errors="replace") as f: + log = f.read() + except OSError: + return f"(no log at {path})" + if len(log) <= _LOG_TAIL_CHARS: + return log + return ( + f"... {len(log) - _LOG_TAIL_CHARS} earlier characters omitted ...\n{log[-_LOG_TAIL_CHARS:]}" + ) + + +def _terminate_process_groups(processes): + for sig in (signal.SIGTERM, signal.SIGKILL): + # ALWAYS signal the group, even when the mpirun leader already + # exited: its ranks stay in the group and would otherwise leak GPU + # memory into the following tests (same contract as the harness's + # test_terminate_process_groups_signals_group_after_leader_exit). + for proc in processes: + try: + os.killpg(proc.pid, sig) + except ProcessLookupError: + pass + deadline = time.monotonic() + _TERMINATE_GRACE_SECONDS + for proc in processes: + if proc.poll() is None: + try: + proc.wait(timeout=max(0.0, deadline - time.monotonic())) + except subprocess.TimeoutExpired: + pass + + +def _disagg_yaml(num_ctx, num_gen, ctx_tp, gen_tp, request_lengths=(64,)): + """Minimal disagg perf-sanity yaml shaped like the checked-in configs.""" + tokens_per_block = 32 + + def side(tp): + return { + "tensor_parallel_size": tp, + "pipeline_parallel_size": 1, + "kv_cache_config": { + "dtype": "bf16", + "tokens_per_block": tokens_per_block, + # Same combination the merged single-node harness test runs: + # NIXL + PYTHON transceiver requires the V2 manager. + "use_kv_cache_manager_v2": True, + }, + "cache_transceiver_config": { + "backend": "NIXL", + "transceiver_runtime": "PYTHON", + "max_tokens_in_buffer": 512, + }, + } + + return { + "metadata": {"model_dir_name": "tiny-llama"}, + "benchmark": {"mode": "e2e", "input_length": 64, "output_length": 8}, + "hardware": { + "gpus_per_node": 1, + "num_ctx_servers": num_ctx, + "num_gen_servers": num_gen, + }, + "worker_config": {"ctx": side(ctx_tp), "gen": side(gen_tp)}, + "cache_transceiver_precheck": { + "request_lengths": list(request_lengths), + "num_requests": 1, + "warmup_requests": 1, + "wave_timeout_s": 60, + "wireup_timeout_s": 30, + "rendezvous_timeout_s": 90, + }, + } + + +def _write_inputs(tmp_path, cfg, name="precheck"): + models_root = tmp_path / "models" + model_dir = models_root / "tiny-llama" + model_dir.mkdir(parents=True, exist_ok=True) + (model_dir / "config.json").write_text(json.dumps(TINY_MODEL_CONFIG)) + config_path = tmp_path / f"{name}.yaml" + config_path.write_text(yaml.safe_dump(cfg)) + return str(config_path), str(models_root) + + +def _launch_instances(tmp_path, jobs, models_root): + """One mpirun per (role, idx, world, config) job; returns [(name, proc, log)].""" + mpirun = _find_mpirun() + work_dir = str(tmp_path / "work") + log_dir = tmp_path / "logs" + log_dir.mkdir(exist_ok=True) + + env = os.environ.copy() # carries the module-level UCX/NIXL pinning + env["CUDA_VISIBLE_DEVICES"] = env.get("CUDA_VISIBLE_DEVICES", "0").split(",")[0] + # run_precheck resolves metadata.model_dir_name under LLM_MODELS_ROOT. + env["LLM_MODELS_ROOT"] = models_root + + launched = [] + try: + for role, idx, world, config_path in jobs: + _launch_one(launched, mpirun, role, idx, world, config_path, work_dir, log_dir, env) + except BaseException: + # A failed Popen mid-list must not leak the instances already + # started: they never reach _wait_all's cleanup. + _terminate_process_groups([p for _, p, _ in launched]) + raise + return work_dir, launched + + +def _launch_one(launched, mpirun, role, idx, world, config_path, work_dir, log_dir, env): + name = f"{role}_{idx}" + log_path = str(log_dir / f"{name}.log") + cmd = [ + mpirun, + "--allow-run-as-root", + "--oversubscribe", + "-np", + str(world), + sys.executable, + DRIVER_SCRIPT, + "--role", + role, + "--server-idx", + str(idx), + "--config", + config_path, + "--work-dir", + work_dir, + ] + with open(log_path, "wb") as log_file: + proc = subprocess.Popen( + cmd, + env=env, + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + launched.append((name, proc, log_path)) + + +def _wait_all(launched): + procs = [p for _, p, _ in launched] + deadline = time.monotonic() + _PROCESS_TIMEOUT_SECONDS + try: + for name, proc, log_path in launched: + try: + proc.wait(timeout=max(0.0, deadline - time.monotonic())) + except subprocess.TimeoutExpired: + tails = "\n".join(f"----- {n} -----\n{_log_tail(lp)}" for n, _, lp in launched) + pytest.fail(f"precheck instance {name} timed out\n{tails}") + finally: + _terminate_process_groups(procs) + + +def _read_status(work_dir, name): + text_path = os.path.join(work_dir, "status", f"{name}.status") + json_path = os.path.join(work_dir, "status", f"{name}.json") + with open(text_path) as f: + text = f.read() + with open(json_path) as f: + doc = json.load(f) + return text, doc + + +def _assert_all_passed(work_dir, launched): + failures = [] + for name, proc, log_path in launched: + if proc.returncode != 0: + failures.append(f"{name} exited rc={proc.returncode}:\n{_log_tail(log_path)}") + if failures: + pytest.fail("\n".join(failures)) + for name, _, log_path in launched: + text, doc = _read_status(work_dir, name) + assert text.startswith(f"PASS {name}"), ( + f"{name} status not PASS: {text}\n{_log_tail(log_path)}" + ) + assert doc["overall"] == "PASS" + assert doc["transceiver_runtime"] == "PYTHON" + assert doc["kv_cache_manager"] == "V2" + + +def _jobs(cfg, config_path): + ctx_world = cfg["worker_config"]["ctx"]["tensor_parallel_size"] + gen_world = cfg["worker_config"]["gen"]["tensor_parallel_size"] + jobs = [] + for i in range(cfg["hardware"]["num_ctx_servers"]): + jobs.append(("ctx", i, ctx_world, config_path)) + for i in range(cfg["hardware"]["num_gen_servers"]): + jobs.append(("gen", i, gen_world, config_path)) + return jobs + + +@pytest.mark.timeout(300) +@pytest.mark.parametrize( + "num_ctx,num_gen,ctx_tp,gen_tp", + [ + pytest.param(1, 1, 1, 1, id="symmetric_1x1"), + pytest.param(2, 2, 1, 1, id="multi_instance_2x2"), + pytest.param(1, 1, 2, 1, id="asymmetric_tp2_to_tp1"), + ], +) +def test_precheck_passes(tmp_path, num_ctx, num_gen, ctx_tp, gen_tp): + """All (ctx, gen) pairs transfer + verify on one GPU; every verdict PASS.""" + pytest.importorskip("mpi4py") + cfg = _disagg_yaml(num_ctx, num_gen, ctx_tp, gen_tp) + config_path, models_root = _write_inputs(tmp_path, cfg) + work_dir, launched = _launch_instances(tmp_path, _jobs(cfg, config_path), models_root) + _wait_all(launched) + _assert_all_passed(work_dir, launched) + # Every gen instance must have exercised every ctx peer. + for gj in range(num_gen): + _, doc = _read_status(work_dir, f"gen_{gj}") + peers = {c["peer"] for c in doc["cases"] if c["status"] == "PASS"} + assert peers == {f"ctx_{ci}" for ci in range(num_ctx)} + + +@pytest.mark.timeout(300) +def test_precheck_fails_fast_on_fingerprint_mismatch(tmp_path): + """Mismatched ctx/gen yamls must produce FAIL verdicts, not a hang. + + This is the failure-attribution path the SLURM gate consumes: non-zero + exit codes plus .status files whose first line names the root cause. + """ + pytest.importorskip("mpi4py") + ctx_cfg = _disagg_yaml(1, 1, 1, 1, request_lengths=(64,)) + gen_cfg = _disagg_yaml(1, 1, 1, 1, request_lengths=(32, 64)) # different fingerprint + ctx_path, models_root = _write_inputs(tmp_path, ctx_cfg, name="ctx") + gen_path, _ = _write_inputs(tmp_path, gen_cfg, name="gen") + jobs = [("ctx", 0, 1, ctx_path), ("gen", 0, 1, gen_path)] + work_dir, launched = _launch_instances(tmp_path, jobs, models_root) + _wait_all(launched) + + for name, proc, log_path in launched: + assert proc.returncode != 0, ( + f"{name} unexpectedly passed with mismatched yamls\n{_log_tail(log_path)}" + ) + text, doc = _read_status(work_dir, name) + assert text.startswith(f"FAIL {name}"), f"{name} status: {text}" + assert doc["overall"] == "FAIL" + # The gen driver saw the ctx abort its handshake: the reason must name it. + text, _ = _read_status(work_dir, "gen_0") + assert "fingerprint" in text or "abort" in text.lower() diff --git a/tests/unittest/others/test_cache_transceiver_precheck_run.py b/tests/unittest/others/test_cache_transceiver_precheck_run.py new file mode 100644 index 000000000000..d0f52a8107fd --- /dev/null +++ b/tests/unittest/others/test_cache_transceiver_precheck_run.py @@ -0,0 +1,419 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Unit tests for the disagg cache-transceiver precheck DRIVER. + +Target: tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py + +Two halves: + +- Pure-logic tests (no torch / tensorrt_llm / MPI): wire format, rid/seed + scheme, rendezvous + abort-flag files, StatusRecorder, bandwidth CSV + parsing, schedule/timeout derivation. + +- Internal-API contract tests (import tensorrt_llm, no GPU work): the + precheck drives TRT-LLM internals directly (_torch.pyexecutor.*, + bindings.internal.*, private llm_utils resolvers), which carry no stability + promise. run_precheck.load_internal_apis() is the single owner of those + imports; these tests exercise it plus the constructor/signature shapes the + driver relies on, so an upstream refactor fails HERE in pre-merge CI + instead of aborting the SLURM disagg perf pipeline at runtime. +""" + +import base64 +import json +import os +import sys +import types + +import pytest + +_PRECHECK_DIR = os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "..", + "..", + "..", + "tests", + "scripts", + "perf-sanity", + "cache_transceiver_precheck", +) +sys.path.insert(0, os.path.abspath(_PRECHECK_DIR)) + +import run_precheck as rp # noqa: E402 (stdlib-only at import time) + +KEY = b"k" * 32 + + +# --------------------------------------------------------------------------- # +# rid / seed scheme +# --------------------------------------------------------------------------- # +def test_make_rid_unique_and_dense_within_session(): + num_ctx, num_gen, seqs = 3, 2, 50 + rids = set() + for gj in range(num_gen): + for ci in range(num_ctx): + session = [rp.make_rid(ci, gj, num_ctx, s) for s in range(seqs)] + # Dense in-session sequence: consecutive rids -> unique low-12-bit + # tags among any 4096 consecutive requests (tagFromRequestId). + assert session == list(range(session[0], session[0] + seqs)) + rids.update(session) + assert len(rids) == num_ctx * num_gen * seqs + assert all(r >= 1 for r in rids) + + +def test_seed_for_deterministic_and_distinct(): + assert rp.seed_for(7, 3) == rp.seed_for(7, 3) # rank-independent by construction + seeds = {rp.seed_for(rid, layer) for rid in (1, 2, 3) for layer in (0, 1, 2)} + assert len(seeds) == 9 + assert all(0 <= s <= 0x7FFFFFFF for s in seeds) + + +# --------------------------------------------------------------------------- # +# HMAC control-channel wire format +# --------------------------------------------------------------------------- # +def test_pack_unpack_roundtrip(): + obj = ["go", {"li": 0, "rep": 1, "wave": [0, 1]}] + assert rp.unpack_msg(rp.pack_msg(obj, KEY), KEY) == obj + + +def test_unpack_rejects_tampered_frame(): + raw = bytearray(rp.pack_msg(["hello", {}], KEY)) + raw[0] ^= 0xFF + with pytest.raises(rp._TransferError, match="HMAC"): + rp.unpack_msg(bytes(raw), KEY) + + +def test_unpack_rejects_wrong_key(): + raw = rp.pack_msg(["hello", {}], KEY) + with pytest.raises(rp._TransferError, match="HMAC"): + rp.unpack_msg(raw, b"x" * 32) + + +def test_unpack_rejects_short_frame(): + with pytest.raises(rp._TransferError, match="too short"): + rp.unpack_msg(b"tiny", KEY) + + +def test_params_to_wire_is_json_safe(): + p = types.SimpleNamespace( + first_gen_tokens=[1, 2], + req_id=42, + opaque_state=b"\x00\x01binary", + draft_tokens=None, + ctx_dp_rank=3, + disagg_info_endpoint="tcp://h:1", + ) + wire = rp.params_to_wire(p) + decoded = json.loads(json.dumps(wire)) # must survive the ZMQ JSON hop + assert base64.b64decode(decoded["opaque_state"]) == p.opaque_state + assert decoded["req_id"] == 42 + assert decoded["ctx_dp_rank"] == 3 + assert decoded["ctx_info_endpoint"] == "tcp://h:1" + + +# --------------------------------------------------------------------------- # +# Rendezvous + abort-flag files +# --------------------------------------------------------------------------- # +def test_addr_roundtrip(tmp_path, monkeypatch): + monkeypatch.setenv("SLURM_JOB_ID", "123") + path = rp.addr_path(str(tmp_path), 0, 1) + rp.write_addr(path, {"host": "h", "port": 5, "key": KEY.hex()}) + assert os.stat(path).st_mode & 0o777 == 0o600 # carries the HMAC key + payload = rp.wait_for_addr(path, timeout_s=5) + assert (payload["host"], payload["port"], payload["job"]) == ("h", 5, "123") + + +def test_wait_for_addr_rejects_stale_job(tmp_path, monkeypatch): + monkeypatch.setenv("SLURM_JOB_ID", "old-run") + path = rp.addr_path(str(tmp_path), 0, 0) + rp.write_addr(path, {"host": "h", "port": 5, "key": KEY.hex()}) + monkeypatch.setenv("SLURM_JOB_ID", "new-run") # requeued job, reused work dir + with pytest.raises(rp._Timeout): + rp.wait_for_addr(path, timeout_s=1.5) + + +def test_wait_for_addr_times_out_on_missing_file(tmp_path): + with pytest.raises(rp._Timeout): + rp.wait_for_addr(rp.addr_path(str(tmp_path), 0, 0), timeout_s=0) + + +def test_abort_flag_roundtrip_and_write_once(tmp_path, monkeypatch): + monkeypatch.setenv("SLURM_JOB_ID", "1") + work = str(tmp_path) + assert rp.abort_flag_reason(work) is None + rp.raise_abort_flag(work, "first failure\nsecond line ignored") + assert rp.abort_flag_reason(work) == "first failure" + rp.raise_abort_flag(work, "later failure") # write-once: first reason wins + assert rp.abort_flag_reason(work) == "first failure" + + +def test_abort_flag_stale_job_ignored(tmp_path, monkeypatch): + monkeypatch.setenv("SLURM_JOB_ID", "old") + rp.raise_abort_flag(str(tmp_path), "stale") + monkeypatch.setenv("SLURM_JOB_ID", "new") + assert rp.abort_flag_reason(str(tmp_path)) is None + + +# --------------------------------------------------------------------------- # +# StatusRecorder +# --------------------------------------------------------------------------- # +def _read_status(tmp_path, name): + with open(os.path.join(str(tmp_path), "status", f"{name}.status")) as f: + text = f.read() + with open(os.path.join(str(tmp_path), "status", f"{name}.json")) as f: + doc = json.load(f) + return text, doc + + +def test_recorder_pass(tmp_path): + rec = rp.StatusRecorder(str(tmp_path), "gen", 0, is_leader=True) + rec.record("ctx_0", 1024, "PASS") + text, doc = _read_status(tmp_path, "gen_0") + assert text.startswith("RUNNING") # not final yet: a SIGKILL must not read as PASS + assert doc["overall"] == "RUNNING" + rec.finalize(extra={"transceiver_runtime": "CPP"}) + text, doc = _read_status(tmp_path, "gen_0") + assert text.startswith("PASS gen_0") + assert doc["overall"] == "PASS" + assert doc["transceiver_runtime"] == "CPP" + + +def test_recorder_failure_summary_first_line_only(tmp_path): + rec = rp.StatusRecorder(str(tmp_path), "ctx", 1, is_leader=True) + rec.record("gen_0", 1024, "PASS") + rec.record("gen_1", 2048, "TRANSFER_ERROR", "boom\ntraceback line\nmore") + rec.finalize() + text, doc = _read_status(tmp_path, "ctx_1") + assert text.startswith("FAIL ctx_1") + assert "boom | traceback line" in text + assert "more" not in text # full reason only in the json + assert doc["overall"] == "FAIL" + assert doc["cases"][1]["reason"].endswith("more") + + +def test_recorder_skip_is_not_a_failure(tmp_path): + rec = rp.StatusRecorder(str(tmp_path), "gen", 0, is_leader=True) + rec.record("ctx_0", 0, "SKIP", "fail-fast") + assert rec.failed_cases() == [] + rec.finalize() + text, _ = _read_status(tmp_path, "gen_0") + assert text.startswith("PASS") + + +def test_recorder_non_leader_writes_nothing(tmp_path): + rec = rp.StatusRecorder(str(tmp_path), "gen", 0, is_leader=False) + rec.record("ctx_0", 0, "TRANSFER_ERROR", "x") + rec.finalize() + assert not os.path.exists(os.path.join(str(tmp_path), "status")) + + +# --------------------------------------------------------------------------- # +# Bandwidth CSV parsing +# --------------------------------------------------------------------------- # +def test_parse_bandwidth_gbps_median(tmp_path): + path = tmp_path / "rank_2_recv.csv" + path.write_text("RequestID,Bandwidth(Gbps),Delay(ms)\n1,80,0\n2,160,0\n3,240,0\n") + # Gbps -> GB/s (/8); median of [10, 20, 30] + assert rp.parse_bandwidth_gbps(str(tmp_path), 2) == 20.0 + + +def test_parse_bandwidth_gbps_missing_or_malformed(tmp_path): + assert rp.parse_bandwidth_gbps(str(tmp_path), 0) is None + (tmp_path / "rank_0_recv.csv").write_text("RequestID,Delay(ms)\n1,0\n") + assert rp.parse_bandwidth_gbps(str(tmp_path), 0) is None + + +def test_parse_python_bandwidth_gbps(tmp_path): + (tmp_path / "perf_a_0.csv").write_text( + "task_type,throughput_mbs\nKVSendTask,1024\nKVRecvTask,\n" + ) + (tmp_path / "perf_a_1.csv").write_text("task_type,throughput_mbs\nKVSendTask,3072\n") + # MB/s -> GB/s (/1024); median of [1, 3] + assert rp.parse_python_bandwidth_gbps(str(tmp_path)) == 2.0 + assert rp.parse_python_bandwidth_gbps(str(tmp_path / "nowhere")) is None + + +# --------------------------------------------------------------------------- # +# Schedule / timeout derivation +# --------------------------------------------------------------------------- # +def _plan(**overrides): + plan = { + "request_lengths": [64, 128], + "warmup_requests": 1, + "num_requests": 2, + "n_pairs": 3, + "wave_size": 2, + "rendezvous_timeout_s": 600, + "wireup_timeout_s": 300, + "wave_timeout_s": 180, + } + plan.update(overrides) + return plan + + +def test_schedule_covers_all_cells_in_lockstep_order(): + plan = _plan() + sched = rp._schedule(plan) + # 2 lengths x (1 warmup + 2 measured) reps x 2 waves ([0,1] and [2]) + assert len(sched) == 2 * 3 * 2 + assert sched[0] == (0, 64, 0, [0, 1]) + assert sched[1] == (0, 64, 0, [2]) + assert sched[-1] == (1, 128, 2, [2]) + + +def test_timeout_budgets(): + plan = _plan() + # Handshakes serialize across peers: rendezvous + per-peer slack. + assert rp.hello_timeout_s(plan, 2) == 600 + 2 * (300 + 300) + # Only the schedule's FIRST rep pays the NIXL wire-up allowance. + assert rp.wave_timeout_s(plan, 0, 0) == 180 + 300 + assert rp.wave_timeout_s(plan, 0, 1) == 180 + assert rp.wave_timeout_s(plan, 1, 0) == 180 + + +# --------------------------------------------------------------------------- # +# Internal-API contract (imports tensorrt_llm; no GPU work) +# --------------------------------------------------------------------------- # +class TestInternalApiContract: + @pytest.fixture(scope="class") + def api(self): + pytest.importorskip("tensorrt_llm") + return rp.load_internal_apis() + + def test_loader_caches(self, api): + assert rp.load_internal_apis() is api + + def test_create_kv_cache_transceiver_signature(self, api): + import inspect + + params = inspect.signature(api.create_kv_cache_transceiver).parameters + # Exactly the positional call shape PrecheckRunner.setup uses. + assert list(params)[:5] == [ + "mapping", + "dist", + "kv_cache_manager", + "attention_type", + "cache_transceiver_config", + ] + + def test_transceiver_interface_methods(self, api): + import importlib + + mod = importlib.import_module(api.create_kv_cache_transceiver.__module__) + base = mod.KvCacheTransceiver + for meth in ( + "respond_and_send_async", + "request_and_receive_async", + "check_context_transfer_status", + "check_gen_transfer_status", + ): + assert hasattr(base, meth), f"KvCacheTransceiver lost {meth}" + + @pytest.mark.parametrize("manager_attr", ["KVCacheManager", "KVCacheManagerV2"]) + def test_kv_cache_manager_ctor_kwargs(self, api, manager_attr): + import inspect + + params = inspect.signature(getattr(api, manager_attr).__init__).parameters + needed = { + "num_layers", + "num_kv_heads", + "head_dim", + "tokens_per_block", + "max_seq_len", + "max_batch_size", + "mapping", + "dtype", + "spec_config", + } + if manager_attr == "KVCacheManagerV2": + needed |= {"vocab_size", "is_disagg"} + missing = needed - set(params) + assert not missing, f"{manager_attr} ctor lost kwargs: {sorted(missing)}" + + def test_serving_resolvers(self, api): + import inspect + + assert len(inspect.signature(api.resolve_kv_cache_manager_v2_auto).parameters) == 2 + rt = inspect.signature(api.resolve_transceiver_runtime_auto).parameters + assert list(rt)[:1] == ["llm_args"] and len(rt) >= 3 + + def test_enum_members(self, api): + for enum, members in ( + (api.DataType, ("FP8", "HALF", "BF16")), + (api.CacheTypeCpp, ("SELF", "SELFKONLY")), + (api.AttentionTypeCpp, ("DEFAULT", "MLA")), + (api.LlmRequestState, ("DISAGG_GENERATION_TRANS_COMPLETE", "DISAGG_TRANS_ERROR")), + ( + api.LlmRequestType, + ("LLMREQUEST_TYPE_CONTEXT_ONLY", "LLMREQUEST_TYPE_GENERATION_ONLY"), + ), + ): + for m in members: + assert hasattr(enum, m), f"{enum} lost member {m}" + + def test_hang_detector_surface(self, api): + import inspect + + params = inspect.signature(api.HangDetector.__init__).parameters + assert {"timeout", "on_detected"} <= set(params) + for meth in ("start", "checkpoint", "cancel_task", "stop"): + assert hasattr(api.HangDetector, meth), f"HangDetector lost {meth}" + + def test_config_constructors(self, api): + cache_cfg = api.CacheTransceiverConfig(backend="UCX", max_tokens_in_buffer=1024) + assert hasattr(cache_cfg, "transceiver_runtime") + api.KvCacheConfigCpp(max_tokens=64, enable_block_reuse=False) + api.MTPDecodingConfig(num_nextn_predict_layers=1) + api.Mapping( + world_size=1, + rank=0, + gpus_per_node=1, + tp_size=1, + pp_size=1, + cp_size=1, + enable_attention_dp=False, + ) + assert hasattr(api.Distributed, "get") + + def test_params_wire_roundtrip_through_real_bindings(self, api): + # opaque_state must DESERIALIZE in the ContextPhaseParams ctor + # (arbitrary bytes -> std::bad_alloc), so use the serialized empty + # state: b"" is re-encoded by the bindings into its canonical form. + source = api.DisaggregatedParams( + ctx_request_id=42, + first_gen_tokens=[7, 8], + opaque_state=b"", + draft_tokens=[9], + ctx_dp_rank=1, + ctx_info_endpoint="tcp://host:1234", + ).get_context_phase_params() + restored = rp.params_from_wire(rp.params_to_wire(source)) + assert rp.params_to_wire(restored) == rp.params_to_wire(source) + + def test_make_request_shapes(self, api): + ctx_req = rp.make_request(True, rid=11, req_len=8, runtime="CPP") + assert ctx_req.py_request_id == 11 + py_ctx = rp.make_request(True, rid=12, req_len=8, runtime="PYTHON") + assert py_ctx.py_disaggregated_params.request_type == "context_only" + ctx_params = api.DisaggregatedParams( + ctx_request_id=13, + first_gen_tokens=[1], + opaque_state=b"", + ctx_dp_rank=0, + ).get_context_phase_params() + gen_req = rp.make_request(False, rid=13, req_len=8, runtime="CPP", ctx_params=ctx_params) + assert gen_req.py_request_id == 13 + py_gen = rp.make_request(False, rid=14, req_len=8, runtime="PYTHON", ctx_params=ctx_params) + assert py_gen.py_disaggregated_params.request_type == "generation_only" From 222c21a6f77a5bd816d0d733869fe66dc9a02891 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Tue, 28 Jul 2026 08:17:47 +0000 Subject: [PATCH 20/21] [None][fix] Parse the renamed transceiver perf CSVs in the harness and precheck PerfLogManager now gives TRTLLM_KVCACHE_TIME_OUTPUT_PATH top priority and names Python task CSVs '_.csv'; the C++ transceiver likewise moved to '__.csv'. The harness report and the precheck driver still globbed the old 'py_*'/'perf_*'/'rank_*' names, so transfers PASSed while every bandwidth field silently came back empty. - report.py / run_precheck.py: identify Python task CSVs by their header columns (unique_rid/task_type + throughput_mbs) instead of a filename prefix; match C++ CSVs by the '__.csv' suffix - run_precheck.py: read repeated Bandwidth(Gbps) columns with csv.reader (DictReader kept only the last transmission) and fix the MiB/s -> GB/s conversion; report.py: count only KVSendTask rows so aux sends cannot drag the median down - drop the now-dead legacy perf env vars; one env var drives both sides - unit tests updated to the real filenames plus new negative coverage; both e2e tests now assert per-GPU bandwidth is present so a future writer/parser naming mismatch fails loudly instead of silently Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../slurm/cache_transceiver_test/report.py | 24 ++++-- .../run_cache_transceiver_test.py | 13 ++-- .../cache_transceiver_precheck/README.md | 4 +- .../run_precheck.py | 78 +++++++++++++------ .../test_cache_transceiver_harness.py | 18 +++++ .../test_cache_transceiver_harness_report.py | 55 ++++++++++++- .../test_cache_transceiver_precheck_e2e.py | 10 +++ .../test_cache_transceiver_precheck_config.py | 21 +++-- .../test_cache_transceiver_precheck_run.py | 56 +++++++++++-- 9 files changed, 226 insertions(+), 53 deletions(-) diff --git a/examples/disaggregated/slurm/cache_transceiver_test/report.py b/examples/disaggregated/slurm/cache_transceiver_test/report.py index 00ecab6b4e49..d1cf66203a6d 100644 --- a/examples/disaggregated/slurm/cache_transceiver_test/report.py +++ b/examples/disaggregated/slurm/cache_transceiver_test/report.py @@ -227,7 +227,16 @@ def _parse_cpp_recv_csvs(csv_dir): def _parse_python_csvs(csv_dir): - """Return {rid: [per_rank_GBps, ...]} from perf_logger py_*_*.csv files. + """Return {rid: [per_rank_GBps, ...]} from Python perf_logger task CSVs. + + PerfLogManager names its per-task CSV by priority of the env vars it sees: + with TRTLLM_KVCACHE_TIME_OUTPUT_PATH set (which this harness always sets + for the C++ transceiver), files are "{dir}/{instanceUuid}_{rank}.csv"; + only the legacy TLLM_KV_TRANSFER_PERF_LOG_FILE fallback produces the old + "py_{instanceUuid}_{rank}.csv" prefix. So identify Python task CSVs by + their header columns (unique_rid + throughput_mbs) rather than by file + name -- C++ send/recv CSVs and the gen_transfer_summary CSV have neither + column and are skipped. `throughput_mbs` is recorded for send tasks in MiB/s (perf_logger divides by 1024*1024). We convert MiB/s -> GB/s (`* 1024^2 / 1e9`) so it matches the @@ -235,14 +244,17 @@ def _parse_python_csvs(csv_dir): `_parse_cpp_recv_csvs`. """ per_rid = {} # rid -> list[per-rank throughput GB/s] - for path in glob.glob(os.path.join(csv_dir, "py_*_*.csv")): + for path in glob.glob(os.path.join(csv_dir, "*.csv")): with open(path) as f: reader = csv.DictReader(f) + fields = reader.fieldnames or [] + if "unique_rid" not in fields or "throughput_mbs" not in fields: + continue # not a Python perf_logger task CSV (e.g. C++ send/recv) for row in reader: - if "unique_rid" not in row or "throughput_mbs" not in row: - continue - if "Send" not in row.get("task_type", ""): - continue # throughput is meaningful only for send tasks + if row.get("task_type") != "KVSendTask": + continue # only KV sends carry meaningful throughput; + # AuxSendTask rows are tiny metadata sends that would + # drag the median down, KVRecvTask rows have none try: rid = int(float(row["unique_rid"])) mbs = float(row["throughput_mbs"]) diff --git a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py index 138a43c96063..e8577cb6ceaa 100644 --- a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py +++ b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py @@ -25,7 +25,8 @@ received blocks regenerate to the same pattern. Bandwidth is emitted by the transceivers themselves into per-rank CSVs (parsed later by report.py): C++ -> TRTLLM_KVCACHE_TIME_OUTPUT_PATH (_*_send.csv / _*_recv.csv) - Py -> TLLM_KV_TRANSFER_PERF_LOG_FILE (py_*_*.csv, throughput_mbs) + Py -> same env var (PerfLogManager gives it top priority): + _.csv, throughput_mbs column This driver mirrors the single-process test (tests/unittest/others/ test_kv_cache_transceiver.py) and the multi-process Python test @@ -487,12 +488,12 @@ def main(): work_dir = cfg["environment"]["work_dir"] csv_dir = os.path.join(work_dir, "csv", str(sweep), role) os.makedirs(csv_dir, exist_ok=True) + # One env var drives both transceivers: C++ caches it on first read, and + # Python's PerfLogManager gives it top priority (enabling perf logging and + # writing task CSVs as _.csv in csv_dir; the + # per-transceiver UUID avoids cross-combination collisions, and report.py + # identifies these files by header columns, not name). os.environ["TRTLLM_KVCACHE_TIME_OUTPUT_PATH"] = csv_dir - # Python perf logging (singleton reads these once; C++ ignores them). Set at - # startup so it is enabled regardless of combination ordering. Per-transceiver UUID - # filenames (py__.csv) avoid cross-combination collisions. - os.environ["TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO"] = "1" - os.environ["TLLM_KV_TRANSFER_PERF_LOG_FILE"] = os.path.join(csv_dir, "py") status_dir = os.path.join(work_dir, "status") os.makedirs(status_dir, exist_ok=True) diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md index 9ec289a4f230..2588c1362157 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md @@ -68,7 +68,9 @@ status/{ctx,gen}_.json # per-case detail: PASS/TIMEOUT/TRANSFER_ERROR/ # MISMATCH/INIT_ERROR + reason + UCX env snapshot status/{ctx,gen}_.status # one-line verdict (parsed by the launch script) logs/{ctx,gen}_.log # merged per-rank logs (UCX_PROTO_INFO=used table) -csv/gen_/rank_*_recv.csv # C++ transceiver per-request bandwidth +csv/gen_/__recv.csv # C++ transceiver per-request bandwidth +csv/ctx_/_.csv # Python transceiver per-task perf + # (KVSendTask throughput on the ctx side) ``` ## Manual runs diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index 85f769167566..58dbd8fa5e7c 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -718,44 +718,75 @@ def parse_bandwidth_gbps(csv_dir, rank, tag="recv"): """Median per-request bandwidth in GB/s (bytes/1e9), best-effort. Parsed from the C++ transceiver CSV this rank wrote via - TRTLLM_KVCACHE_TIME_OUTPUT_PATH. + TRTLLM_KVCACHE_TIME_OUTPUT_PATH, named "__.csv" + (instanceId is a runtime UUID), so match by the "__.csv" + suffix. The suffix's leading "_" keeps rank 1 from matching rank 11. + + Each row repeats the Bandwidth(Gbps) column once per transmission, so use + csv.reader (DictReader would keep only the last duplicate) and take the + mean transmission bandwidth as that request's value, then the median + across requests -- same semantics as the harness report. """ import csv as csv_mod import statistics - path = os.path.join(csv_dir, f"rank_{rank}_{tag}.csv") + suffix = f"_{rank}_{tag}.csv" try: - with open(path) as f: - rows = list(csv_mod.DictReader(f)) - col = next((c for c in (rows[0] or {}) if "Bandwidth" in c), None) - if col is None: - return None - vals = [float(r[col]) / 8.0 for r in rows if r.get(col)] - return statistics.median(vals) if vals else None - except (OSError, ValueError, IndexError, StopIteration): + names = [n for n in os.listdir(csv_dir) if n.endswith(suffix)] + except OSError: return None + vals = [] + for name in names: + try: + with open(os.path.join(csv_dir, name)) as f: + reader = csv_mod.reader(f) + header = next(reader, None) + if not header: + continue + bw_cols = [i for i, c in enumerate(header) if "Bandwidth" in c] + if not bw_cols: + continue + for row in reader: + bws = [] + for i in bw_cols: + if i < len(row) and row[i]: + try: + bws.append(float(row[i]) / 8.0) # Gbps -> GB/s + except ValueError: + pass + if bws: + vals.append(sum(bws) / len(bws)) + except OSError: + continue + return statistics.median(vals) if vals else None def parse_python_bandwidth_gbps(csv_dir): """Median KV-send throughput in GB/s from the Python transceiver's perf CSVs. - Written by perf_logger.py; enabled via TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO=1 - + TLLM_KV_TRANSFER_PERF_LOG_FILE=, which writes - __.csv. throughput_mbs (MB/s) is on the SENDER - (ctx) side, task_type=KVSendTask; receiver rows have no throughput. - Best-effort -- returns None when no perf CSV exists. + Written by perf_logger.py, which gives TRTLLM_KVCACHE_TIME_OUTPUT_PATH top + priority and names files "{dir}/{instanceUuid}_{rank}.csv" -- so identify + the CSVs by their header columns (task_type + throughput_mbs) rather than + by name; C++ send/recv CSVs have neither column. throughput_mbs (MiB/s) is + on the SENDER (ctx) side, task_type=KVSendTask; receiver rows have no + throughput. Best-effort -- returns None when no perf CSV exists. """ import csv as csv_mod import glob import statistics vals = [] - for path in glob.glob(os.path.join(csv_dir, "perf_*_*.csv")): + for path in glob.glob(os.path.join(csv_dir, "*.csv")): try: with open(path) as f: - for r in csv_mod.DictReader(f): + reader = csv_mod.DictReader(f) + fields = reader.fieldnames or [] + if "task_type" not in fields or "throughput_mbs" not in fields: + continue + for r in reader: if r.get("task_type") == "KVSendTask" and r.get("throughput_mbs"): - vals.append(float(r["throughput_mbs"]) / 1024.0) # MB/s -> GB/s + # MiB/s -> GB/s + vals.append(float(r["throughput_mbs"]) * 1024.0 * 1024.0 / 1e9) except (OSError, ValueError): continue return statistics.median(vals) if vals else None @@ -816,13 +847,12 @@ def setup(self, kv_shape, max_req_len): enable_attention_dp=par["enable_attention_dp"], ) os.makedirs(self.csv_dir, exist_ok=True) - # C++ transceiver bandwidth CSV (per-rank recv). + # One env var drives both transceivers' bandwidth CSVs: C++ writes + # per-rank "__send/recv.csv", and Python's + # PerfLogManager gives the same var top priority, writing task CSVs + # as "_.csv" (KVSendTask throughput on the ctx + # side). os.environ["TRTLLM_KVCACHE_TIME_OUTPUT_PATH"] = self.csv_dir - # Python transceiver bandwidth CSV (perf_logger.py): writes - # {base}_{instance}_{rank}.csv with KVSendTask throughput on the ctx - # side. Harmless to set unconditionally -- the C++ path ignores it. - os.environ["TLLM_ENABLE_CACHE_TRANSFER_PERF_INFO"] = "1" - os.environ["TLLM_KV_TRANSFER_PERF_LOG_FILE"] = os.path.join(self.csv_dir, "perf") # Built VERBATIM from the disagg yaml's cache_transceiver_config so # backend/max_tokens_in_buffer/timeouts match the real test exactly. diff --git a/tests/unittest/disaggregated/test_cache_transceiver_harness.py b/tests/unittest/disaggregated/test_cache_transceiver_harness.py index baaa3f421f35..cb37e89b4660 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_harness.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_harness.py @@ -353,6 +353,24 @@ def test_single_node_transfer(tmp_path): f"ctx log tail:\n{_log_tail(ctx_log)}\n" f"gen log tail:\n{_log_tail(gen_log)}" ) + # Bandwidth must actually be parsed from the perf CSVs, not just + # the transfer verified. The PYTHON+NIXL combination writes its + # perf CSVs via PerfLogManager into TRTLLM_KVCACHE_TIME_OUTPUT_PATH + # (set by the driver) as "_.csv"; a naming or + # glob mismatch between perf_logger.py and report.py leaves + # per_gpu_BW_GBps None while status stays PASS, so assert on it + # explicitly. + ctx_csv_dir = os.path.join(work_dir, "csv", "0", "ctx") + csv_listing = os.listdir(ctx_csv_dir) if os.path.isdir(ctx_csv_dir) else [] + assert sweep["per_gpu_BW_GBps"] is not None and sweep["per_gpu_BW_GBps"] > 0, ( + f"per_gpu_BW_GBps missing for {combo['combination']} " + f"sweep={sweep['sweep']} — perf CSVs were not parsed " + f"(ctx csv dir contents: {csv_listing})" + ) + assert sweep["num_samples"] > 0, ( + f"num_samples is 0 for {combo['combination']} sweep={sweep['sweep']} " + f"(ctx csv dir contents: {csv_listing})" + ) best_path = os.path.splitext(results_path)[0] + ".best.json" assert os.path.exists(best_path), "results.best.json was not created" diff --git a/tests/unittest/disaggregated/test_cache_transceiver_harness_report.py b/tests/unittest/disaggregated/test_cache_transceiver_harness_report.py index 491eab1e0b67..ecd8744cea1a 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_harness_report.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_harness_report.py @@ -227,7 +227,7 @@ def test_basic(self, tmp_path): { "unique_rid": "0", "throughput_mbs": "1048.576", - "task_type": "Send", + "task_type": "KVSendTask", "other": "", } ) @@ -236,7 +236,17 @@ def test_basic(self, tmp_path): { "unique_rid": "0", "throughput_mbs": "2000.0", - "task_type": "Recv", + "task_type": "KVRecvTask", + "other": "", + } + ) + # Aux sends are tiny metadata transfers; they must not drag the + # KV bandwidth stats down. + w.writerow( + { + "unique_rid": "0", + "throughput_mbs": "1.0", + "task_type": "AuxSendTask", "other": "", } ) @@ -246,6 +256,47 @@ def test_basic(self, tmp_path): assert len(result[0]) == 1 assert abs(result[0][0] - 1048.576 * 1024 * 1024 / 1e9) < 0.001 + def test_cpp_naming_without_py_prefix(self, tmp_path): + """PerfLogManager prefers the C++ output-path env var for naming. + + With TRTLLM_KVCACHE_TIME_OUTPUT_PATH set, Python task CSVs are named + "_.csv" (no py_ prefix); the parser must find them + by header columns instead of file name. + """ + csv_path = tmp_path / "cd93dae6-1111-2222-3333-444455556666_3.csv" + with open(csv_path, "w", newline="") as f: + w = csv.DictWriter( + f, + fieldnames=["unique_rid", "throughput_mbs", "task_type"], + ) + w.writeheader() + w.writerow({"unique_rid": "0", "throughput_mbs": "1048.576", "task_type": "KVSendTask"}) + result = _parse_python_csvs(str(tmp_path)) + assert 0 in result + assert abs(result[0][0] - 1048.576 * 1024 * 1024 / 1e9) < 0.001 + + def test_ignores_cpp_csvs_in_same_dir(self, tmp_path): + """C++ CSVs sharing the directory must not contribute samples. + + C++ send/recv and gen-summary CSVs lack the unique_rid/throughput_mbs + columns, so the header check skips them. + """ + with open( + tmp_path / "3c9f0e2a-aaaa-bbbb-cccc-ddddeeeeffff_0_recv.csv", "w", newline="" + ) as f: + w = csv.writer(f) + w.writerow(["RequestID", "Bandwidth(Gbps)"]) + w.writerow([0, 10.0]) + with open( + tmp_path / "3c9f0e2a-aaaa-bbbb-cccc-ddddeeeeffff_0_gen_transfer_summary.csv", + "w", + newline="", + ) as f: + w = csv.writer(f) + w.writerow(["timestamp", "RequestID", "gen_side_transfer_time(ms)", "kv_cache_size"]) + w.writerow(["2026-01-01 00:00:00.000", 0, 1.0, 1024]) + assert _parse_python_csvs(str(tmp_path)) == {} + def test_empty_dir(self, tmp_path): assert _parse_python_csvs(str(tmp_path)) == {} diff --git a/tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py b/tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py index 7abde75cd12e..e8a92f99c99c 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_precheck_e2e.py @@ -264,6 +264,16 @@ def _assert_all_passed(work_dir, launched): assert doc["overall"] == "PASS" assert doc["transceiver_runtime"] == "PYTHON" assert doc["kv_cache_manager"] == "V2" + # The Python transceiver records bandwidth on the ctx (sender) leader + # via PerfLogManager CSVs in TRTLLM_KVCACHE_TIME_OUTPUT_PATH. Asserting + # it here ties the whole chain (perf_logger naming -> CSV -> parser) + # into the e2e check, so a writer/parser naming mismatch fails loudly + # instead of silently dropping the bandwidth field. + if name.startswith("ctx_"): + assert doc.get("per_gpu_bw_gbps", 0) > 0, ( + f"{name} verdict has no per_gpu_bw_gbps — Python perf CSVs were " + f"not written or not parsed\n{_log_tail(log_path)}" + ) def _jobs(cfg, config_path): diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py index 9212f17bfc05..0bcff41adb48 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_config.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py @@ -743,8 +743,10 @@ def test_sparse_attention_model_uses_simplified_mla_pool(tmp_path): def test_python_transceiver_bandwidth_csv(tmp_path): """Bandwidth from the Python transceiver's perf_logger CSVs. - Median over KVSendTask throughput_mbs (MB/s -> GB/s), receiver rows - ignored. + PerfLogManager names them "_.csv" (it gives + TRTLLM_KVCACHE_TIME_OUTPUT_PATH top priority); the parser identifies them + by header columns. Median over KVSendTask throughput_mbs (MiB/s -> GB/s), + receiver rows ignored. """ header = ( "timestamp,task_type,unique_rid,peer_rank,transfer_size_bytes," @@ -752,17 +754,20 @@ def test_python_transceiver_bandwidth_csv(tmp_path): "queue_latency_ms,transfer_latency_ms,task_latency_ms,throughput_mbs" ) # two ctx ranks, each its own perf file; KVSendTask rows carry throughput - (tmp_path / "perf_abc_0.csv").write_text( + (tmp_path / "cd93dae6-9d75-4b0e-8a89-2c9e2f0f1a2b_0.csv").write_text( header + "\n" - "t,KVSendTask,1,0,1000,,,0,0,0,0,102400.00\n" # 100 GB/s + "t,KVSendTask,1,0,1000,,,0,0,0,0,102400.00\n" "t,AuxSendTask,1,0,10,,,0,0,0,0,1.00\n" # tiny metadata, ignored ) - (tmp_path / "perf_abc_1.csv").write_text( - header + "\nt,KVSendTask,2,1,1000,,,0,0,0,0,204800.00\n" # 200 GB/s + (tmp_path / "cd93dae6-9d75-4b0e-8a89-2c9e2f0f1a2b_1.csv").write_text( + header + "\nt,KVSendTask,2,1,1000,,,0,0,0,0,204800.00\n" ) # a receiver file (no throughput) must not contribute - (tmp_path / "perf_def_8.csv").write_text(header + "\nt,KVRecvTask,3,0,,,,,,,5.0,\n") + (tmp_path / "5d7b1f80-aaaa-bbbb-cccc-ddddeeeeffff_8.csv").write_text( + header + "\nt,KVRecvTask,3,0,,,,,,,5.0,\n" + ) bw = rp.parse_python_bandwidth_gbps(str(tmp_path)) - assert bw == 150.0 # median(100, 200) + # median(102400, 204800) = 153600 MiB/s -> GB/s (*1024^2/1e9) + assert abs(bw - 153600 * 1024 * 1024 / 1e9) < 1e-9 # no perf files -> None assert rp.parse_python_bandwidth_gbps(str(tmp_path / "empty")) is None diff --git a/tests/unittest/others/test_cache_transceiver_precheck_run.py b/tests/unittest/others/test_cache_transceiver_precheck_run.py index d0f52a8107fd..a46fcb1276c7 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_run.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_run.py @@ -223,28 +223,72 @@ def test_recorder_non_leader_writes_nothing(tmp_path): # Bandwidth CSV parsing # --------------------------------------------------------------------------- # def test_parse_bandwidth_gbps_median(tmp_path): - path = tmp_path / "rank_2_recv.csv" + # C++ names timing CSVs "__.csv" (instanceId is a + # runtime UUID), so the parser must suffix-match, not expect "rank_*". + path = tmp_path / "3c9f0e2a-1111-2222-3333-444455556666_2_recv.csv" path.write_text("RequestID,Bandwidth(Gbps),Delay(ms)\n1,80,0\n2,160,0\n3,240,0\n") # Gbps -> GB/s (/8); median of [10, 20, 30] assert rp.parse_bandwidth_gbps(str(tmp_path), 2) == 20.0 +def test_parse_bandwidth_gbps_duplicate_columns_mean(tmp_path): + # C++ repeats the Bandwidth(Gbps) column once per transmission; the parser + # must average them per request (DictReader would keep only the last). + (tmp_path / "uuid_0_recv.csv").write_text( + "RequestID,Bandwidth(Gbps),Bandwidth(Gbps)\n1,80,240\n" + ) + # mean(80, 240) = 160 Gbps -> /8 = 20 GB/s + assert rp.parse_bandwidth_gbps(str(tmp_path), 0) == 20.0 + + +def test_parse_bandwidth_gbps_rank_suffix_no_cross_match(tmp_path): + # Rank 1 must not pick up rank 11's file (the suffix's leading "_"). + (tmp_path / "uuid_11_recv.csv").write_text("RequestID,Bandwidth(Gbps)\n1,80\n") + assert rp.parse_bandwidth_gbps(str(tmp_path), 1) is None + assert rp.parse_bandwidth_gbps(str(tmp_path), 11) == 10.0 + + def test_parse_bandwidth_gbps_missing_or_malformed(tmp_path): assert rp.parse_bandwidth_gbps(str(tmp_path), 0) is None - (tmp_path / "rank_0_recv.csv").write_text("RequestID,Delay(ms)\n1,0\n") + (tmp_path / "uuid_0_recv.csv").write_text("RequestID,Delay(ms)\n1,0\n") assert rp.parse_bandwidth_gbps(str(tmp_path), 0) is None def test_parse_python_bandwidth_gbps(tmp_path): - (tmp_path / "perf_a_0.csv").write_text( + # PerfLogManager gives TRTLLM_KVCACHE_TIME_OUTPUT_PATH top priority and + # names task CSVs "_.csv" (no fixed prefix); the + # parser identifies them by header columns, not name. + (tmp_path / "cd93dae6-9d75-4b0e-8a89-2c9e2f0f1a2b_0.csv").write_text( "task_type,throughput_mbs\nKVSendTask,1024\nKVRecvTask,\n" ) - (tmp_path / "perf_a_1.csv").write_text("task_type,throughput_mbs\nKVSendTask,3072\n") - # MB/s -> GB/s (/1024); median of [1, 3] - assert rp.parse_python_bandwidth_gbps(str(tmp_path)) == 2.0 + (tmp_path / "cd93dae6-9d75-4b0e-8a89-2c9e2f0f1a2b_1.csv").write_text( + "task_type,throughput_mbs\nKVSendTask,3072\n" + ) + # MiB/s -> GB/s (*1024^2/1e9); median of [1024, 3072] MiB/s = 2048 MiB/s + expected = 2048 * 1024 * 1024 / 1e9 + assert abs(rp.parse_python_bandwidth_gbps(str(tmp_path)) - expected) < 1e-9 assert rp.parse_python_bandwidth_gbps(str(tmp_path / "nowhere")) is None +def test_parse_python_bandwidth_gbps_ignores_cpp_csvs(tmp_path): + # C++ send/recv and gen-summary CSVs share csv_dir; they lack the + # task_type/throughput_mbs columns and must not contribute samples. + (tmp_path / "uuid_0_recv.csv").write_text("RequestID,Bandwidth(Gbps)\n1,80\n") + (tmp_path / "uuid_0_gen_transfer_summary.csv").write_text( + "RequestID,gen_side_transfer_time(ms),kv_cache_size\n1,1.0,1024\n" + ) + assert rp.parse_python_bandwidth_gbps(str(tmp_path)) is None + + +def test_parse_bandwidth_gbps_ignores_gen_summary(tmp_path): + # "__gen_transfer_summary.csv" must not match the + # "__recv.csv" suffix. + (tmp_path / "uuid_0_gen_transfer_summary.csv").write_text( + "RequestID,gen_side_transfer_time(ms),kv_cache_size\n1,1.0,1024\n" + ) + assert rp.parse_bandwidth_gbps(str(tmp_path), 0) is None + + # --------------------------------------------------------------------------- # # Schedule / timeout derivation # --------------------------------------------------------------------------- # From 427b1abd2f34c612490838d6576f85ce06191dd9 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Wed, 29 Jul 2026 01:34:24 +0000 Subject: [PATCH 21/21] [None][fix] Track the V2 resolver's new signature and disagg gating in the precheck Main's d6a2d25c7a (#16598) grew _resolve_kv_cache_manager_v2_auto a third optional param and taught it to read llm_args.cache_transceiver_config to gate a V2 model default on the NIXL Python transceiver. The precheck's contract test pinned the old 2-param signature, so it failed in CI, and the driver's shim lacked cache_transceiver_config, so a V2-default model would have silently fallen back to V1 via the except path. Resolve the transceiver runtime before the V2 setting (the resolver treats an unresolved 'auto' runtime as non-PYTHON), hand the resolver the cache config through the shim, and re-anchor the contract test on the call shape the driver actually uses: first two positional params fixed, any later addition must carry a default. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../run_precheck.py | 29 +++++++++++-------- .../test_cache_transceiver_precheck_run.py | 6 +++- 2 files changed, 22 insertions(+), 13 deletions(-) diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index 58dbd8fa5e7c..19c6bf70c856 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -334,6 +334,19 @@ def resolve_model_prefs(model_dir, side, cache_cfg): api = load_internal_apis() model_cls, hf_view = _lookup_model_cls(model_dir) + # Runtime BEFORE V2, like serving: the V2 resolver's disagg gating reads + # cache_cfg.transceiver_runtime and treats an unresolved "auto" as non-PYTHON. + if getattr(cache_cfg, "transceiver_runtime", None) == "auto": + try: + shim = types.SimpleNamespace(cache_transceiver_config=cache_cfg) + api.resolve_transceiver_runtime_auto(shim, model_cls, hf_view) + except Exception as e: # noqa: BLE001 - fall back to the create() default (CPP) + print( + f"[precheck] WARNING: transceiver_runtime 'auto' resolution failed " + f"({e!r}); create_kv_cache_transceiver will fall back to CPP", + flush=True, + ) + setting = side["use_kv_cache_manager_v2"] if setting == "auto": defaults = {} @@ -348,8 +361,11 @@ def resolve_model_prefs(model_dir, side, cache_cfg): try: # The REAL serving resolver, via the same shim pattern as the # runtime resolution below -- one owner for the 'auto' semantics. + # cache_transceiver_config feeds the resolver's disagg gating + # (a V2 model default requires the NIXL Python transceiver). shim = types.SimpleNamespace( - kv_cache_config=types.SimpleNamespace(use_kv_cache_manager_v2="auto") + kv_cache_config=types.SimpleNamespace(use_kv_cache_manager_v2="auto"), + cache_transceiver_config=cache_cfg, ) use_v2 = bool(api.resolve_kv_cache_manager_v2_auto(shim, defaults)) except Exception as e: # noqa: BLE001 - fall back like a missing model @@ -359,17 +375,6 @@ def resolve_model_prefs(model_dir, side, cache_cfg): use_v2 = False else: use_v2 = bool(setting) - - if getattr(cache_cfg, "transceiver_runtime", None) == "auto": - try: - shim = types.SimpleNamespace(cache_transceiver_config=cache_cfg) - api.resolve_transceiver_runtime_auto(shim, model_cls, hf_view) - except Exception as e: # noqa: BLE001 - fall back to the create() default (CPP) - print( - f"[precheck] WARNING: transceiver_runtime 'auto' resolution failed " - f"({e!r}); create_kv_cache_transceiver will fall back to CPP", - flush=True, - ) return use_v2 diff --git a/tests/unittest/others/test_cache_transceiver_precheck_run.py b/tests/unittest/others/test_cache_transceiver_precheck_run.py index a46fcb1276c7..ed060d1e9053 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_run.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_run.py @@ -389,7 +389,11 @@ def test_kv_cache_manager_ctor_kwargs(self, api, manager_attr): def test_serving_resolvers(self, api): import inspect - assert len(inspect.signature(api.resolve_kv_cache_manager_v2_auto).parameters) == 2 + # The driver calls resolve_kv_cache_manager_v2_auto(shim, defaults): + # the first two params are fixed, anything added later must default. + v2 = inspect.signature(api.resolve_kv_cache_manager_v2_auto).parameters + assert list(v2)[:2] == ["llm_args", "model_defaults_dict"] + assert all(p.default is not inspect.Parameter.empty for p in list(v2.values())[2:]) rt = inspect.signature(api.resolve_transceiver_runtime_auto).parameters assert list(rt)[:1] == ["llm_args"] and len(rt) >= 3