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 9920dc740de7..c42433190658 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 @@ -56,7 +56,10 @@ from tensorrt_llm._torch.distributed import Distributed 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.kv_cache_transceiver import ( + create_kv_cache_transceiver, + maybe_enable_fabric_memory_for_python_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 BlockReuseConfig, CacheTransceiverConfig @@ -830,6 +833,28 @@ def open_sock(): zmq_sock = open_sock() cases = build_cases(cfg) + # The C++ fabric-memory env getter is cached on its first KV pool + # allocation. Enable the default before any matrix case builds a pool, even + # when a C++ transceiver case appears before Python+V1 in the matrix. + python_v1_case = next( + (case for case in cases if case["runtime"] == "PYTHON" and case["cache_manager"] == "V1"), + None, + ) + if python_v1_case is not None: + maybe_enable_fabric_memory_for_python_transceiver( + CacheTransceiverConfig( + backend=python_v1_case["backend"], + transceiver_runtime="PYTHON", + ), + KVCacheManager, + ) + print( + f"[{role} rank={rank}] PYTHON+V1 case in matrix: " + "TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY=" + f"{os.environ.get('TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY')} " + "applies to every case in this run, including C++ transceiver ones", + flush=True, + ) req_lens = cfg["test_matrix"]["request_lengths"] warmup = cfg["test_matrix"]["warmup_requests"] num_req = cfg["test_matrix"]["num_requests_per_length"] diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 64a685887549..e13d976db38f 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -2373,8 +2373,10 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG fi done - # Kill tail -f process - kill \$tailPid + # Stop and reap the log follower. It may have already exited + # when the remote log stream closes; that is not a test failure. + kill \$tailPid 2>/dev/null || true + wait \$tailPid 2>/dev/null || true # Wait briefly to ensure accounting is consistent sleep 10 diff --git a/jenkins/scripts/perf/cluster_env.py b/jenkins/scripts/perf/cluster_env.py index cbc342aa2d24..465299790dc3 100644 --- a/jenkins/scripts/perf/cluster_env.py +++ b/jenkins/scripts/perf/cluster_env.py @@ -44,15 +44,19 @@ "rocep198s0:1,rocep199s0:1,rocep205s0:1,rocep206s0:1" " UCX_IB_GID_INDEX=auto UCX_IB_TRAFFIC_CLASS=52 UCX_IB_SL=0", ), - # oci-aga: avoid transports that fail on this VF fabric, disable DEVX to - # avoid UAR allocation failures, and pin the GPU-connected rail VFs. + # oci-aga: use TCP over IPv4 alongside the local CUDA/shared-memory + # transports. ( "oci-aga*", "*", - "export UCX_TLS=^tcp,rc_gda,gga UCX_IB_MLX5_DEVX=n " - "UCX_NET_DEVICES=" - "rdma_vf_rail0:1,rdma_vf_rail1:1,rdma_vf_rail2:1,rdma_vf_rail3:1 " - "UCX_IB_TRAFFIC_CLASS=96 TRTLLM_NIXL_NUM_THREADS=1", + "export UCX_TLS=cuda_ipc,cuda_copy,sm,self,tcp UCX_TCP_AF_PRIO=inet", + ), + # oci-hsg: UCX picks wrong RDMA devices; pin the usable mlx5 ports and + # keep eth0 as the TCP fallback device. + ( + "oci-hsg*", + "*", + "export UCX_NET_DEVICES=mlx5_0:1,mlx5_1:1,mlx5_3:1,mlx5_4:1,eth0", ), # nsc-svg: UCX picks wrong RDMA devices; pin the usable mlx5 ports. ( @@ -61,9 +65,15 @@ "export UCX_NET_DEVICES=" "mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,mlx5_4:1,mlx5_5:1,mlx5_10:1,mlx5_11:1", ), - # aws-cmh: UCX transport auto-selection hangs on this fabric; pin the - # working transport set explicitly. - ("aws-cmh*", "*", "export UCX_TLS=cuda_ipc,cuda_copy,sm,self,tcp"), + # aws-cmh: UCX transport/device auto-selection hangs on this fabric; pin + # the working transport set and Ethernet/RDMA devices explicitly. + ( + "aws-cmh*", + "*", + "export UCX_TLS=cuda_ipc,cuda_copy,sm,self,tcp " + "UCX_NET_DEVICES=eth0,mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1," + "mlx5_4:1,mlx5_5:1,mlx5_6:1,mlx5_7:1", + ), # aws-dfw: gdr_copy is broken on this cluster; exclude it. ("aws-dfw*", "*", "export UCX_TLS=^gdr_copy"), # Default: base unset only. diff --git a/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh b/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh index e3b02f88741d..857fe00c2f26 100644 --- a/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh +++ b/jenkins/scripts/perf/disaggregated/slurm_ct_precheck_gate.sh @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + # 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 @@ -21,7 +24,7 @@ ct_xml_escape() { # 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" \ + 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|WATCHDOG_KILL|EXTERNAL_TIMEOUT|EXTERNAL_KILL|PROCESS_EXIT" \ "$1" 2>/dev/null || true } @@ -40,6 +43,29 @@ ct_step_excerpt() { tail -n 60 "$stepLog" 2>/dev/null || true } +# The driver normally owns verdict files. If the external total-runtime +# backstop or srun kills it first, synthesize a precise verdict here so junit +# never degrades the failure to NO_STATUS. +ct_record_step_exit() { + local name=$1 rc=$2 + local statusFile="$precheckDir/status/$name.status" + # Preserve a driver's specific failure, but never let a stale/premature + # PASS mask the non-zero process exit observed by the gate. + if [ -f "$statusFile" ] && ! grep -q '^PASS' "$statusFile"; then + return 0 + fi + if [ "$rc" -eq 124 ]; then + printf 'EXTERNAL_TIMEOUT %s: step exceeded the %ss total-runtime backstop\n' \ + "$name" "${ctPrecheckTimeout:-900}" > "$statusFile" + elif [ "$rc" -eq 137 ]; then + printf 'EXTERNAL_KILL %s: step received SIGKILL (possibly timeout -k escalation) before the driver wrote a verdict\n' \ + "$name" > "$statusFile" + else + printf 'PROCESS_EXIT %s: srun exited with code %s before the driver wrote a verdict\n' \ + "$name" "$rc" > "$statusFile" + fi +} + # 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. @@ -110,7 +136,7 @@ run_cache_transceiver_precheck() { fi echo "Starting cache transceiver precheck..." precheckDir="$testOutputDir/cache_transceiver_precheck" - mkdir -p "$precheckDir/logs" + mkdir -p "$precheckDir/logs" "$precheckDir/status" # A reused work dir (Slurm requeue reruns this batch script with the same # 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 @@ -120,17 +146,21 @@ run_cache_transceiver_precheck() { # (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 "$precheckDir"/precheck.abort 2>/dev/null || true + rm -f "$precheckDir"/rendezvous/*.addr "$precheckDir"/progress/*.json \ + "$precheckDir"/status/*.status "$precheckDir"/status/*.json \ + "$precheckDir"/precheck.abort 2>/dev/null || true rm -rf "$precheckDir"/csv 2>/dev/null || true precheckPids=() precheckNames=() # 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}" \ + # Scope the launch identity and command to this precheck process. The + # parent shell subsequently launches the real perf workers, so these + # must not replace its DISAGG_SERVING_TYPE/pytestCommand values. + DISAGG_SERVING_TYPE="${role^^}_PRECHECK_$i" \ + pytestCommand="$pytestCmd --server-idx $i" \ + timeout -k 60 "${ctPrecheckTimeout:-900}" \ srun "${srunArgs[@]}" --mpi=pmix --kill-on-bad-exit=1 \ -N "$nodes" \ -w "$nodeList" \ @@ -158,6 +188,7 @@ run_cache_transceiver_precheck() { else rc=$? echo "Precheck step ${precheckNames[$k]} FAILED (exit $rc; 124 = external timeout)" + ct_record_step_exit "${precheckNames[$k]}" "$rc" precheckFailed=1 fi done diff --git a/jenkins/scripts/perf/submit.py b/jenkins/scripts/perf/submit.py index 0a9ad3ec3ab6..e30ecc72ec54 100755 --- a/jenkins/scripts/perf/submit.py +++ b/jenkins/scripts/perf/submit.py @@ -137,12 +137,53 @@ def _test_nodeid(test_line): The pytest node ID from the entry. """ return re.split( - r"\s+(?:XFAIL|SKIP|UNSTABLE|TIMEOUT)(?:\s|$)", + r"\s+(?:XFAIL|SKIP|UNSTABLE|TIMEOUT)(?=[\s(]|$)", test_line, maxsplit=1, )[0] +def _test_marker(test_line): + """Return a test-list execution marker, or ``None`` when absent.""" + line = test_line.partition("#")[0].strip() + match = re.search(r"\s+(XFAIL|SKIP|UNSTABLE|TIMEOUT)(?=[\s(]|$)", line) + return match.group(1) if match else None + + +def selected_test_is_skip_waived(selected_test_line, waives_file, test_prefix=None): + """Whether pytest will skip the selected case before executing its body. + + The CI pipeline merges remote waives into the repository waives file + before invoking this launcher. Mirror the exact-nodeid SKIP decision here + so a skipped test does not run an otherwise unrelated precheck first. + """ + if _test_marker(selected_test_line) == "SKIP": + return True + + selected_nodeid = _test_nodeid(selected_test_line).strip() + try: + waive_lines = _read_test_list_lines(waives_file) + except (FileNotFoundError, ValueError): + return False + + for line in waive_lines: + if _test_marker(line) != "SKIP": + continue + waived_nodeid = _test_nodeid(line).strip() + if waived_nodeid.startswith("full:"): + scope, separator, waived_nodeid = waived_nodeid[5:].partition("/") + if not separator or not test_prefix: + continue + # Match the platform-prefix handling in test_list_parser. SM + # waives require runtime GPU discovery and remain pytest-owned. + platform_prefix = test_prefix.split("-", 1)[0] + if scope.startswith("sm") or platform_prefix not in scope: + continue + if waived_nodeid == selected_nodeid: + return True + return False + + def _load_pytest_split_durations(tokens, llm_src): """Load pytest-split duration data using the launcher's path fallback. @@ -759,6 +800,14 @@ def main(): script_prefix_lines, args.split_group, ) + pytest_tokens = _pytest_command_tokens(script_prefix_lines) + selected_test_skipped = selected_test_is_skip_waived( + selected_test_line, + os.path.join(args.llm_src, "tests", "integration", "test_lists", "waives.txt"), + test_prefix=_pytest_option(pytest_tokens, "--test-prefix"), + ) + if selected_test_skipped: + print("Selected test is SKIP-waived; cache-transceiver precheck will not run") config_yaml, server_name, benchmark_mode, runtime_mode = parse_test_case_name( args.llm_src, selected_test_line, @@ -903,9 +952,22 @@ def main(): # (single owner, shared with the local flow). pcfg = _import_precheck_config(args.llm_src) precheck_enabled = pcfg.precheck_enabled(config) - llm_models_root = ( - _resolve_llm_models_root(script_prefix_lines) if precheck_enabled else None - ) + precheck_will_run = precheck_enabled and not selected_test_skipped + # The model root is only consumed by the precheck (auto KV-cache-manager + # resolution needs the model config). Fail fast only when the precheck + # will actually run; otherwise degrade to a warning so stages whose + # pytestCommand does not carry LLM_MODELS_ROOT inline keep submitting. + llm_models_root = None + if precheck_enabled: + try: + llm_models_root = _resolve_llm_models_root(script_prefix_lines) + except ValueError as e: + if precheck_will_run: + raise + print( + f"WARNING: {e}; " + "cache-transceiver precheck is skipped for this config so continuing" + ) script_prefix_lines.extend( pcfg.precheck_prefix_lines( config, @@ -918,6 +980,7 @@ def main(): ), stage_name=args.stage_name, llm_models_root=llm_models_root, + skip_precheck=selected_test_skipped, ) ) srun_args_lines.extend( diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 6000aecbc1b3..e473d3de352e 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -57,7 +57,9 @@ from .dwdp import DwdpManager from .guided_decoder import GuidedDecoder from .kv_cache_manager_v2 import KVCacheManagerV2 -from .kv_cache_transceiver import AttentionTypeCpp, create_kv_cache_transceiver +from .kv_cache_transceiver import ( + AttentionTypeCpp, create_kv_cache_transceiver, + maybe_enable_fabric_memory_for_python_transceiver) from .llm_request import ExecutorResponse, LlmRequestState from .mamba_cache_manager import (BaseMambaCacheManager, CppMambaHybridCacheManager, @@ -589,28 +591,8 @@ def __init__( self._maybe_enable_fabric_memory_for_python_transceiver() def _maybe_enable_fabric_memory_for_python_transceiver(self) -> None: - """Default TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY=1 for the Python - transceiver on the C++ V1 KV cache manager. - - The Python transceiver (KvCacheTransceiverV2) transfers KV blocks - directly out of the C++ pool, so the pool should be allocated with - fabric memory to enable MNNVL transfers. This must run before any - pool allocation because the C++ env getter caches the value on first - read. Explicit user settings are respected, and platforms without - fabric memory support fall back to standard allocation in C++. - """ - if (self._cache_transceiver_config is None - or self._cache_transceiver_config.backend is None or - self._cache_transceiver_config.transceiver_runtime != "PYTHON"): - return - if not issubclass(self._kv_cache_manager_cls, KVCacheManager): - return - if os.environ.get("TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY") is None: - os.environ["TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY"] = "1" - logger.info( - "Python cache transceiver with C++ KV cache manager detected; " - "defaulting TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY=1 (set it " - "to 0 explicitly to disable)") + maybe_enable_fabric_memory_for_python_transceiver( + self._cache_transceiver_config, self._kv_cache_manager_cls) def _get_model_kv_cache_manager_cls( self, diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index d4868acf0226..65f613337fff 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 from abc import ABC, abstractmethod -from os import getenv +from os import environ, getenv from typing import Any, Dict, List, Optional import tensorrt_llm @@ -30,10 +30,40 @@ _DISABLE_KV_CACHE_TRANSFER_OVERLAP_ENV = "TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP" _DISAGG_LAYERWISE_ENV = "TRTLLM_DISAGG_LAYERWISE" _TRY_ZCOPY_FOR_KV_CACHE_TRANSFER_ENV = "TRTLLM_TRY_ZCOPY_FOR_KVCACHE_TRANSFER" +_KVCACHE_POOL_USE_FABRIC_MEMORY_ENV = "TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY" _SUPPORTED_INFLIGHT_CANCEL_NIXL_BACKEND = "UCX" _disagg_inflight_cancel_enabled_cache: Optional[bool] = None +def maybe_enable_fabric_memory_for_python_transceiver( + cache_transceiver_config: Optional[CacheTransceiverConfig], + kv_cache_manager_cls: type) -> None: + """Default the C++ V1 KV pool to fabric memory for the Python transceiver. + + This must run before any KV pool allocation because the C++ environment + getter caches the value on first read. Explicit user settings are always + respected. + + Args: + cache_transceiver_config: Configuration used to select the cache + transceiver runtime and backend. + kv_cache_manager_cls: KV-cache manager class to check for C++ V1 pool + allocation. + """ + if (cache_transceiver_config is None + or cache_transceiver_config.backend is None + or cache_transceiver_config.transceiver_runtime != "PYTHON"): + return + if not issubclass(kv_cache_manager_cls, KVCacheManager): + return + if getenv(_KVCACHE_POOL_USE_FABRIC_MEMORY_ENV) is None: + environ[_KVCACHE_POOL_USE_FABRIC_MEMORY_ENV] = "1" + logger.info( + "Python cache transceiver with C++ KV cache manager detected; " + f"defaulting {_KVCACHE_POOL_USE_FABRIC_MEMORY_ENV}=1 (set it " + "to 0 explicitly to disable)") + + def is_disagg_inflight_cancel_enabled() -> bool: """Return whether disaggregated in-flight KV transfer cancellation is enabled.""" global _disagg_inflight_cancel_enabled_cache diff --git a/tests/integration/defs/perf/test_perf_sanity.py b/tests/integration/defs/perf/test_perf_sanity.py index c5fb1594c2d0..8111dfee8c0c 100644 --- a/tests/integration/defs/perf/test_perf_sanity.py +++ b/tests/integration/defs/perf/test_perf_sanity.py @@ -469,6 +469,28 @@ def add_host_port_to_cmd(cmd: List[str], host: str, port: int) -> List[str]: return cmd + ["--host", host, "--port", str(port)] +def _run_benchmark_with_log(cmd: List[str], env: Dict[str, str], log_path: str) -> str: + """Run a benchmark while streaming its combined output to an artifact log.""" + benchmark_env = env.copy() + benchmark_env.setdefault("PYTHONUNBUFFERED", "1") + with open(log_path, "wb") as log_file: + result = subprocess.run( + cmd, + env=benchmark_env, + stdout=log_file, + stderr=subprocess.STDOUT, + check=False, + ) + + with open(log_path, "rb") as log_file: + raw_output = log_file.read() + + if result.returncode != 0: + raise subprocess.CalledProcessError(result.returncode, cmd, output=raw_output) + + return raw_output.decode("utf-8", errors="replace") + + class ServerConfig: """Configurations of trtllm-server.""" @@ -1085,7 +1107,10 @@ class AggrTestCmds(NamedTuple): def get_server_logs(self, server_idx) -> List[str]: server_file_path = os.path.join(self.test_output_dir, f"trtllm-serve.{server_idx}.log") - return [server_file_path] + benchmark_logs = sorted( + glob.glob(os.path.join(self.test_output_dir, f"trtllm-benchmark.{server_idx}.*.log")) + ) + return [server_file_path, *benchmark_logs] def run_cmd(self, server_idx: int) -> List[str]: """Run all clients for a server and return outputs. @@ -1155,14 +1180,11 @@ def run_cmd(self, server_idx: int) -> List[str]: client_env = copy.deepcopy(os.environ) if client_config: client_env.update(client_config.to_env()) - output = subprocess.check_output( + output = _run_benchmark_with_log( client_cmd_with_port, - stderr=subprocess.STDOUT, - env=client_env, - ).decode() - - with open(client_file_path, "w") as client_ctx: - client_ctx.write(output) + client_env, + client_file_path, + ) outputs.append(output) else: print_info( @@ -1439,6 +1461,13 @@ def get_server_logs(self, server_idx: int) -> List[str]: os.path.join(self.test_output_dir, f"trtllm-serve.DISAGG_SERVER.{server_idx}.log") ) server_logs.append(os.path.join(self.test_output_dir, "disagg_server.log")) + server_logs.extend( + sorted( + glob.glob( + os.path.join(self.test_output_dir, f"trtllm-benchmark.{server_idx}.*.log") + ) + ) + ) return server_logs @staticmethod @@ -1607,14 +1636,11 @@ def run_cmd(self, server_idx: int) -> List[str]: bench_env = copy.deepcopy(os.environ) if client_config: bench_env.update(client_config.to_env()) - output = subprocess.check_output( + output = _run_benchmark_with_log( client_cmd_with_port, - env=bench_env, - stderr=subprocess.STDOUT, - ).decode() - - with open(benchmark_file_path, "w") as benchmark_ctx: - benchmark_ctx.write(output) + bench_env, + benchmark_file_path, + ) outputs.append(output) if collect_device_step_time: diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index d6c5f42db13f..7a6b4022dc62 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -330,19 +330,15 @@ perf/test_perf_sanity.py::test_e2e[aggr_upload-deepseek_r1_fp8_blackwell-r1_fp8_ perf/test_perf_sanity.py::test_e2e[aggr_upload-dynamo_gpt_oss_120b_fp4_blackwell-gpt_oss_fp4_tep4_adp_cutlass_8k1k] SKIP (https://nvbugs/6374910) perf/test_perf_sanity.py::test_e2e[aggr_upload-glm5_fp4_blackwell-glm5_fp4_tep8_mtp3_8k1k] SKIP (https://nvbugs/6329155) perf/test_perf_sanity.py::test_e2e[aggr_upload-super_ad_blackwell-super_ad_ws1_1k1k] SKIP (https://nvbugs/6153575) -perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb300_deepseek-r1-fp4_128k8k_con256_ctx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL] SKIP (https://nvbugs/6572843) perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb300_deepseek-v4-pro-fp4_8k1k_con180_ctx3_dep4_gen1_dep32_eplb384_mtp3_ccb-NIXL] SKIP (https://nvbugs/6601537) perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb300_deepseek-v4-pro-fp4_8k1k_con4301_ctx12_dep4_gen1_dep8_eplb384_mtp1_ccb-NIXL] SKIP (https://nvbugs/6581075) perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb300_deepseek-v4-pro-fp4_8k1k_con666_ctx6_dep4_gen1_dep16_eplb384_mtp3_ccb-NIXL] SKIP (https://nvbugs/6581075) perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb300_deepseek-v4-pro-fp4_8k1k_con8_ctx1_dep4_gen4_tep8_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6601537) -perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-gb300_glm-5-fp4_8k1k_con1024_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL] SKIP (https://nvbugs/6566777) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL] SKIP (https://nvbugs/6581075) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_deepseek-r1-fp4_128k8k_con256_ctx1_pp4_gen1_dep8_eplb0_mtp1_ccb-NIXL] SKIP (https://nvbugs/6581075) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_deepseek-v4-pro-fp4_8k1k_con180_ctx3_dep4_gen1_dep32_eplb384_mtp3_ccb-NIXL] SKIP (https://nvbugs/6601537) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_deepseek-v4-pro-fp4_8k1k_con4301_ctx12_dep4_gen1_dep8_eplb384_mtp1_ccb-NIXL] SKIP (https://nvbugs/6581075) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_deepseek-v4-pro-fp4_8k1k_con666_ctx6_dep4_gen1_dep16_eplb384_mtp3_ccb-NIXL] SKIP (https://nvbugs/6581075) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_deepseek-v4-pro-fp4_8k1k_con8_ctx1_dep4_gen4_tep8_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6601537) -perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_glm-5-fp4_8k1k_con1024_ctx1_dep2_gen1_dep8_eplb256_mtp1_ccb-NIXL] SKIP (https://nvbugs/6581075) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_glm-5-fp4_8k1k_con1_ctx1_dep2_gen1_tep8_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6581075) perf/test_perf_sanity.py::test_e2e[disagg_upload-gen_only-gb300_glm-5-fp4_8k1k_con512_ctx1_dep2_gen1_dep32_eplb0_mtp3_ccb-NIXL] SKIP (https://nvbugs/6581075) test_e2e.py::test_ptp_quickstart_advanced[Nemotron-Nano-9B-v2-nvfp4-NVIDIA-Nemotron-Nano-9B-v2-NVFP4] SKIP (https://nvbugs/6624972) diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md index 027a4852774c..49f8e19c3597 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/README.md @@ -26,21 +26,19 @@ layers) is read from the real model's `config.json` under `$LLM_MODELS_ROOT`. ## Enabling / disabling -**Off by default** until the gate is validated on the post-merge stages -(the precheck is a launch-script gate, not a pytest case, so it cannot be -waived in `waives.txt` — this default is the waive). To opt in: +**On by default** for every disaggregated perf-sanity test. To opt out: - 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 # 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) + # wireup_timeout_s: 600 # first-rep NIXL agent wire-up allowance + # step_timeout_s: 1200 # default; yaml overrides are capped at 1800 ``` - or globally at launch-script generation time: `TRTLLM_DISAGG_CT_PRECHECK=0` @@ -50,13 +48,19 @@ waived in `waives.txt` — this default is the waive). To opt in: ## 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 +`wireup_timeout_s` (default `min(600, 150 * max world size)`). "Wire-up" is +the one-time exchange of agent and registered-memory metadata and creation of +the peer endpoints before payload transfer. 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 `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. +so the precheck does too. The complete precheck has a topology-independent +20-minute default external limit; exceptional yaml overrides may extend it to +at most 30 minutes. Exceeding the configured limit is a network/environment +failure, even with multiple peers. 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. The Python sender's `kv_transfer_timeout_ms` is also a real request deadline. If block-all returns without every expected request completed—or any rank 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 b5fb567daf43..fb7767e8b9be 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/precheck_config.py @@ -32,6 +32,9 @@ # Optional per-yaml overrides live under a `cache_transceiver_precheck:` block. PRECHECK_DEFAULTS = { + # Enabled centrally for every disaggregated perf-sanity case. Individual + # yamls may opt out for a documented exception. + "enabled": True, # Request lengths to transfer. None -> derived: [1024, benchmark ISL], # clamped to cache_transceiver_config.max_tokens_in_buffer and to # max_request_length below. @@ -50,12 +53,21 @@ "max_concurrent_pairs": 8, # signal.alarm / hang-detector budget for one wave of transfers. "wave_timeout_s": 180, + # KV-pool + native NIXL/UCX agent construction. Python's signal handler + # cannot interrupt a native call, so the HangDetector enforces this + # budget when the extension releases the GIL. Ten minutes accommodates + # the 1-3 minute cold starts observed on loaded Blackwell CI nodes without + # turning a stuck setup into a half-hour gate. + "setup_timeout_s": 600, # 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). + # precheck must too. Keep the no-progress allowance bounded at ten minutes + # rather than multiplying it into every serialized peer wait; a yaml can + # override this for a topology with measured slower remote-MD setup. None + # -> derived: min(600, 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). @@ -64,6 +76,16 @@ "verify_data": True, } +# Timeout layering. The phase watchdog is reset at every checkpoint; the +# external step timeout is a bounded total-runtime backstop for GIL-held +# native hangs. Keep these in the config module so launch and driver budgets +# cannot silently drift apart again. +MAX_WIREUP_TIMEOUT_S = 600 +SERIALIZED_WAIT_GRACE_S = 120 +WATCHDOG_GRACE_S = 60 +DEFAULT_STEP_TIMEOUT_S = 1200 +MAX_STEP_TIMEOUT_S = 1800 + # Fallback KV shape for dry-runs and explicitly selected manager versions when # the model directory cannot be resolved. Manager-version "auto" resolution # fails fast instead of silently pairing this shape with V1. @@ -118,44 +140,102 @@ def _spec_nextn(side): def wireup_timeout_s(max_world): """First-rep NIXL agent wire-up allowance (see PRECHECK_DEFAULTS).""" - return min(1800, 150 * int(max_world)) + return min(MAX_WIREUP_TIMEOUT_S, 150 * int(max_world)) + + +def peer_progress_timeout_s(setup_timeout_s, rendezvous_timeout_s, wave_timeout_s, wireup_s): + """Maximum no-progress interval while a peer is queued behind sessions. + + Serialized peer waits are refreshed whenever the target ctx or gen + instance advances to another phase or wave. The budget therefore covers one + legitimate active phase, not the cumulative duration of every earlier + session, which would grow to hours on 6/12-ctx topologies. + """ + longest_active_phase_s = max( + int(setup_timeout_s), + int(rendezvous_timeout_s), + int(wave_timeout_s) + int(wireup_s), + ) + return longest_active_phase_s + SERIALIZED_WAIT_GRACE_S -def default_step_timeout_s(max_world): +def timeout_budget(cfg, max_world): + """Resolve aligned phase, watchdog, and external-step timeout budgets. + + The default external limit is intentionally topology-independent at twenty + minutes; yaml overrides may extend exceptional cases to at most thirty + minutes. Adding peers alone must not inflate the default. + """ + knobs = dict(PRECHECK_DEFAULTS) + knobs.update(cfg.get("cache_transceiver_precheck", {}) or {}) + wireup_s = int( + knobs["wireup_timeout_s"] + if knobs["wireup_timeout_s"] is not None + else wireup_timeout_s(max_world) + ) + peer_progress_s = peer_progress_timeout_s( + knobs["setup_timeout_s"], + knobs["rendezvous_timeout_s"], + knobs["wave_timeout_s"], + wireup_s, + ) + longest_phase_s = max( + int(knobs["setup_timeout_s"]), + int(knobs["wave_timeout_s"]) + wireup_s, + peer_progress_s, + ) + watchdog_s = longest_phase_s + WATCHDOG_GRACE_S + step_s = int(knobs.get("step_timeout_s", DEFAULT_STEP_TIMEOUT_S)) + if step_s > MAX_STEP_TIMEOUT_S: + raise ValueError( + "cache_transceiver_precheck.step_timeout_s must not exceed the " + f"global health-check limit ({MAX_STEP_TIMEOUT_S}s), got {step_s}s" + ) + if step_s <= watchdog_s: + raise ValueError( + "cache_transceiver_precheck.step_timeout_s must exceed the longest " + f"phase watchdog ({watchdog_s}s), got {step_s}s" + ) + return { + "wireup_timeout_s": wireup_s, + "longest_phase_timeout_s": longest_phase_s, + "watchdog_timeout_s": watchdog_s, + "step_timeout_s": step_s, + } + + +def default_step_timeout_s(cfg, max_world): """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. + Used by precheck_prefix_lines(), which is shared by both launch + generators. It always exceeds the longest internal phase watchdog while + retaining the 20-minute default and 30-minute maximum health-check limits. """ - return 900 + wireup_timeout_s(max_world) + return timeout_budget(cfg, max_world)["step_timeout_s"] def precheck_enabled(cfg): - """Resolve the shared yaml/environment enable policy.""" - knobs = cfg.get("cache_transceiver_precheck", {}) or {} - # Off by default until the gate is validated on the post-merge stages - # (the precheck is a launch-script gate, not a pytest case, so it cannot - # be waived in waives.txt — this default is the waive). Yaml opts in per - # test; the env var (when set) overrides 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. + """Resolve the precheck enable/kill-switch policy for a test config. + + On by default; yaml opts out per test; the env var (when set) overrides + 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", False)) - return enabled + return True + if val in ("0", "false", "off", "no"): + return False + raise ValueError( + "TRTLLM_DISAGG_CT_PRECHECK must be a boolean " + f"(1/0/true/false/on/off/yes/no), got {env!r}" + ) + knobs = dict(PRECHECK_DEFAULTS) + knobs.update(cfg.get("cache_transceiver_precheck", {}) or {}) + return bool(knobs["enabled"]) def precheck_prefix_lines( @@ -166,6 +246,7 @@ def precheck_prefix_lines( max_world: int, stage_name: str = "", llm_models_root: str | None = None, + skip_precheck: bool = False, ) -> list[str]: """Launch-script export lines wiring the precheck gate. @@ -182,6 +263,7 @@ def precheck_prefix_lines( max_world: Largest role world size, used to derive the step timeout. stage_name: Optional stage name for the synthetic JUnit report. llm_models_root: Model-root path exported when the precheck is enabled. + skip_precheck: Disable the gate for this config even when enabled. Returns: Generated launch-script export lines shared by both submit modules. @@ -189,25 +271,32 @@ def precheck_prefix_lines( Raises: ValueError: If the precheck is enabled without a nonempty model root. """ - knobs = cfg.get("cache_transceiver_precheck", {}) or {} - enabled = precheck_enabled(cfg) + # A selected pytest case that is SKIP-waived must not run its precheck. + # This takes precedence over the global force-enable knob because there + # will be no corresponding test execution to gate. + enabled = precheck_enabled(cfg) and not skip_precheck 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" ) + model_root_env = f"LLM_MODELS_ROOT={shlex.quote(llm_models_root)}" if llm_models_root else "" 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)))}", + # Single-source timeout derivation keeps the external backstop above + # every phase watchdog without allowing topology-linear inflation. + f"export ctPrecheckTimeout={default_step_timeout_s(cfg, 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"', + # Quote the complete assignment as an export value. The launch script + # expands it into pytestCommand*Precheck, whose later eval interprets + # the inner shlex-quoted model path. + f"export ctPrecheckModelRootEnv={shlex.quote(model_root_env)}", + f'export pytestCommandCTXPrecheck="{ucx_tls_cmd} $ctPrecheckModelRootEnv ' + f'$CTX_WORKER_ENV_VARS $PYTEST_COMMON_VARS {cmd} --role ctx"', + f'export pytestCommandGENPrecheck="{ucx_tls_cmd} $ctPrecheckModelRootEnv ' + f'$GEN_WORKER_ENV_VARS $PYTEST_COMMON_VARS {cmd} --role gen"', ] if enabled: if not llm_models_root: @@ -315,6 +404,15 @@ def resolve_plan(cfg, benchmark_mode="e2e"): n_pairs = max(ctx["dp_size"], gen["dp_size"], 1) wave_size = max(1, min(n_pairs, int(knobs["max_concurrent_pairs"]))) + setup_timeout_s = int(knobs["setup_timeout_s"]) + wave_timeout_s = int(knobs["wave_timeout_s"]) + wireup_s = int( + knobs["wireup_timeout_s"] + if knobs["wireup_timeout_s"] is not None + else wireup_timeout_s(max(ctx["world_size"], gen["world_size"])) + ) + rendezvous_timeout_s = int(knobs["rendezvous_timeout_s"]) + plan = { "skip": False, "num_ctx_servers": num_ctx_servers, @@ -328,13 +426,16 @@ def resolve_plan(cfg, benchmark_mode="e2e"): "warmup_requests": int(knobs["warmup_requests"]), "n_pairs": n_pairs, "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 - else wireup_timeout_s(max(ctx["world_size"], gen["world_size"])) + "wave_timeout_s": wave_timeout_s, + "setup_timeout_s": setup_timeout_s, + "wireup_timeout_s": wireup_s, + "peer_progress_timeout_s": peer_progress_timeout_s( + setup_timeout_s, + rendezvous_timeout_s, + wave_timeout_s, + wireup_s, ), - "rendezvous_timeout_s": int(knobs["rendezvous_timeout_s"]), + "rendezvous_timeout_s": rendezvous_timeout_s, "verify_data": bool(knobs["verify_data"]), } for role, side, xcvr in (("ctx", ctx_side, ctx_xcvr), ("gen", gen_side, gen_xcvr)): 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 8b9e1eda4846..9a2c7c5b131f 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -85,6 +85,7 @@ # 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 +CONTROL_POLL_INTERVAL_MS = 5_000 ABORT_COORDINATION_TIMEOUT_S = 2.0 @@ -231,7 +232,10 @@ def load_internal_apis(): from tensorrt_llm._torch.models.modeling_utils import get_registered_model_class 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.kv_cache_transceiver import ( + create_kv_cache_transceiver, + maybe_enable_fabric_memory_for_python_transceiver, + ) from tensorrt_llm._torch.pyexecutor.llm_request import ( LlmRequest, LlmRequestState, @@ -265,6 +269,9 @@ def load_internal_apis(): KVCacheManager=KVCacheManager, KVCacheManagerV2=KVCacheManagerV2, create_kv_cache_transceiver=create_kv_cache_transceiver, + maybe_enable_fabric_memory_for_python_transceiver=( + maybe_enable_fabric_memory_for_python_transceiver + ), LlmRequest=LlmRequest, LlmRequestState=LlmRequestState, LlmRequestType=LlmRequestType, @@ -368,7 +375,9 @@ class and adopt its get_preferred_kv_cache_manager_version() value 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. + Resolution errors propagate so the precheck cannot silently exercise a + different runtime/cache-manager combination from serving. Returns the + effective use_v2 bool. """ import types @@ -640,6 +649,55 @@ def wait_for_addr(path, timeout_s): raise _Timeout(f"rendezvous file {path} not published within {timeout_s}s") +def peer_progress_path(work_dir, role, server_idx): + """Shared progress marker for one ctx or gen instance.""" + return os.path.join(work_dir, "progress", f"{role}_{server_idx}.json") + + +def publish_peer_progress(runner, phase): + """Best-effort atomic phase marker used by queued peer instances. + + Only instance leaders write. A new sequence value means the target peer is + advancing, so queued hello/bye waits may refresh their no-progress + watchdog without budgeting earlier serialized sessions cumulatively. + """ + if not runner.is_leader: + return + path = peer_progress_path(runner.work_dir, runner.role, runner.server_idx) + try: + runner._precheck_progress_seq = getattr(runner, "_precheck_progress_seq", 0) + 1 + os.makedirs(os.path.dirname(path), exist_ok=True) + tmp = f"{path}.tmp.{os.getpid()}" + with open(tmp, "w") as f: + json.dump( + { + "job": run_token(), + "phase": str(phase)[:400], + "seq": runner._precheck_progress_seq, + }, + f, + ) + os.replace(tmp, path) + except OSError: + # Missing progress only removes watchdog refreshes; the normal bounded + # timeout and external gate remain authoritative. + pass + + +def read_peer_progress(work_dir, role, server_idx): + """Return this run's peer progress sequence, or None if absent/stale.""" + try: + with open(peer_progress_path(work_dir, role, server_idx)) 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("seq") + + def abort_flag_path(work_dir): return os.path.join(work_dir, "precheck.abort") @@ -925,6 +983,8 @@ def setup(self, kv_shape, max_req_len): "(the C++ transceiver only supports the V1 manager)" ) + manager_cls = api.KVCacheManagerV2 if self.use_v2 else api.KVCacheManager + api.maybe_enable_fabric_memory_for_python_transceiver(cache_cfg, manager_cls) self.kvm = build_kv_cache_manager( kv_shape, self.plan, self.side, self.mapping, max_req_len, self.use_v2 ) @@ -1032,6 +1092,12 @@ def ctx_finish_wave(self, reqs): completed, failed = self.xcvr.check_context_transfer_status(None) # block-all completed_rids = set(completed) failed_rids = set(failed) + # Python's block-all wait is bounded by the kv_transfer_timeout + # deadline; a request still nonterminal on return exceeded it. + # The precheck payloads are far smaller than real requests, so the + # gate classifies that as a transfer failure — while retaining the + # KV pages, because deadline expiry does not prove the peer + # quiesced (the fatal path below skips ordinary teardown). missing = [ p for p, req in reqs.items() @@ -1123,6 +1189,10 @@ def gen_run_wave(self, peer_idx, li, req_len, rep, wave, params_by_pair): failed_rids = set(failed) cancelled_rids = {req.py_request_id for req in cancelled} expected_rids = {req.py_request_id for req in reqs.values()} + # A request still nonterminal after the block-all deadline + # (missing) is a gate failure like failed/cancelled; the fatal + # path retains its KV pages because the peer may not have + # quiesced. missing_rids = expected_rids - completed_rids - failed_rids - cancelled_rids if failed_rids or cancelled_rids or missing_rids: raise _TransferError( @@ -1229,6 +1299,107 @@ def _leader_send_recv(self, sock, obj, key): raise _TransferError(f"ZMQ control channel failed: {err}") return reply + def _leader_send_recv_with_progress( + self, + sock, + obj, + key, + *, + peer_role, + peer_idx, + what, + timeout_s, + arm, + ): + """REQ round-trip refreshed by progress from the queued target peer.""" + timeout_s = int(timeout_s) + arm(what, seconds=timeout_s, python_alarm=False) + send_err = None + if self.is_leader: + try: + sock.send(pack_msg(obj, key)) + except Exception as e: # noqa: BLE001 - shared via bcast below + send_err = repr(e) + send_err = self.comm.bcast(send_err, root=0) + if send_err: + raise _TransferError(f"ZMQ control send failed: {send_err}") + + return _collective_recv_with_progress( + runner=self, + sock=sock, + key=key, + peer_role=peer_role, + peer_idx=peer_idx, + what=what, + timeout_s=timeout_s, + arm=arm, + refresh_from_peer_progress=True, + ) + + +def _collective_recv_with_progress( + runner, + sock, + key, + peer_role, + peer_idx, + what, + timeout_s, + arm, + *, + refresh_from_peer_progress, +): + """Collectively receive control; target-peer progress resets the deadline. + + The caller arms the phase watchdog before entering this loop. Requiring an + explicit refresh policy keeps active transfer waves on a hard deadline. + """ + timeout_s = int(timeout_s) + deadline = time.monotonic() + timeout_s + last_progress = ( + read_peer_progress(runner.work_dir, peer_role, peer_idx) + if runner.is_leader and refresh_from_peer_progress + else None + ) + + while True: + event = None + if runner.is_leader: + zmq, _ = runner._zmq() + try: + event = ("message", unpack_msg(sock.recv(), key)) + except zmq.Again: + progress = ( + read_peer_progress(runner.work_dir, peer_role, peer_idx) + if refresh_from_peer_progress + else None + ) + if progress is not None and progress != last_progress: + event = ("progress", progress) + elif time.monotonic() >= deadline: + event = ("timeout", None) + else: + event = ("poll", None) + except Exception as e: # noqa: BLE001 - shared via bcast below + event = ("error", repr(e)) + + kind, payload = runner.comm.bcast(event, root=0) + if kind == "message": + return payload + if kind == "error": + raise _TransferError(f"ZMQ recv from {peer_role}_{peer_idx} failed: {payload}") + if kind == "timeout": + raise _Timeout(f"{what} made no progress for {timeout_s}s") + if kind == "progress": + last_progress = payload + deadline = time.monotonic() + timeout_s + arm( + what, + seconds=timeout_s, + python_alarm=False, + publish_progress=False, + ) + def _schedule(plan): """Deterministic (li, req_len, rep, wave) schedule both sides iterate.""" @@ -1241,15 +1412,15 @@ def _schedule(plan): return out -def hello_timeout_s(plan, num_peers): - """Timeout budget for session handshakes. +def hello_timeout_s(plan): + """No-progress budget for serialized hello/bye waits. - 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 (including - the peer's first-rep wire-up). + The target peer's progress marker refreshes this wait after each + phase/wave, so the budget covers one slow active phase rather than + incorrectly using either side's local peer count or multiplying every + earlier session. """ - return plan["rendezvous_timeout_s"] + num_peers * (300 + plan["wireup_timeout_s"]) + return plan["peer_progress_timeout_s"] def wave_timeout_s(plan, li, rep): @@ -1262,6 +1433,38 @@ def wave_timeout_s(plan, li, rep): return plan["wave_timeout_s"] + extra +def _recv_ctx_control( + runner, + sock, + key, + peer_idx, + what, + timeout_s, + arm, + refresh_from_gen_progress=False, +): + """Collectively receive one ctx-side control message with progress refresh. + + Only the leader owns the ZMQ socket; every poll result is broadcast so all + ranks execute the same watchdog checkpoints and failure branch. For queued + hello/bye waits, observed progress from the target gen resets the deadline. + Active wave waits do not refresh from unrelated progress. + """ + timeout_s = int(timeout_s) + arm(what, seconds=timeout_s, python_alarm=False) + return _collective_recv_with_progress( + runner=runner, + sock=sock, + key=key, + peer_role="gen", + peer_idx=peer_idx, + what=what, + timeout_s=timeout_s, + arm=arm, + refresh_from_peer_progress=refresh_from_gen_progress, + ) + + # --------------------------------------------------------------------------- # # ctx / gen session loops # @@ -1287,39 +1490,36 @@ def wave_timeout_s(plan, li, rep): 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 - - def leader_recv(): - msg, err, timed_out = None, None, False - if runner.is_leader: - try: - msg = unpack_msg(sock.recv(), key) - except _Timeout as e: - timed_out = True - err = repr(e) - except Exception as e: # noqa: BLE001 - err = repr(e) - timed_out, err, msg = comm.bcast((timed_out, err, msg), root=0) - if timed_out: - raise _Timeout(err) - 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(pack_msg(obj, key)) - arm(f"hello gen_{peer_idx}", seconds=hello_timeout_s(plan, runner.side["num_peers"])) - msg = leader_recv() + msg = _recv_ctx_control( + runner, + sock, + key, + peer_idx, + f"hello gen_{peer_idx}", + hello_timeout_s(plan), + arm, + refresh_from_gen_progress=True, + ) 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, wave in _schedule(plan): - arm(f"gen_{peer_idx} len={req_len} rep={rep}", seconds=wave_timeout_s(plan, li, rep)) - msg = leader_recv() + msg = _recv_ctx_control( + runner, + sock, + key, + peer_idx, + f"gen_{peer_idx} len={req_len} rep={rep}", + wave_timeout_s(plan, li, rep), + arm, + ) 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). @@ -1370,9 +1570,18 @@ def leader_reply(obj): # 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() + # therefore span the gen's remaining sessions: use the same progress-aware + # no-progress budget as the initial handshake. + msg = _recv_ctx_control( + runner, + sock, + key, + peer_idx, + f"bye gen_{peer_idx}", + hello_timeout_s(plan), + arm, + refresh_from_gen_progress=True, + ) if msg[0] != "done": raise _TransferError(f"expected done from gen_{peer_idx}, got {msg[:1]}") leader_reply(("bye", {})) @@ -1389,7 +1598,7 @@ def _gen_open_session(runner, peer_idx, arm): """ plan = runner.plan comm = runner.comm - hello_s = hello_timeout_s(plan, runner.side["num_peers"]) + hello_s = hello_timeout_s(plan) sock, key, err = None, None, None arm(f"rendezvous ctx_{peer_idx}", seconds=hello_s) @@ -1405,7 +1614,9 @@ def _gen_open_session(runner, peer_idx, arm): zmq, zctx = runner._zmq() sock = zctx.socket(zmq.REQ) sock.setsockopt(zmq.LINGER, 0) - sock.setsockopt(zmq.RCVTIMEO, hello_s * 1000) + # Poll while queued behind another gen session so ctx progress can + # refresh the no-progress deadline. + sock.setsockopt(zmq.RCVTIMEO, CONTROL_POLL_INTERVAL_MS) sock.connect(f"tcp://{addr['host']}:{addr['port']}") except Exception as e: # noqa: BLE001 - shared via bcast below err = repr(e) @@ -1414,16 +1625,23 @@ def _gen_open_session(runner, peer_idx, arm): 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( + reply = runner._leader_send_recv_with_progress( sock, ("hello", {"gen_idx": runner.server_idx, "fingerprint": plan["fingerprint"]}), key, + peer_role="ctx", + peer_idx=peer_idx, + what=f"hello ctx_{peer_idx}", + timeout_s=hello_s, + arm=arm, ) 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]}") + if runner.is_leader: + zmq, _ = runner._zmq() + sock.setsockopt(zmq.RCVTIMEO, hello_s * 1000) return sock, key except BaseException: if sock is not None: @@ -1511,7 +1729,7 @@ def gen_abort_peer(runner, peer_idx, reason, arm, disarm): 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"])) + arm(f"abort ctx_{peer_idx}", seconds=hello_timeout_s(runner.plan)) runner._leader_send_recv(sock, ("abort", f"peer fail-fast: {reason}"), key) disarm() finally: @@ -1598,19 +1816,27 @@ def _on_hang(): finally: 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. + # `arm` updates this timeout before each checkpoint. This initial value is + # replaced before the first task is scheduled. hang_detector = load_internal_apis().HangDetector( - timeout=hello_timeout_s(plan, runner.side["num_peers"]) + plan["wave_timeout_s"] + 60, + timeout=plan["setup_timeout_s"] + pcfg.WATCHDOG_GRACE_S, on_detected=_on_hang, ) hang_detector.start() - def arm(what, seconds=None): + def arm(what, seconds=None, python_alarm=True, publish_progress=True): current_cell["what"] = what - signal.alarm(seconds or plan["wave_timeout_s"]) + phase_timeout_s = int(plan["wave_timeout_s"] if seconds is None else seconds) + signal.alarm(phase_timeout_s if python_alarm else 0) + # A Python alarm cannot interrupt a native extension. Re-arm the + # thread-based detector with the same phase budget plus a small grace, + # rather than one global timeout inflated by unrelated later phases. + hang_detector.timeout = phase_timeout_s + pcfg.WATCHDOG_GRACE_S hang_detector.checkpoint() + # Progress-derived watchdog refreshes must not echo a marker back to + # the peer: reciprocal echoes could keep a genuinely stuck pair alive. + if publish_progress: + publish_peer_progress(runner, what) def disarm(): signal.alarm(0) @@ -1710,8 +1936,9 @@ def _serve_gen_peers(runner, plan, arm, disarm, record_peer_failure): 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) + # Poll so queued ctx sessions can observe target-gen progress and + # refresh a no-progress watchdog without cumulative peer budgets. + s.setsockopt(zmq.RCVTIMEO, CONTROL_POLL_INTERVAL_MS) port = s.bind_to_random_port("tcp://*") keys[gj] = secrets.token_bytes(32) write_addr( @@ -1840,7 +2067,7 @@ def main(argv=None): # --- setup: KV pool + transceiver (same config as the real test) --------- setup_err = None try: - arm("kv pool + transceiver setup") + arm("kv pool + transceiver setup", seconds=plan["setup_timeout_s"]) runner.setup(kv_shape, max_req_len=max(plan["request_lengths"])) disarm() except Exception as e: # noqa: BLE001 - recorded and gated below diff --git a/tests/unittest/others/test_cache_transceiver_precheck_config.py b/tests/unittest/others/test_cache_transceiver_precheck_config.py index 55c9cb7a6580..00d224f5e6a6 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_config.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_config.py @@ -19,8 +19,10 @@ import json import os +import shlex import subprocess import sys +import time import types import pytest @@ -474,13 +476,121 @@ 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) + assert plan["wireup_timeout_s"] == 600 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_timeout_budget_is_bounded_and_layered(): + one_peer = pcfg.timeout_budget(_disagg_yaml(), max_world=16) + assert one_peer == { + "wireup_timeout_s": 600, + "longest_phase_timeout_s": 900, + "watchdog_timeout_s": 960, + "step_timeout_s": 1200, + } + + cfg = _disagg_yaml(hardware={"gpus_per_node": 4, "num_ctx_servers": 12, "num_gen_servers": 1}) + multi_peer = pcfg.timeout_budget(cfg, max_world=32) + # Even 12 serialized ctx peers retain the 20-minute default health limit. + # Progress refreshes protect legitimate per-peer phases without scaling + # the external backstop with topology. + assert multi_peer["longest_phase_timeout_s"] == 900 + assert multi_peer["step_timeout_s"] == 1200 + + cfg["cache_transceiver_precheck"] = {"step_timeout_s": 900} + with pytest.raises(ValueError, match="must exceed the longest phase watchdog"): + pcfg.timeout_budget(cfg, max_world=32) + + cfg["cache_transceiver_precheck"] = {"step_timeout_s": 1800} + assert pcfg.timeout_budget(cfg, max_world=32)["step_timeout_s"] == 1800 + + cfg["cache_transceiver_precheck"] = {"step_timeout_s": 1801} + with pytest.raises(ValueError, match="must not exceed the global health-check limit"): + pcfg.timeout_budget(cfg, max_world=32) + + +@pytest.mark.parametrize("model_root", ["/models with spaces", "/models/o'hare"]) +def test_precheck_commands_propagate_model_root(monkeypatch, model_root): + # CI provides the model root inside its inbound pytest command, not in the + # environment of the Python process that generates the launch script. + monkeypatch.delenv("LLM_MODELS_ROOT", raising=False) + monkeypatch.delenv("TRTLLM_DISAGG_CT_PRECHECK", raising=False) + lines = pcfg.precheck_prefix_lines( + {}, + "e2e", + "$config", + "unset UCX_TLS &&", + max_world=8, + llm_models_root=model_root, + ) + shell_script = "\n".join( + [ + "CTX_WORKER_ENV_VARS=", + "GEN_WORKER_ENV_VARS=", + "PYTEST_COMMON_VARS=", + "llmSrcNode=/repo", + "testOutputDir=/tmp/output", + "config=/tmp/config.yaml", + *lines, + "printf '%s\\n' \"$pytestCommandCTXPrecheck\"", + "printf '%s\\n' \"$pytestCommandGENPrecheck\"", + ] + ) + + result = subprocess.run( + ["bash"], input=shell_script, capture_output=True, check=True, text=True + ) + commands = result.stdout.splitlines() + + assert len(commands) == 2 + for command in commands: + tokens = shlex.split(command) + assignment = f"LLM_MODELS_ROOT={model_root}" + assert assignment in tokens + assert tokens.index(assignment) < tokens.index("python3") + + +def test_precheck_commands_split_pytest_common_vars(monkeypatch): + # $PYTEST_COMMON_VARS is spliced unquoted on purpose: bash word splitting + # must yield separate K=V env-assignment tokens ahead of the executable. + # Values containing spaces are unsupported by design — this pins the + # expected splitting behavior. + monkeypatch.delenv("LLM_MODELS_ROOT", raising=False) + monkeypatch.delenv("TRTLLM_DISAGG_CT_PRECHECK", raising=False) + lines = pcfg.precheck_prefix_lines( + {}, + "e2e", + "$config", + "unset UCX_TLS &&", + max_world=8, + llm_models_root="/models", + ) + shell_script = "\n".join( + [ + "CTX_WORKER_ENV_VARS=", + "GEN_WORKER_ENV_VARS=", + 'PYTEST_COMMON_VARS="FOO=1 BAR=two"', + "llmSrcNode=/repo", + "testOutputDir=/tmp/output", + "config=/tmp/config.yaml", + *lines, + "printf '%s\\n' \"$pytestCommandCTXPrecheck\"", + ] + ) + + result = subprocess.run( + ["bash"], input=shell_script, capture_output=True, check=True, text=True + ) + tokens = shlex.split(result.stdout.splitlines()[0]) + + python_index = tokens.index("python3") + for assignment in ("FOO=1", "BAR=two"): + assert tokens.index(assignment) < python_index + + def _enabled_line(cfg): lines = pcfg.precheck_prefix_lines( cfg, @@ -531,7 +641,7 @@ def test_disabled_precheck_does_not_require_or_export_model_root(monkeypatch): monkeypatch.delenv("TRTLLM_DISAGG_CT_PRECHECK", raising=False) lines = pcfg.precheck_prefix_lines( - _disagg_yaml(), + _disagg_yaml(cache_transceiver_precheck={"enabled": False}), "e2e", "$config", "unset UCX_TLS &&", @@ -557,6 +667,32 @@ def test_enabled_precheck_requires_model_root(monkeypatch, llm_models_root): ) +def test_precheck_enabled_helper(monkeypatch): + # submit.py consults this helper to decide whether a missing model root is + # fatal — it must mirror the policy encoded in ctPrecheckEnabled. + monkeypatch.delenv("TRTLLM_DISAGG_CT_PRECHECK", raising=False) + assert pcfg.PRECHECK_DEFAULTS["enabled"] is True + assert pcfg.precheck_enabled({}) is True + assert pcfg.precheck_enabled({"cache_transceiver_precheck": {"enabled": False}}) is False + monkeypatch.setenv("TRTLLM_DISAGG_CT_PRECHECK", "0") + assert pcfg.precheck_enabled({}) is False + monkeypatch.setenv("TRTLLM_DISAGG_CT_PRECHECK", "true") + assert pcfg.precheck_enabled({"cache_transceiver_precheck": {"enabled": False}}) is True + + +def test_skip_waived_case_overrides_force_enable(monkeypatch): + monkeypatch.setenv("TRTLLM_DISAGG_CT_PRECHECK", "1") + lines = pcfg.precheck_prefix_lines( + {}, + "e2e", + "$c", + "unset &&", + max_world=8, + skip_precheck=True, + ) + assert next(x for x in lines if x.startswith("export ctPrecheckEnabled")).endswith("=0") + + def test_precheck_env_kill_switch_truthy(monkeypatch): """The TRTLLM_DISAGG_CT_PRECHECK kill switch parses the usual boolean spellings. @@ -566,7 +702,7 @@ def test_precheck_env_kill_switch_truthy(monkeypatch): cfg = {"cache_transceiver_precheck": {"enabled": True}} monkeypatch.delenv("TRTLLM_DISAGG_CT_PRECHECK", raising=False) assert _enabled_line(cfg).endswith("=1") # yaml opt-in - assert _enabled_line({}).endswith("=0") # off by default (waived) + assert _enabled_line({}).endswith("=1") # on by default for v in ("1", "true", "on", "YES", " True "): monkeypatch.setenv("TRTLLM_DISAGG_CT_PRECHECK", v) assert _enabled_line(cfg).endswith("=1"), v @@ -608,6 +744,114 @@ def test_gate_library_content(tmp_path): pcfg.gate_library_content("/nowhere/draft.sh", str(tmp_path / "empty")) +@pytest.mark.parametrize( + ("exit_code", "expected_verdict"), + ((124, "EXTERNAL_TIMEOUT"), (137, "EXTERNAL_KILL")), +) +def test_gate_records_external_timeout_verdict(tmp_path, exit_code, expected_verdict): + gate = os.path.join( + os.path.dirname(_PRECHECK_DIR), + "..", + "..", + "..", + "jenkins", + "scripts", + "perf", + "disaggregated", + "slurm_ct_precheck_gate.sh", + ) + gate = os.path.abspath(gate) + shell_script = f""" +source {shlex.quote(gate)} +timeout() {{ return {exit_code}; }} +sleep() {{ :; }} +cleanup_on_failure() {{ :; }} +ctPrecheckEnabled=1 +ctPrecheckTimeout=1200 +testOutputDir={shlex.quote(str(tmp_path / "output"))} +jobWorkspace={shlex.quote(str(tmp_path / "workspace"))} +mkdir -p "$jobWorkspace" +numGenServers=1 +numCtxServers=1 +nodesPerGenServer=1 +nodesPerCtxServer=1 +gpusPerNodePerGenServer=1 +gpusPerNodePerCtxServer=1 +genNodeLists=(gen-node) +ctxNodeLists=(ctx-node) +srunArgs=() +pytestCommandGENPrecheck=gen-command +pytestCommandCTXPrecheck=ctx-command +precheckRunScript=/unused +run_cache_transceiver_precheck +""" + subprocess.run(["bash"], input=shell_script, capture_output=True, check=True, text=True) + + status_dir = tmp_path / "output" / "cache_transceiver_precheck" / "status" + for name in ("gen_0", "ctx_0"): + verdict = (status_dir / f"{name}.status").read_text() + assert verdict.startswith(f"{expected_verdict} {name}:") + if exit_code == 124: + assert "1200s total-runtime backstop" in verdict + else: + assert "possibly timeout -k escalation" in verdict + junit = (tmp_path / "workspace" / "results-ct-precheck.xml").read_text() + assert expected_verdict in junit + assert "NO_STATUS" not in junit + + +def test_gate_scopes_precheck_launch_environment(tmp_path): + gate = os.path.join( + os.path.dirname(_PRECHECK_DIR), + "..", + "..", + "..", + "jenkins", + "scripts", + "perf", + "disaggregated", + "slurm_ct_precheck_gate.sh", + ) + gate = os.path.abspath(gate) + shell_script = f""" +source {shlex.quote(gate)} +timeout() {{ + case "$DISAGG_SERVING_TYPE:$pytestCommand" in + "GEN_PRECHECK_0:gen-command --server-idx 0"|\ + "CTX_PRECHECK_0:ctx-command --server-idx 0") return 0 ;; + *) return 99 ;; + esac +}} +sleep() {{ :; }} +cleanup_on_failure() {{ return 98; }} +ctPrecheckEnabled=1 +ctPrecheckTimeout=1200 +testOutputDir={shlex.quote(str(tmp_path / "output"))} +jobWorkspace={shlex.quote(str(tmp_path / "workspace"))} +mkdir -p "$jobWorkspace" +numGenServers=1 +numCtxServers=1 +nodesPerGenServer=1 +nodesPerCtxServer=1 +gpusPerNodePerGenServer=1 +gpusPerNodePerCtxServer=1 +genNodeLists=(gen-node) +ctxNodeLists=(ctx-node) +srunArgs=() +pytestCommandGENPrecheck=gen-command +pytestCommandCTXPrecheck=ctx-command +precheckRunScript=/unused +DISAGG_SERVING_TYPE=REAL_PERF_PARENT +pytestCommand=real-perf-command +run_cache_transceiver_precheck +printf '%s\n%s\n' "$DISAGG_SERVING_TYPE" "$pytestCommand" +""" + result = subprocess.run( + ["bash"], input=shell_script, capture_output=True, check=True, text=True + ) + assert result.stdout.splitlines()[-2:] == ["REAL_PERF_PARENT", "real-perf-command"] + + def test_rid_tags_dense_within_session(): """Rids must be dense within a (ctx, gen) session. @@ -636,7 +880,7 @@ def session_rids(ctx_idx, gen_idx): class TestMultiPeerOrchestration: - """CPU-only end-to-end run of the 2-ctx x 1-gen session protocol. + """CPU-only end-to-end runs of the multi-peer session protocol. Exercises the exact multi-instance logic of the hardware "B" topology: real ZMQ sockets + HMAC frames + StatusRecorder + rendezvous files via @@ -668,7 +912,16 @@ class _FakeParams: ctx_dp_rank = 0 disagg_info_endpoint = None - def _mk_runner(self, role, server_idx, plan, work_dir, monkeypatch): + def _mk_runner( + self, + role, + server_idx, + plan, + work_dir, + monkeypatch, + fail_ctx=False, + wave_delay_s=0, + ): import sys import types @@ -685,6 +938,10 @@ def _mk_runner(self, role, server_idx, plan, work_dir, monkeypatch): calls = {"waves": 0} def ctx_run_wave(peer_idx, li, req_len, rep, wave): + if fail_ctx: + raise rp._TransferError("injected ctx failure") + if wave_delay_s: + time.sleep(wave_delay_s) calls["waves"] += 1 return {p: self._FakeParams() for p in wave}, {} @@ -694,7 +951,16 @@ def ctx_run_wave(peer_idx, li, req_len, rep, wave): runner._calls = calls return runner - def _run(self, tmp_path, monkeypatch, fail_peer_idx=None): + def _run( + self, + tmp_path, + monkeypatch, + fail_peer_idx=None, + fail_ctx_idx=None, + num_ctx_servers=2, + first_ctx_wave_delay_s=0, + peer_progress_timeout_s=None, + ): import threading monkeypatch.setenv("SLURM_JOB_ID", "777") @@ -702,7 +968,11 @@ def _run(self, tmp_path, monkeypatch, fail_peer_idx=None): # 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}, + hardware={ + "gpus_per_node": 4, + "num_ctx_servers": num_ctx_servers, + "num_gen_servers": 1, + }, cache_transceiver_precheck={ "request_lengths": [32], "num_requests": 1, @@ -713,11 +983,25 @@ def _run(self, tmp_path, monkeypatch, fail_peer_idx=None): }, ) plan = pcfg.resolve_plan(cfg) + if peer_progress_timeout_s is not None: + plan["peer_progress_timeout_s"] = peer_progress_timeout_s + monkeypatch.setattr(rp, "CONTROL_POLL_INTERVAL_MS", 50) 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) for i in range(2)] + ctxs = [ + self._mk_runner( + "ctx", + i, + plan, + work, + monkeypatch, + fail_ctx=(fail_ctx_idx is not None and i == fail_ctx_idx), + wave_delay_s=first_ctx_wave_delay_s if i == 0 else 0, + ) + for i in range(num_ctx_servers) + ] if fail_peer_idx is not None: real_gen_run_peer = rp.gen_run_peer @@ -752,16 +1036,21 @@ def rec(peer, exc): ] for t in threads: t.start() + + def gen_arm(what, publish_progress=True, **kwargs): + if publish_progress: + rp.publish_peer_progress(gen, what) + try: rp._drive_ctx_peers( - gen, noop, noop, rp._make_peer_failure_recorder(gen, noop, {"what": "test"}) + gen, gen_arm, noop, rp._make_peer_failure_recorder(gen, noop, {"what": "test"}) ) finally: # Always join every peer, even when the driver raises. Asserting # inside the loop can itself strand later peers and trip CI's # pytest-threadleak hook. for thread in threads: - thread.join(timeout=5) + thread.join(timeout=60) leaked = [thread.name for thread in threads if thread.is_alive()] assert not leaked, f"ctx serve threads wedged: {leaked}" return plan, gen, ctxs, failures @@ -781,6 +1070,123 @@ 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_four_ctx_full_pass(self, tmp_path, monkeypatch): + # ctx_0's four delayed waves take longer than the one-second queued + # wait budget. gen phase progress must refresh ctx_1..ctx_3 rather than + # letting their hello watchdog expire cumulatively behind ctx_0. + plan, gen, ctxs, failures = self._run( + tmp_path, + monkeypatch, + num_ctx_servers=4, + first_ctx_wave_delay_s=0.3, + peer_progress_timeout_s=1, + ) + assert not failures + assert [c["peer"] for c in gen.recorder.cases] == [ + "ctx_0", + "ctx_1", + "ctx_2", + "ctx_3", + ] + assert all(c["status"] == "PASS" for c in gen.recorder.cases) + assert all([case["status"] for case in ctx.recorder.cases] == ["PASS"] for ctx in ctxs) + + def test_one_ctx_four_gen_full_pass(self, tmp_path, monkeypatch): + """Queued gen instances refresh from the active ctx's progress.""" + import threading + + monkeypatch.setenv("SLURM_JOB_ID", "778") + monkeypatch.setenv("SLURMD_NODENAME", "127.0.0.1") + monkeypatch.setattr(rp, "CONTROL_POLL_INTERVAL_MS", 50) + cfg = _disagg_yaml( + hardware={ + "gpus_per_node": 4, + "num_ctx_servers": 1, + "num_gen_servers": 4, + }, + cache_transceiver_precheck={ + "request_lengths": [32], + "num_requests": 1, + "warmup_requests": 1, + "rendezvous_timeout_s": 30, + "wave_timeout_s": 30, + "wireup_timeout_s": 0, + }, + ) + plan = pcfg.resolve_plan(cfg) + # One ctx session takes four 0.3s waves, so later gen instances wait + # longer than this budget and need ctx progress to avoid false timeout. + plan["peer_progress_timeout_s"] = 1 + work = str(tmp_path) + noop = lambda *args, **kwargs: None # noqa: E731 - signal.alarm needs main thread + + ctx = self._mk_runner( + "ctx", + 0, + plan, + work, + monkeypatch, + wave_delay_s=0.3, + ) + gens = [self._mk_runner("gen", i, plan, work, monkeypatch) for i in range(4)] + peer_failures = [] + thread_errors = [] + + def record_ctx_peer_failure(peer, exc): + peer_failures.append((peer, type(exc).__name__)) + + def progress_arm(runner): + def arm(what, publish_progress=True, **kwargs): + if publish_progress: + rp.publish_peer_progress(runner, what) + + return arm + + def run_ctx(): + try: + rp._serve_gen_peers( + ctx, + plan, + progress_arm(ctx), + noop, + record_ctx_peer_failure, + ) + except Exception as exc: # noqa: BLE001 - surface thread failures in the test + thread_errors.append(("ctx", exc)) + + def run_gen(gen): + try: + rp._drive_ctx_peers( + gen, + progress_arm(gen), + noop, + rp._make_peer_failure_recorder(gen, noop, {"what": "test"}), + ) + except Exception as exc: # noqa: BLE001 - surface thread failures in the test + thread_errors.append((f"gen_{gen.server_idx}", exc)) + + threads = [threading.Thread(target=run_ctx, name="ctx_0", daemon=True)] + threads.extend( + threading.Thread( + target=run_gen, + args=(gen,), + name=f"gen_{gen.server_idx}", + daemon=True, + ) + for gen in gens + ) + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=60) + + leaked = [thread.name for thread in threads if thread.is_alive()] + assert not leaked, f"orchestration threads wedged: {leaked}" + assert not thread_errors + assert not peer_failures + assert [case["status"] for case in ctx.recorder.cases] == ["PASS"] * 4 + assert all([case["status"] for case in gen.recorder.cases] == ["PASS"] for gen in gens) + 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. diff --git a/tests/unittest/others/test_cache_transceiver_precheck_run.py b/tests/unittest/others/test_cache_transceiver_precheck_run.py index 1520c9179e9e..0b8911073592 100644 --- a/tests/unittest/others/test_cache_transceiver_precheck_run.py +++ b/tests/unittest/others/test_cache_transceiver_precheck_run.py @@ -351,6 +351,8 @@ def _plan(**overrides): "rendezvous_timeout_s": 600, "wireup_timeout_s": 300, "wave_timeout_s": 180, + "setup_timeout_s": 600, + "peer_progress_timeout_s": 900, } plan.update(overrides) return plan @@ -368,49 +370,160 @@ def test_schedule_covers_all_cells_in_lockstep_order(): 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) + # Peer count is deliberately absent: active peer progress refreshes this + # bounded no-progress interval for every serialized waiter. + assert rp.hello_timeout_s(plan) == 900 # 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 -# --------------------------------------------------------------------------- # -# Model preference resolution -# --------------------------------------------------------------------------- # -def test_resolve_model_prefs_allows_registered_class_without_preference_hook(monkeypatch): - model_cls = type("ModelWithoutPreferenceHook", (), {}) - cache_cfg = types.SimpleNamespace(transceiver_runtime="CPP") - calls = [] +@pytest.mark.parametrize("role", ["ctx", "gen"]) +def test_peer_progress_marker_is_atomic_and_run_stamped(tmp_path, monkeypatch, role): + monkeypatch.setenv("SLURM_JOB_ID", "111") + runner = types.SimpleNamespace( + role=role, + is_leader=True, + work_dir=str(tmp_path), + server_idx=2, + ) + rp.publish_peer_progress(runner, "peer_0 first wave") + first = rp.read_peer_progress(str(tmp_path), role, 2) + assert isinstance(first, int) - def resolve_v2(shim, resolved_model_cls, pretrained_config): - calls.append((shim, resolved_model_cls, pretrained_config)) - return False + rp.publish_peer_progress(runner, "peer_0 second wave") + assert rp.read_peer_progress(str(tmp_path), role, 2) != first - hf_view = object() - monkeypatch.setattr(rp, "_lookup_model_cls", lambda _model_dir: (model_cls, hf_view)) - monkeypatch.setattr( - rp, - "load_internal_apis", - lambda: types.SimpleNamespace( - TorchLlmArgs=lambda **kwargs: types.SimpleNamespace(**kwargs), - resolve_kv_cache_manager_v2_auto=resolve_v2, - ), + monkeypatch.setenv("SLURM_JOB_ID", "222") + assert rp.read_peer_progress(str(tmp_path), role, 2) is None + + +def test_ctx_control_wait_refreshes_only_on_target_gen_progress(monkeypatch): + class Again(Exception): + pass + + class FakeSocket: + def __init__(self): + self.calls = 0 + + def recv(self): + self.calls += 1 + if self.calls <= 2: + raise Again + return rp.pack_msg(("done", {}), KEY) + + runner = types.SimpleNamespace( + is_leader=True, + work_dir="/unused", + comm=types.SimpleNamespace(bcast=lambda value, root=0: value), + _zmq=lambda: (types.SimpleNamespace(Again=Again), None), ) + progress = iter((10, 11, 11)) + monotonic = iter((0.0, 1.0, 2.0, 3.0)) + progress_reads = [] + + def read_progress(work_dir, role, server_idx): + progress_reads.append((role, server_idx)) + return next(progress) + + monkeypatch.setattr(rp, "read_peer_progress", read_progress) + monkeypatch.setattr(rp.time, "monotonic", lambda: next(monotonic)) + arms = [] + + def arm(what, seconds, python_alarm, publish_progress=True): + arms.append((what, seconds, python_alarm, publish_progress)) + + msg = rp._recv_ctx_control( + runner, + FakeSocket(), + KEY, + peer_idx=3, + what="bye gen_3", + timeout_s=5, + arm=arm, + refresh_from_gen_progress=True, + ) + assert msg[0] == "done" + assert arms == [ + ("bye gen_3", 5, False, True), + ("bye gen_3", 5, False, False), + ] + assert progress_reads == [("gen", 3), ("gen", 3), ("gen", 3)] - use_v2 = rp.resolve_model_prefs( - "/models/example", - { - "use_kv_cache_manager_v2": "auto", - "parallel": {"tp": 1, "pp": 1, "cp": 1}, - }, - cache_cfg, + +def test_ctx_control_wait_times_out_without_progress(monkeypatch): + class Again(Exception): + pass + + class FakeSocket: + def recv(self): + raise Again + + runner = types.SimpleNamespace( + is_leader=True, + work_dir="/unused", + comm=types.SimpleNamespace(bcast=lambda value, root=0: value), + _zmq=lambda: (types.SimpleNamespace(Again=Again), None), ) + monotonic = iter((0.0, 6.0)) + monkeypatch.setattr(rp, "read_peer_progress", lambda work_dir, role, server_idx: None) + monkeypatch.setattr(rp.time, "monotonic", lambda: next(monotonic)) + with pytest.raises(rp._Timeout, match="made no progress for 5s"): + rp._recv_ctx_control( + runner, + FakeSocket(), + KEY, + peer_idx=3, + what="hello gen_3", + timeout_s=5, + arm=lambda *args, **kwargs: None, + refresh_from_gen_progress=True, + ) - assert use_v2 is False - assert len(calls) == 1 - assert calls[0][1:] == (model_cls, hf_view) + +def test_watchdog_tracks_each_phase_budget(monkeypatch): + class FakeHangDetector: + instance = None + + def __init__(self, timeout, on_detected): + self.timeout = timeout + self.on_detected = on_detected + self.checkpoints = [] + FakeHangDetector.instance = self + + def start(self): + pass + + def checkpoint(self): + self.checkpoints.append(self.timeout) + + def cancel_task(self): + pass + + def stop(self): + pass + + api = types.SimpleNamespace(HangDetector=FakeHangDetector) + monkeypatch.setattr(rp, "load_internal_apis", lambda: api) + runner = types.SimpleNamespace( + role="gen", + is_leader=False, + server_idx=0, + side={"num_peers": 1}, + recorder=types.SimpleNamespace(record=lambda *args: None, finalize=lambda: None), + ) + previous_alarm_handler = rp.signal.getsignal(rp.signal.SIGALRM) + arm, disarm, stop, _ = rp._install_watchdog(runner, _plan(), rank=0) + try: + arm("setup", seconds=600) + arm("first wave", seconds=480) + arm("steady wave") + assert FakeHangDetector.instance.checkpoints == [660, 540, 240] + finally: + disarm() + stop() + rp.signal.signal(rp.signal.SIGALRM, previous_alarm_handler) # --------------------------------------------------------------------------- # @@ -433,27 +546,35 @@ def _ctx_finish_runner(monkeypatch, check_status): return runner, events -def test_ctx_finish_wave_frees_only_after_block_all_returns_every_request(monkeypatch): - events = [] +def test_ctx_finish_wave_frees_only_after_every_request_completes(monkeypatch): + runner, events = _ctx_finish_runner(monkeypatch, lambda _n: ([101, 102], [])) + reqs = { + 0: types.SimpleNamespace(py_request_id=101, state="in_progress"), + 1: types.SimpleNamespace(py_request_id=102, state="in_progress"), + } - def check_status(at_least_request_num): - events.append(("block_all", at_least_request_num)) - return [101, 102], [] + runner.ctx_finish_wave(reqs) - runner, free_events = _ctx_finish_runner(monkeypatch, check_status) + assert events == [("free", [0, 1])] + + +def test_ctx_finish_wave_retains_pages_after_transfer_failure(monkeypatch): + runner, events = _ctx_finish_runner(monkeypatch, lambda _n: ([101], [102])) reqs = { 0: types.SimpleNamespace(py_request_id=101, state="in_progress"), 1: types.SimpleNamespace(py_request_id=102, state="in_progress"), } - runner._free_all = lambda owned: events.append(("free", sorted(owned))) - runner.ctx_finish_wave(reqs) + with pytest.raises(rp._FatalTransferError, match=r"ctx transfer failed for pairs \[1\]"): + runner.ctx_finish_wave(reqs) - assert events == [("block_all", None), ("free", [0, 1])] - assert free_events == [] + assert events == [] -def test_ctx_finish_wave_retains_pages_when_block_all_omits_request(monkeypatch): +def test_ctx_finish_wave_block_all_requires_terminal_status(monkeypatch): + # A request still nonterminal after block-all returns exceeded the + # kv_transfer_timeout deadline (Python) or violated the true block-all + # contract (C++); both classify as a page-retaining gate failure. runner, events = _ctx_finish_runner(monkeypatch, lambda _n: ([101], [])) reqs = { 0: types.SimpleNamespace(py_request_id=101, state="in_progress"), @@ -466,27 +587,175 @@ def test_ctx_finish_wave_retains_pages_when_block_all_omits_request(monkeypatch) assert events == [] -def test_ctx_finish_wave_retains_pages_when_block_all_raises(monkeypatch): - def check_status(_n): - raise RuntimeError("interrupted") +def _gen_run_wave_runner(monkeypatch, outcome): + requests = {} + events = [] + states = types.SimpleNamespace( + DISAGG_GENERATION_TRANS_COMPLETE="complete", + DISAGG_TRANS_ERROR="error", + ) + monkeypatch.setitem( + sys.modules, + "torch", + types.SimpleNamespace( + cuda=types.SimpleNamespace(synchronize=lambda: events.append("cuda_sync")) + ), + ) + monkeypatch.setitem( + sys.modules, + "tensorrt_llm", + types.SimpleNamespace(logger=types.SimpleNamespace(info=lambda *_args, **_kwargs: None)), + ) - runner, events = _ctx_finish_runner(monkeypatch, check_status) - reqs = {0: types.SimpleNamespace(py_request_id=101, state="in_progress")} + def make_request(_is_ctx, rid, _req_len, _runtime, ctx_params=None): + req = types.SimpleNamespace(py_request_id=rid, state="in_progress") + requests[rid] = req + return req - with pytest.raises(rp._FatalTransferError, match="interrupted"): - runner.ctx_finish_wave(reqs) + monkeypatch.setattr(rp, "make_request", make_request) + monkeypatch.setattr(rp, "add_sequence", lambda *_args: None) + + def check_status(_at_least_request_num): + completed, failed, cancelled = outcome(requests) + for rid in completed: + requests[rid].state = states.DISAGG_GENERATION_TRANS_COMPLETE + for rid in failed: + requests[rid].state = states.DISAGG_TRANS_ERROR + return completed, failed, [requests[rid] for rid in cancelled] + + runner = object.__new__(rp.PrecheckRunner) + runner.runtime = "PYTHON" + runner.kvm = object() + runner.use_v2 = True + runner.server_idx = 0 + runner.rank = 0 + runner.llm_request_state = states + runner.plan = {"verify_data": False, "warmup_requests": 0} + runner.xcvr = types.SimpleNamespace( + request_and_receive_async=lambda _req: None, + check_gen_transfer_status=check_status, + ) + runner.comm = types.SimpleNamespace(allgather=lambda value: [value]) + runner._owned = lambda _wave: [0, 1] + runner._pair_rid = lambda _peer, _li, _rep, pair: 101 + pair + runner._consensus_error = lambda err: None if err is None else repr(err) + runner._free_all = lambda reqs: events.append(("free", sorted(reqs))) + return runner, events + + +def test_gen_run_wave_frees_only_after_every_receive_completes(monkeypatch): + runner, events = _gen_run_wave_runner(monkeypatch, lambda _reqs: ([101, 102], [], [])) + + ok, detail = runner.gen_run_wave(0, 0, 64, 0, [0, 1], {0: object(), 1: object()}) + + assert ok and not detail + assert events == ["cuda_sync", ("free", [0, 1])] + + +@pytest.mark.parametrize( + ("outcome", "message"), + ( + (lambda _reqs: ([101], [102], []), r"failed=\[102\]"), + (lambda _reqs: ([101], [], [102]), r"cancelled=\[102\]"), + # Nonterminal after the block-all deadline: a gate failure, not a + # keep-polling condition. + (lambda _reqs: ([101], [], []), r"missing=\[102\]"), + ), +) +def test_gen_run_wave_retains_pages_without_all_successes(monkeypatch, outcome, message): + runner, events = _gen_run_wave_runner(monkeypatch, outcome) + + with pytest.raises(rp._FatalTransferError, match=message): + runner.gen_run_wave(0, 0, 64, 0, [0, 1], {0: object(), 1: object()}) assert events == [] -def test_ctx_finish_wave_retains_pages_when_request_failed(monkeypatch): - runner, events = _ctx_finish_runner(monkeypatch, lambda _n: ([101], [102])) - reqs = { - 0: types.SimpleNamespace(py_request_id=101, state="in_progress"), - 1: types.SimpleNamespace(py_request_id=102, state="error"), +@pytest.mark.parametrize( + ("failing_resolver", "message"), + ( + ("runtime", "refusing to validate a runtime"), + ("manager", "refusing to assume V1"), + ), +) +def test_model_preference_resolution_fails_closed(monkeypatch, failing_resolver, message): + def resolve_runtime(_shim, _model_cls, _hf_view): + if failing_resolver == "runtime": + raise RuntimeError("runtime resolution failed") + + def resolve_manager(_args, _model_cls, _hf_view): + if failing_resolver == "manager": + raise RuntimeError("manager resolution failed") + return True + + api = types.SimpleNamespace( + resolve_transceiver_runtime_auto=resolve_runtime, + resolve_kv_cache_manager_v2_auto=resolve_manager, + TorchLlmArgs=lambda **kwargs: types.SimpleNamespace(**kwargs), + MTPDecodingConfig=lambda **kwargs: types.SimpleNamespace(**kwargs), + ) + monkeypatch.setattr(rp, "load_internal_apis", lambda: api) + monkeypatch.setattr(rp, "_lookup_model_cls", lambda _model_dir: (object(), object())) + cache_cfg = types.SimpleNamespace(transceiver_runtime="auto") + side = { + "use_kv_cache_manager_v2": "auto", + "parallel": {"tp": 1, "pp": 1, "cp": 1}, } - with pytest.raises(rp._FatalTransferError, match=r"ctx transfer failed for pairs \[1\]"): + with pytest.raises(RuntimeError, match=message): + rp.resolve_model_prefs("/model", side, cache_cfg) + + +# --------------------------------------------------------------------------- # +# Model preference resolution +# --------------------------------------------------------------------------- # +def test_resolve_model_prefs_allows_registered_class_without_preference_hook(monkeypatch): + model_cls = type("ModelWithoutPreferenceHook", (), {}) + cache_cfg = types.SimpleNamespace(transceiver_runtime="CPP") + calls = [] + + def resolve_v2(shim, resolved_model_cls, pretrained_config): + calls.append((shim, resolved_model_cls, pretrained_config)) + return False + + hf_view = object() + monkeypatch.setattr(rp, "_lookup_model_cls", lambda _model_dir: (model_cls, hf_view)) + monkeypatch.setattr( + rp, + "load_internal_apis", + lambda: types.SimpleNamespace( + TorchLlmArgs=lambda **kwargs: types.SimpleNamespace(**kwargs), + resolve_kv_cache_manager_v2_auto=resolve_v2, + ), + ) + + use_v2 = rp.resolve_model_prefs( + "/models/example", + { + "use_kv_cache_manager_v2": "auto", + "parallel": {"tp": 1, "pp": 1, "cp": 1}, + }, + cache_cfg, + ) + + assert use_v2 is False + assert len(calls) == 1 + assert calls[0][1:] == (model_cls, hf_view) + + +# --------------------------------------------------------------------------- # +# Transfer ownership +# --------------------------------------------------------------------------- # + + +def test_ctx_finish_wave_retains_pages_when_block_all_raises(monkeypatch): + def check_status(_n): + raise RuntimeError("interrupted") + + runner, events = _ctx_finish_runner(monkeypatch, check_status) + reqs = {0: types.SimpleNamespace(py_request_id=101, state="in_progress")} + + with pytest.raises(rp._FatalTransferError, match="interrupted"): runner.ctx_finish_wave(reqs) assert events == [] @@ -623,73 +892,6 @@ def test_ctx_run_wave_post_dispatch_collective_error_is_fatal(monkeypatch): runner.ctx_run_wave(0, 0, 64, 0, [0]) -def _gen_run_wave_runner(monkeypatch, outcome): - requests = {} - events = [] - states = types.SimpleNamespace( - DISAGG_GENERATION_TRANS_COMPLETE="complete", - DISAGG_TRANS_ERROR="error", - ) - monkeypatch.setitem( - sys.modules, - "torch", - types.SimpleNamespace( - cuda=types.SimpleNamespace(synchronize=lambda: events.append("cuda_sync")) - ), - ) - monkeypatch.setitem( - sys.modules, - "tensorrt_llm", - types.SimpleNamespace(logger=types.SimpleNamespace(info=lambda *_args, **_kwargs: None)), - ) - - def make_request(_is_ctx, rid, _req_len, _runtime, ctx_params=None): - req = types.SimpleNamespace(py_request_id=rid, state="in_progress") - requests[rid] = req - return req - - monkeypatch.setattr(rp, "make_request", make_request) - monkeypatch.setattr(rp, "add_sequence", lambda *_args: None) - - def check_status(_at_least_request_num): - completed, failed, cancelled = outcome(requests) - for rid in completed: - requests[rid].state = states.DISAGG_GENERATION_TRANS_COMPLETE - for rid in failed: - requests[rid].state = states.DISAGG_TRANS_ERROR - return completed, failed, [requests[rid] for rid in cancelled] - - runner = object.__new__(rp.PrecheckRunner) - runner.runtime = "PYTHON" - runner.kvm = object() - runner.use_v2 = True - runner.server_idx = 0 - runner.rank = 0 - runner.side = {"parallel": {"enable_attention_dp": False}} - runner.mapping = types.SimpleNamespace(pp_rank=0) - runner.llm_request_state = states - runner.plan = {"verify_data": False, "warmup_requests": 0} - runner.xcvr = types.SimpleNamespace( - request_and_receive_async=lambda _req: None, - check_gen_transfer_status=check_status, - ) - runner.comm = types.SimpleNamespace(allgather=lambda value: [value]) - runner._owned = lambda _wave: [0, 1] - runner._pair_rid = lambda _peer, _li, _rep, pair: 101 + pair - runner._consensus_error = lambda err: None if err is None else repr(err) - runner._free_all = lambda owned: events.append(("free", sorted(owned))) - return runner, events - - -def test_gen_run_wave_frees_only_after_every_python_receive_completes(monkeypatch): - runner, events = _gen_run_wave_runner(monkeypatch, lambda _reqs: ([101, 102], [], [])) - - ok, detail = runner.gen_run_wave(0, 0, 64, 0, [0, 1], {0: object(), 1: object()}) - - assert ok and not detail - assert events == ["cuda_sync", ("free", [0, 1])] - - def test_gen_run_wave_checks_python_status_on_empty_owner_rank(monkeypatch): calls = [] @@ -726,23 +928,6 @@ def receive(_req): assert events == [] -@pytest.mark.parametrize( - ("outcome", "message"), - ( - (lambda _reqs: ([101], [102], []), r"failed=\[102\]"), - (lambda _reqs: ([101], [], [102]), r"cancelled=\[102\]"), - (lambda _reqs: ([101], [], []), r"missing=\[102\]"), - ), -) -def test_gen_run_wave_retains_pages_without_all_successes(monkeypatch, outcome, message): - runner, events = _gen_run_wave_runner(monkeypatch, outcome) - - with pytest.raises(rp._FatalTransferError, match=message): - runner.gen_run_wave(0, 0, 64, 0, [0, 1], {0: object(), 1: object()}) - - assert events == [] - - def test_gen_run_wave_does_not_free_when_peer_rank_is_unsafe(monkeypatch): runner, events = _gen_run_wave_runner(monkeypatch, lambda _reqs: ([101, 102], [], [])) consensus_calls = [] @@ -1042,6 +1227,20 @@ def get_preferred_kv_cache_manager_version(cls, pretrained_config=None): cache_cfg, ) + def test_deepseek_v4_auto_selects_kv_cache_manager_v2(self, api, tmp_path): + model_dir = tmp_path / "deepseek-v4" + model_dir.mkdir() + (model_dir / "config.json").write_text( + json.dumps({"architectures": ["DeepseekV4ForCausalLM"]}) + ) + side = { + "use_kv_cache_manager_v2": "auto", + "parallel": {"tp": 1, "pp": 1, "cp": 1}, + } + cache_cfg = api.CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON") + + assert rp.resolve_model_prefs(str(model_dir), side, cache_cfg) is True + def test_enum_members(self, api): for enum, members in ( (api.DataType, ("FP8", "HALF", "BF16")), diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index 241a4aee0d09..8737091bbd2d 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -21,8 +21,10 @@ import tensorrt_llm.bindings import tensorrt_llm.bindings.executor as trtllm 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.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import ( + create_kv_cache_transceiver, + maybe_enable_fabric_memory_for_python_transceiver) from tensorrt_llm._torch.pyexecutor.llm_request import (LlmRequest, LlmRequestState) from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import \ @@ -42,6 +44,34 @@ KV_TRANSFER_COMPLETION_MARGIN_S = 10.0 +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "runtime,manager_cls,initial_value,expected_value", + [ + ("PYTHON", KVCacheManager, None, "1"), + ("PYTHON", KVCacheManager, "0", "0"), + ("PYTHON", KVCacheManagerV2, None, None), + ("CPP", KVCacheManager, None, None), + # "auto" is resolved to PYTHON/CPP by serving before this helper runs; + # callers that construct CacheTransceiverConfig directly must resolve + # it first — the helper deliberately leaves "auto" untouched. + ("auto", KVCacheManager, None, None), + ], +) +def test_maybe_enable_fabric_memory_for_python_transceiver( + monkeypatch, runtime, manager_cls, initial_value, expected_value): + env_name = "TRTLLM_KVCACHE_POOL_USE_FABRIC_MEMORY" + if initial_value is None: + monkeypatch.delenv(env_name, raising=False) + else: + monkeypatch.setenv(env_name, initial_value) + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime=runtime) + + maybe_enable_fabric_memory_for_python_transceiver(config, manager_cls) + + assert os.environ.get(env_name) == expected_value + + @pytest.mark.parametrize("transceiver_runtime", ["CPP", "auto"]) def test_cpp_transceiver_rejects_mixed_mamba_manager(transceiver_runtime): config = CacheTransceiverConfig(backend="NIXL", diff --git a/tests/unittest/scripts/test_cluster_env.py b/tests/unittest/scripts/test_cluster_env.py index 26905068b10f..7623794040cb 100644 --- a/tests/unittest/scripts/test_cluster_env.py +++ b/tests/unittest/scripts/test_cluster_env.py @@ -81,14 +81,20 @@ def test_gpu_type_from_supported_gpus( "export UCX_NET_DEVICES=mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1," "mlx5_4:1,mlx5_5:1,mlx5_10:1,mlx5_11:1", ), + ( + "oci-hsg-cs-001", + "export UCX_NET_DEVICES=mlx5_0:1,mlx5_1:1,mlx5_3:1,mlx5_4:1,eth0", + ), ( "oci-aga-cs-001", - "export UCX_TLS=^tcp,rc_gda,gga UCX_IB_MLX5_DEVX=n " - "UCX_NET_DEVICES=" - "rdma_vf_rail0:1,rdma_vf_rail1:1,rdma_vf_rail2:1,rdma_vf_rail3:1 " - "UCX_IB_TRAFFIC_CLASS=96 TRTLLM_NIXL_NUM_THREADS=1", + "export UCX_TLS=cuda_ipc,cuda_copy,sm,self,tcp UCX_TCP_AF_PRIO=inet", + ), + ( + "aws-cmh", + "export UCX_TLS=cuda_ipc,cuda_copy,sm,self,tcp " + "UCX_NET_DEVICES=eth0,mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1," + "mlx5_4:1,mlx5_5:1,mlx5_6:1,mlx5_7:1", ), - ("aws-cmh", "export UCX_TLS=cuda_ipc,cuda_copy,sm,self,tcp"), ("aws-dfw-prod", "export UCX_TLS=^gdr_copy"), ), ) diff --git a/tests/unittest/scripts/test_perf_sanity_helpers.py b/tests/unittest/scripts/test_perf_sanity_helpers.py index f344529c2f35..941e787e357b 100644 --- a/tests/unittest/scripts/test_perf_sanity_helpers.py +++ b/tests/unittest/scripts/test_perf_sanity_helpers.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import subprocess import sys from pathlib import Path @@ -26,6 +27,64 @@ from defs.perf import test_perf_sanity as perf_sanity # noqa: E402 +def test_run_benchmark_with_log_returns_successful_output(tmp_path: Path) -> None: + benchmark_log = tmp_path / "trtllm-benchmark.0.0.log" + command = [sys.executable, "-c", "print('benchmark succeeded')"] + + output = perf_sanity._run_benchmark_with_log(command, {}, str(benchmark_log)) + + assert output == "benchmark succeeded\n" + assert benchmark_log.read_text(encoding="utf-8") == output + + +def test_run_benchmark_with_log_preserves_failed_output(tmp_path: Path) -> None: + benchmark_log = tmp_path / "trtllm-benchmark.0.0.log" + command = [ + sys.executable, + "-c", + ( + "import sys; " + "print('benchmark stdout'); " + "print('benchmark stderr', file=sys.stderr); " + "sys.exit(7)" + ), + ] + + with pytest.raises(subprocess.CalledProcessError) as exc_info: + perf_sanity._run_benchmark_with_log(command, {}, str(benchmark_log)) + + expected_output = "benchmark stdout\nbenchmark stderr\n" + assert exc_info.value.returncode == 7 + assert exc_info.value.output.decode() == expected_output + assert benchmark_log.read_text(encoding="utf-8") == expected_output + + +def test_benchmark_log_is_included_in_report_logs(tmp_path: Path) -> None: + benchmark_log = tmp_path / "trtllm-benchmark.0.0.log" + benchmark_log.touch() + aggr_commands = perf_sanity.AggrTestCmds( + server_cmds=[[]], + client_cmds={0: [[]]}, + timeout=1, + output_dir=str(tmp_path), + test_output_dir=str(tmp_path), + ) + disagg_commands = perf_sanity.DisaggTestCmds( + server_cmds=[], + client_cmds={}, + timeout=1, + hostname="localhost", + disagg_serving_type="BENCHMARK", + num_ctx_servers=0, + num_gen_servers=0, + output_dir=str(tmp_path), + test_output_dir=str(tmp_path), + ) + + assert str(benchmark_log) in aggr_commands.get_server_logs(0) + assert str(benchmark_log) in disagg_commands.get_server_logs(0) + + def test_sentinel_timeout_falls_back_to_current_gen_logs( monkeypatch: pytest.MonkeyPatch, tmp_path: Path, diff --git a/tests/unittest/scripts/test_perf_submit.py b/tests/unittest/scripts/test_perf_submit.py index de81570e14fd..5473f4768494 100644 --- a/tests/unittest/scripts/test_perf_submit.py +++ b/tests/unittest/scripts/test_perf_submit.py @@ -180,6 +180,49 @@ def test_ci_submit_selects_same_least_duration_shard_as_pytest_split( assert selected == test_lines[1] +@pytest.mark.parametrize( + ("selected_suffix", "waive_line", "expected"), + ( + (" SKIP (inline)", "", True), + (" SKIP(inline)", "", True), + ("", "{nodeid} SKIP (global)", True), + ("", "{nodeid} SKIP(global)", True), + ("", "{nodeid} XFAIL (known failure)", False), + ("", "perf/test_other.py::test_other SKIP (other)", False), + ), +) +def test_ci_submit_skips_precheck_only_for_selected_skip_waive( + ci_submit_module: ModuleType, + tmp_path: Path, + selected_suffix: str, + waive_line: str, + expected: bool, +) -> None: + nodeid = "perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-case]" + waives = tmp_path / "waives.txt" + waives.write_text(waive_line.format(nodeid=nodeid), encoding="utf-8") + + assert ( + ci_submit_module.selected_test_is_skip_waived(f"{nodeid}{selected_suffix}", waives) + is expected + ) + + +def test_ci_submit_honors_matching_platform_scoped_skip_waive( + ci_submit_module: ModuleType, tmp_path: Path +) -> None: + nodeid = "perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-case]" + waives = tmp_path / "waives.txt" + waives.write_text(f"full:GB300/{nodeid} SKIP (platform)\n", encoding="utf-8") + + assert ci_submit_module.selected_test_is_skip_waived( + nodeid, waives, test_prefix="GB300-Disagg-Perf" + ) + assert not ci_submit_module.selected_test_is_skip_waived( + nodeid, waives, test_prefix="DGX_B200-Disagg-Perf" + ) + + def test_ci_submit_selector_matches_installed_pytest_split( ci_submit_module: ModuleType, ) -> None: