diff --git a/tests/integration/defs/perf/README_test_perf_sanity.md b/tests/integration/defs/perf/README_test_perf_sanity.md index 57af410d67b3..9829e2fb927c 100644 --- a/tests/integration/defs/perf/README_test_perf_sanity.md +++ b/tests/integration/defs/perf/README_test_perf_sanity.md @@ -293,6 +293,32 @@ perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-{disagg config file base na perf/test_perf_sanity.py::test_e2e[disagg_upload-e2e-deepseek-r1-fp4_1k1k_ctx1_gen1_dep8] ``` +#### Connection warmup + +Before the measured clients, the harness (`DisaggTestCmds._run_connection_warmup`) +sends `2 x (ctx servers x ctx attention-DP size) x (gen servers x gen attention-DP size)` +short requests (ISL 100, OSL 10), where a server's attention-DP size is its TP size with +attention DP and 1 without. They open the ctx<->gen KV-cache transfer connections +outside the measured window. `examples/disaggregated/slurm/benchmark/submit.py` sizes +the same warmup as `ucx_warmup_requests`. Without it, each ctx rank loads a gen rank's +transfer agent the first time that gen rank requests KV from its ctx server, one agent +at a time on a single listener thread, and the first measured requests wait behind that +queue (nvbugs/6564113). Notes: +- e2e only, and skipped when every client is AgentX (which runs its own warmup) or + accuracy-only. +- If only some warmup requests fail, the harness logs a warning and the measured + clients still run. The test fails instead, rather than letting the lane run cold or + hang, if the warmup's benchmark_serving exits non-zero, runs past + `CONNECTION_WARMUP_TIMEOUT` (capped by the test timeout), or every warmup request + fails. +- Client log: `trtllm-connection-warmup..log`, not parsed for metrics. + The servers' perf_metrics records of the warmup requests are excluded from the time + breakdown (see below). +- Iteration-based nsys ranges (`local/submit.py --ctx-nsys-start-stop` / + `--gen-nsys-start-stop`) count from server start, so they begin with the + connection-warmup iterations; start the range later to profile only measured + iterations. + ### The optional instrumentation modifier The three shapes that read a disaggregated config (2, 3 and 4 above) take one @@ -364,6 +390,17 @@ mid-record) costs only that line, never the whole file: dropping the disagg serv file would silently reroute the per-request groups to same-role worker fallbacks, which produce plausible values for the wrong phase. +The servers also append records of the e2e connection warmup's requests to these +files. The harness therefore waits until the largest file holds a record for every +successful warmup request and the files hold still, then notes each file's line count; the aggregation skips those leading lines and leaves +them out of the census. Unlike the measured run's census, this one is a failure: if +the count is not reached within `PERF_METRICS_SETTLE_TIMEOUT`, the test fails rather +than mixing late warmup records into the breakdown. The log line +`Time breakdown: excluded connection-warmup records from {...}` prints the counts as +JSON. To reproduce the uploaded numbers offline, run +`time_breakdown_metrics.py --mode e2e --output-dir /perf_metrics +--drop-warmup-request --leading-lines-to-skip ''`. + ## CI Test Database Test lists are defined in `tests/integration/test_lists/test-db/`. diff --git a/tests/integration/defs/perf/test_perf_sanity.py b/tests/integration/defs/perf/test_perf_sanity.py index fb14af964298..838b8deb0be8 100644 --- a/tests/integration/defs/perf/test_perf_sanity.py +++ b/tests/integration/defs/perf/test_perf_sanity.py @@ -68,6 +68,9 @@ # (agentx_client.py). Any other non-empty value is rejected at parse time. AGENTX_BENCHMARK_CLIENT = "agentx" +# Modes whose measured client sends benchmark_serving's single un-measured test +# request (it omits --no-test-input). Unrelated to the connection warmup +# (DisaggTestCmds._run_connection_warmup). WARMUP_BENCHMARK_MODES = ("e2e", "ctx_only") @@ -174,6 +177,19 @@ def ensure_bench_serving_repo() -> str: # rather than call-site literals so a test can shorten the window. PERF_METRICS_SETTLE_SECONDS = _TB_SETTLE_SECONDS PERF_METRICS_SETTLE_TIMEOUT = _TB_SETTLE_TIMEOUT +# Connection warmup for disagg e2e lanes; see DisaggTestCmds._run_connection_warmup. +# Shape and count follow ucx_warmup_requests in +# examples/disaggregated/slurm/benchmark/{submit.py,run_benchmark.sh}. +CONNECTION_WARMUP_ISL = 100 +CONNECTION_WARMUP_OSL = 10 +CONNECTION_WARMUP_REQUESTS_PER_RANK_PAIR = 2 +# A healthy warmup takes about a minute even on the largest lanes (768 requests to a +# cold 32-rank gen server finished in 42 s on oci-jhb), so reaching this bound means +# the servers are wedged. +CONNECTION_WARMUP_TIMEOUT = 600 +# benchmark_serving's "Failed requests:" summary line, its only failure signal: +# benchmark_serving exits 0 even when requests fail. +_FAILED_REQUESTS_RE = re.compile(r"^Failed requests:\s+(\d+)", re.MULTILINE) def server_ready_timeout(default: int, mode: str) -> int: @@ -1626,8 +1642,24 @@ def publish_addr_file(path: str, host: str, port: int) -> None: os.unlink(tmp_path) -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.""" +def _run_benchmark_with_log( + cmd: List[str], env: Dict[str, str], log_path: str, timeout: float = DEFAULT_TIMEOUT +) -> str: + """Run a benchmark while streaming its combined output to an artifact log. + + Args: + cmd: Benchmark command line. + env: Subprocess environment; PYTHONUNBUFFERED=1 is added if unset. + log_path: File that receives the combined stdout and stderr. + timeout: Seconds before the benchmark is killed. + + Returns: + The combined output, decoded as UTF-8. + + Raises: + subprocess.CalledProcessError: The benchmark exited non-zero. + subprocess.TimeoutExpired: The benchmark ran longer than ``timeout``. + """ benchmark_env = env.copy() benchmark_env.setdefault("PYTHONUNBUFFERED", "1") with open(log_path, "wb") as log_file: @@ -1637,7 +1669,7 @@ def _run_benchmark_with_log(cmd: List[str], env: Dict[str, str], log_path: str) stdout=log_file, stderr=subprocess.STDOUT, check=False, - timeout=DEFAULT_TIMEOUT, + timeout=timeout, ) with open(log_path, "rb") as log_file: @@ -2344,6 +2376,7 @@ def append_time_breakdown_metrics( pending_time_breakdown: List[dict], outputs: List[str], breakdown_dir: str, + leading_lines_to_skip: Optional[Dict[str, int]] = None, ) -> None: """Aggregate the workers' perf_metrics JSONLs into log lines the parser reads. @@ -2365,6 +2398,10 @@ def append_time_breakdown_metrics( parsed ``Time Breakdown ...`` lines is what check_test_failure hard-fails on, which keeps the diagnosis in one place instead of tearing down the whole session from inside a post-benchmark hook. + + ``leading_lines_to_skip`` maps a JSONL basename to the number of lines it held + before the measured clients started (the connection warmup's records); those lines + are left out of both the census and the breakdown. """ if not pending_time_breakdown: return @@ -2379,6 +2416,7 @@ def append_time_breakdown_metrics( expected_requests=expected_requests or None, stable_seconds=PERF_METRICS_SETTLE_SECONDS, timeout_seconds=PERF_METRICS_SETTLE_TIMEOUT, + leading_lines_to_skip=leading_lines_to_skip, ) for warning in wait_info["warnings"]: print_info(f"Time breakdown: {warning}") @@ -2407,7 +2445,10 @@ def append_time_breakdown_metrics( # The client's warmup request is un-measured and absent from every other # metric on the row, so it is excluded here too -- see _drop_warmup_record. metrics, info = compute_time_breakdown_metrics( - paths, case_type, drop_warmup_request=bool(record.get("warmup")) + paths, + case_type, + drop_warmup_request=bool(record.get("warmup")), + leading_lines_to_skip=leading_lines_to_skip, ) except (OSError, ValueError, KeyError) as exc: print_info(f"Time breakdown aggregation failed for {breakdown_dir}: {exc}") @@ -2418,6 +2459,11 @@ def append_time_breakdown_metrics( print_info(f"Time breakdown ({case_type}) from {len(paths)} file(s): {info['counts']}") if info["warmup_dropped"]: print_info(f"Time breakdown: excluded the warmup request from {info['warmup_dropped']}") + if info["leading_lines_skipped"]: + print_info( + "Time breakdown: excluded connection-warmup records from " + f"{json.dumps(info['leading_lines_skipped'])}" + ) summary_lines = "\n".join(format_metric_log_lines(metrics)) with open(record["benchmark_file_path"], "a") as benchmark_ctx: @@ -2589,6 +2635,11 @@ def get_cmd_str(self, server_idx: int) -> List[str]: return ["aggr_server tests, please check config files"] +def _attention_dp_size(config: "ServerConfig") -> int: + """Attention-DP ranks per worker: its TP size with attention DP, else 1.""" + return config.tp if config.enable_attention_dp else 1 + + class DisaggTestCmds(NamedTuple): """Commands for multi-node disaggregated server perf sanity tests.""" @@ -2955,6 +3006,7 @@ def _append_time_breakdown_metrics( self, pending_time_breakdown: List[dict], outputs: List[str], + leading_lines_to_skip: Optional[Dict[str, int]] = None, ) -> None: """Disagg entry point for the shared aggregation; see the module function. @@ -2968,7 +3020,9 @@ def _append_time_breakdown_metrics( not a method here. Calling the method on self would raise AttributeError after the whole benchmark has already run. """ - append_time_breakdown_metrics(pending_time_breakdown, outputs, self.perf_metrics_output_dir) + append_time_breakdown_metrics( + pending_time_breakdown, outputs, self.perf_metrics_output_dir, leading_lines_to_skip + ) def get_server_logs(self, server_idx: int) -> List[str]: server_logs = [] @@ -2993,6 +3047,11 @@ def get_server_logs(self, server_idx: int) -> List[str]: ) ) ) + # Listed only when present, like the benchmark logs above. If the connection warmup + # fails, report_error then shows its log: the only client-side log at that point. + connection_warmup_log = self._connection_warmup_log_path(server_idx) + if os.path.exists(connection_warmup_log): + server_logs.append(connection_warmup_log) return server_logs @staticmethod @@ -3008,6 +3067,213 @@ def _wait_for_config_file(config_path: str, timeout: int = 600) -> None: print_info(f"Waiting for config file {config_path}, elapsed: {elapsed:.0f}s") time.sleep(1) + def _connection_warmup_log_path(self, server_idx: int) -> str: + return os.path.join(self.test_output_dir, f"trtllm-connection-warmup.{server_idx}.log") + + def _remove_stale_connection_warmup_log(self, server_idx: int) -> None: + """Delete a previous attempt's connection-warmup log. + + A retry reuses test_output_dir, and get_server_logs feeds this log to the readiness + wait's error scan, so an error line left by the previous attempt would fail the new + attempt before it starts. + """ + log_path = self._connection_warmup_log_path(server_idx) + if os.path.exists(log_path): + os.remove(log_path) + + def _measured_client_for_connection_warmup(self, server_idx: int) -> Optional["ClientConfig"]: + """The measured client whose env, backend and trust_remote_code the warmup reuses. + + Returns None to skip the warmup. Only e2e lanes warm up, as in + examples/disaggregated/slurm/benchmark/submit.py; gen_only lanes start measuring + only after the gen queue fills, and warmup requests would count toward that fill. + AgentX clients run their own warmup phase, and accuracy-only clients measure no + performance, so a lane with only those skips it. + """ + configs_for_idx = ( + self.server_configs[server_idx] if server_idx < len(self.server_configs) else None + ) + if configs_for_idx is None or configs_for_idx[2].benchmark_mode != "e2e": + return None + return next( + ( + client + for client in self.client_configs.get(server_idx, []) + if not client.only_run_accuracy + and client.benchmark_client != AGENTX_BENCHMARK_CLIENT + ), + None, + ) + + def _connection_warmup_num_requests(self, server_idx: int) -> int: + """Two connection-warmup requests per (ctx attention-DP rank, gen attention-DP rank) pair. + + This is the count submit.py sends. The router does not guarantee every pair gets + one. What matters is that every (gen attention-DP rank, ctx server) pair gets at + least one: the first such request makes each rank of that ctx server load the gen + rank's transfer agent. This count averages 2 x (ctx attention-DP size) requests per + such pair. + """ + ctx_config, gen_config, _ = self.server_configs[server_idx] + num_ctx_dp_ranks = self.num_ctx_servers * _attention_dp_size(ctx_config) + num_gen_dp_ranks = self.num_gen_servers * _attention_dp_size(gen_config) + return CONNECTION_WARMUP_REQUESTS_PER_RANK_PAIR * num_ctx_dp_ranks * num_gen_dp_ranks + + def _connection_warmup_cmd( + self, measured_client: "ClientConfig", num_requests: int + ) -> List[str]: + """benchmark_serving command that sends all warmup requests at once. + + Uses the lane's model and the measured client's backend and trust_remote_code, + and tokenizes random ids on the client as the measured client does on the random + dataset, so it hits the same endpoint the same way; --no-test-input keeps the + request count exact. The caller adds host and port. + """ + model_dir = get_model_dir(self.model_name) + model_path = model_dir if os.path.exists(model_dir) else self.model_name + cmd = [ + "python", + "-m", + "tensorrt_llm.serve.scripts.benchmark_serving", + "--model", + model_path, + "--tokenizer", + model_path, + "--dataset-name", + "random", + "--random-ids", + "--tokenize-on-client", + "--random-input-len", + str(CONNECTION_WARMUP_ISL), + "--random-output-len", + str(CONNECTION_WARMUP_OSL), + "--num-prompts", + str(num_requests), + "--ignore-eos", + "--non-streaming", + "--no-test-input", + ] + if measured_client.backend: + cmd += ["--backend", measured_client.backend] + if measured_client.trust_remote_code: + cmd.append("--trust-remote-code") + return cmd + + def _run_connection_warmup( + self, server_idx: int, hostname: str, port: int, measured_client: "ClientConfig" + ) -> int: + """Open the KV-cache transfer connections before the measured clients run. + + Each ctx rank loads a gen rank's transfer agent the first time that gen rank + requests KV from its ctx server, and it loads agents one at a time on a single + listener thread. On a large lane the first measured requests would otherwise wait + tens of seconds behind that queue (nvbugs/6564113: 32 gen ranks, about 1.4 s each + on oci-jhb). benchmark_serving's single test request (WARMUP_BENCHMARK_MODES) + reaches only one ctx/gen pair, so it cannot replace this. + + If only some requests fail, the warmup still opened the connections it reached, + so it logs a warning and the measured clients decide whether the test passes. The + test fails instead of letting the lane run cold or hang if benchmark_serving exits + non-zero, runs past CONNECTION_WARMUP_TIMEOUT (capped by the test timeout), or + every request fails. + + Args: + server_idx: Index of the disagg server under test. + hostname: Disagg server host. + port: Disagg server port. + measured_client: The client from _measured_client_for_connection_warmup. + + Returns: + The number of warmup requests that succeeded. + + Raises: + RuntimeError: benchmark_serving exited non-zero, timed out, or had every + request fail. + OSError: benchmark_serving could not be started or its log could not be opened. + """ + num_requests = self._connection_warmup_num_requests(server_idx) + cmd = add_host_port_to_cmd( + self._connection_warmup_cmd(measured_client, num_requests), hostname, port + ) + log_path = self._connection_warmup_log_path(server_idx) + timeout_seconds = min(CONNECTION_WARMUP_TIMEOUT, self.timeout) + # Reuse the measured client's env (proxy settings, HF offline flags, ...): + # without it the warmup may not reach the server or load the tokenizer. + env = copy.deepcopy(os.environ) + env.update(measured_client.to_env()) + + print_info(f"Connection warmup: sending {num_requests} short requests. cmd is {cmd}") + start = time.monotonic() + try: + output = _run_benchmark_with_log(cmd, env, log_path, timeout=timeout_seconds) + except subprocess.CalledProcessError as e: + raise RuntimeError( + f"Connection warmup: benchmark_serving exited with code {e.returncode}. " + f"See {log_path}" + ) from e + except subprocess.TimeoutExpired as e: + raise RuntimeError( + f"Connection warmup: {num_requests} short requests did not finish within " + f"{timeout_seconds}s, so the servers look wedged. See {log_path}" + ) from e + elapsed = time.monotonic() - start + failed_match = _FAILED_REQUESTS_RE.search(output) + failed_count = int(failed_match.group(1)) if failed_match else 0 + if failed_count >= num_requests: + raise RuntimeError( + f"Connection warmup: all {num_requests} requests failed. See {log_path}" + ) + if failed_count: + print_warning( + f"Connection warmup: {failed_count} of {num_requests} requests failed " + f"({elapsed:.0f}s); connections they missed will open during the measured " + f"run. See {log_path}" + ) + else: + print_info(f"Connection warmup: {num_requests} requests done in {elapsed:.0f}s") + return num_requests - failed_count + + def _wait_for_connection_warmup_line_counts( + self, num_succeeded_requests: int + ) -> Dict[str, int]: + """Wait for every warmup perf_metrics record to be written; return per-file line counts. + + The time breakdown skips this many leading lines per file, so a warmup record + written after this snapshot would count as measured. Each server queues a + request's record before it finishes answering that request, so every warmup record + is already queued when the warmup client exits; this waits for them to be written. + First, the largest file must reach num_succeeded_requests lines. That file belongs + to a server that saw every warmup request: the disagg server, or a lone ctx or gen + server. Then the files must hold still for PERF_METRICS_SETTLE_SECONDS, so the + other servers' queued records are written too. + + Args: + num_succeeded_requests: Warmup requests that succeeded. + + Returns: + Line count per perf_metrics JSONL basename. + + Raises: + RuntimeError: That did not happen within PERF_METRICS_SETTLE_TIMEOUT. + """ + _, settle_info = wait_for_perf_metrics_files( + self.perf_metrics_output_dir, + expected_requests=num_succeeded_requests, + wait_for_expected_requests=True, + stable_seconds=PERF_METRICS_SETTLE_SECONDS, + timeout_seconds=PERF_METRICS_SETTLE_TIMEOUT, + ) + if not settle_info["stable"]: + raise RuntimeError( + f"Connection warmup: expected {num_succeeded_requests} warmup records in the " + f"largest perf_metrics file under {self.perf_metrics_output_dir}, but the " + f"files did not reach that count and stop growing within " + f"{PERF_METRICS_SETTLE_TIMEOUT:.0f}s (line counts: " + f"{settle_info['line_counts']}). Failing so that late warmup records are " + "not counted as measured in d_tb_*." + ) + return settle_info["line_counts"] + def run_cmd(self, server_idx: int) -> List[str]: """Run commands for a server and return outputs.""" outputs = [] @@ -3143,7 +3409,10 @@ def run_cmd(self, server_idx: int) -> List[str]: # is the single master switch), so there is no second predicate to # keep in sync with it. collect_time_breakdown = bool(self.perf_metrics_output_dir) + # Filled only when the connection warmup ran with the time breakdown on. + connection_warmup_line_counts: Dict[str, int] = {} try: + self._remove_stale_connection_warmup_log(server_idx) disagg_server_hostname, disagg_server_port = ( self._get_disagg_server_hostname_and_port(server_idx) ) @@ -3158,6 +3427,18 @@ def run_cmd(self, server_idx: int) -> List[str]: observations = self._collect_worker_startup_observations(server_idx) write_startup_observations(self.test_output_dir, server_idx, observations) + measured_client = self._measured_client_for_connection_warmup(server_idx) + if measured_client is not None: + num_succeeded_connection_warmup_requests = self._run_connection_warmup( + server_idx, disagg_server_hostname, disagg_server_port, measured_client + ) + if collect_time_breakdown: + connection_warmup_line_counts = ( + self._wait_for_connection_warmup_line_counts( + num_succeeded_connection_warmup_requests + ) + ) + client_configs = self.client_configs.get(server_idx, []) # Run all clients for this server @@ -3299,7 +3580,11 @@ def run_cmd(self, server_idx: int) -> List[str]: # instead of rescanning the log. Modes outside the tuple leave the # queue empty and skip both steps. self._append_gen_worker_device_step_time(pending_device_step_time, outputs) - self._append_time_breakdown_metrics(pending_time_breakdown, outputs) + self._append_time_breakdown_metrics( + pending_time_breakdown, + outputs, + leading_lines_to_skip=connection_warmup_line_counts, + ) return outputs diff --git a/tests/integration/defs/perf/time_breakdown_metrics.py b/tests/integration/defs/perf/time_breakdown_metrics.py index 29d2e06ccde1..5bbe26532bb1 100644 --- a/tests/integration/defs/perf/time_breakdown_metrics.py +++ b/tests/integration/defs/perf/time_breakdown_metrics.py @@ -71,7 +71,7 @@ import statistics import time from collections import defaultdict -from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple +from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple STATS = ("mean", "median", "p75", "p99") METRIC_PREFIX = "d_tb_" @@ -264,9 +264,13 @@ def _classify(path: str, records: List[Dict[str, Any]]) -> str: return "ctx_worker" -def _read_jsonl(path: str) -> Tuple[List[Dict[str, Any]], int]: +def _read_jsonl(path: str, num_leading_lines_to_skip: int = 0) -> Tuple[List[Dict[str, Any]], int]: """Parse a JSONL file, skipping unparsable lines. Returns ``(records, skipped)``. + The first ``num_leading_lines_to_skip`` lines are ignored without being parsed and are not + counted in ``skipped``, which counts unparsable lines only; see + ``compute_time_breakdown_metrics`` for why they are ignored. + A malformed line is expected rather than exceptional: the client and every worker append to these files while the run is live, so the final line can be a partial write at the moment the aggregator reads (and, for a worker killed mid-flush, can @@ -279,7 +283,9 @@ def _read_jsonl(path: str) -> Tuple[List[Dict[str, Any]], int]: out: List[Dict[str, Any]] = [] skipped = 0 with open(path) as handle: - for line in handle: + for line_idx, line in enumerate(handle): + if line_idx < num_leading_lines_to_skip: + continue line = line.strip() if not line: continue @@ -486,6 +492,7 @@ def compute_time_breakdown_metrics( paths: Sequence[str], benchmark_mode: str, drop_warmup_request: bool = False, + leading_lines_to_skip: Optional[Mapping[str, int]] = None, ) -> Tuple[Dict[str, float], Dict[str, Any]]: """Reduce per-request JSONLs to ``{d_tb__: value_ms}``. @@ -501,6 +508,11 @@ def compute_time_breakdown_metrics( ``--no-test-input`` for ``e2e`` and ``ctx_only``). Discards that request's record so the breakdown covers the same population as every other metric on the row; see :func:`_drop_warmup_record`. Reported as ``info["warmup_dropped"]``. + leading_lines_to_skip: per file basename, the number of leading lines to ignore + because they record traffic sent before the measured clients (perf-sanity's + connection warmup). The servers append every request to the same files, so + those records would otherwise join the population. Reported as + ``info["leading_lines_skipped"]``. Returns: ``(metrics, info)``. ``metrics`` always has exactly ``len(ALL_METRICS) * 4`` keys; @@ -527,12 +539,17 @@ def compute_time_breakdown_metrics( gen_workers: List[Tuple[str, List[Dict[str, Any]]]] = [] skipped_lines: Dict[str, int] = {} + leading_lines_skipped: Dict[str, int] = {} + leading_lines_to_skip = leading_lines_to_skip or {} for path in paths: + num_leading_lines = leading_lines_to_skip.get(os.path.basename(path), 0) try: - records, skipped = _read_jsonl(path) + records, skipped = _read_jsonl(path, num_leading_lines_to_skip=num_leading_lines) except OSError as exc: warnings.append(f"{os.path.basename(path)}: unreadable ({exc})") continue + if num_leading_lines: + leading_lines_skipped[os.path.basename(path)] = num_leading_lines if skipped: skipped_lines[os.path.basename(path)] = skipped warnings.append( @@ -650,6 +667,9 @@ def compute_time_breakdown_metrics( "warmup_dropped": warmup_dropped, # Per file, non-empty only when a line failed to parse (usually a partial write). "skipped_lines": skipped_lines, + # Per file, how many leading lines were ignored as pre-measurement traffic (the + # connection warmup). + "leading_lines_skipped": leading_lines_skipped, "sample_counts": { name: len( (per_instance if name in CHUNK_METRICS + STEP_METRICS else per_request).get( @@ -722,6 +742,14 @@ def _count_lines(path: str) -> int: total += chunk.count(b"\n") +def _census(line_counts: Mapping[str, int], leading_lines_to_skip: Mapping[str, int]) -> int: + """Largest per-file line count after subtracting that file's skipped leading lines.""" + return max( + (count - leading_lines_to_skip.get(name, 0) for name, count in line_counts.items()), + default=0, + ) + + def wait_for_perf_metrics_files( output_dir: str, expected_requests: Optional[int] = None, @@ -730,8 +758,10 @@ def wait_for_perf_metrics_files( poll_seconds: float = COMPLETION_POLL_SECONDS, sleep=time.sleep, monotonic=time.monotonic, + leading_lines_to_skip: Optional[Mapping[str, int]] = None, + wait_for_expected_requests: bool = False, ) -> Tuple[List[str], Dict[str, Any]]: - """Wait for the perf_metrics JSONLs to stop growing, then return them. + """Wait for the perf_metrics JSONLs to hold still (and optionally to reach a count). Positive completion gate for the aggregation. The harness has no completion signal for the context workers or the disaggregated server -- unlike the generation workers, whose @@ -744,16 +774,30 @@ def wait_for_perf_metrics_files( (new files appearing counts as growth), bounded by ``timeout_seconds``. Then, if ``expected_requests`` is known, compare it against the largest file's complete-line count -- the disagg server and the aggregated server each write one record per request, - so that file is the run's census -- and report a shortfall. + so that file is the run's census -- and report a shortfall. Each file's + ``leading_lines_to_skip`` entry is subtracted from its count first, so earlier traffic + cannot mask a shortfall. ``info["line_counts"]`` stays raw (before the subtraction). + + With ``wait_for_expected_requests``, the files also need a census of at least + ``expected_requests`` before they count as settled, and an empty directory is polled + instead of returned at once. On timeout, ``info["stable"]`` is False and the caller + decides. The connection warmup's snapshot uses this mode because it starts as soon as + the warmup client exits. At that point the servers may not have written their last + records, or even created their files. ``sleep`` / ``monotonic`` are injected so the unit tests can drive this without wall time. - Returns ``(paths, info)``; ``info`` carries ``stable`` (bool), ``waited_seconds``, - ``total_bytes``, ``line_counts`` and ``warnings``. + Returns ``(paths, info)``; ``info`` carries ``stable`` (bool: the files held still for + ``stable_seconds`` and, with ``wait_for_expected_requests``, the census reached + ``expected_requests``), ``waited_seconds``, ``total_bytes``, ``line_counts`` and + ``warnings``. """ + if wait_for_expected_requests and not expected_requests: + raise ValueError("wait_for_expected_requests needs a positive expected_requests") warnings: List[str] = [] deadline = monotonic() + timeout_seconds + leading_lines_to_skip = leading_lines_to_skip or {} def snapshot() -> Tuple[List[str], Tuple[Tuple[str, Optional[int]], ...]]: paths = discover_perf_metrics_files(output_dir) @@ -767,12 +811,19 @@ def snapshot() -> Tuple[List[str], Tuple[Tuple[str, Optional[int]], ...]]: sizes.append((os.path.basename(path), None)) return paths, tuple(sizes) + def holds_expected_requests(paths: List[str]) -> bool: + try: + line_counts = {os.path.basename(path): _count_lines(path) for path in paths} + except OSError: + return False + return _census(line_counts, leading_lines_to_skip) >= expected_requests + started = monotonic() paths, fingerprint = snapshot() - if not paths: - # Nothing was ever created, so there is nothing to wait for: the workers write - # their first record long before the client exits. The caller reports the empty - # discovery itself. + if not paths and not wait_for_expected_requests: + # Without a count to wait for, no file means nothing was ever written: the workers + # write their first record long before the client exits. The caller reports the + # empty discovery itself. return [], { "stable": True, "waited_seconds": 0.0, @@ -785,14 +836,18 @@ def snapshot() -> Tuple[List[str], Tuple[Tuple[str, Optional[int]], ...]]: stable = False while True: now = monotonic() - if now - unchanged_since >= stable_seconds: + held_still = now - unchanged_since >= stable_seconds + if held_still and (not wait_for_expected_requests or holds_expected_requests(paths)): stable = True break if now >= deadline: - warnings.append( - f"perf_metrics files under {output_dir} were still growing after " - f"{timeout_seconds:.0f}s; aggregating what is on disk" - ) + # If the files held still, only a short census kept this waiting; the census + # check below reports it when any file exists. + if not held_still: + warnings.append( + f"perf_metrics files under {output_dir} were still growing after " + f"{timeout_seconds:.0f}s; aggregating what is on disk" + ) break sleep(min(poll_seconds, max(0.0, deadline - now))) paths, new_fingerprint = snapshot() @@ -808,7 +863,7 @@ def snapshot() -> Tuple[List[str], Tuple[Tuple[str, Optional[int]], ...]]: warnings.append(f"{os.path.basename(path)}: could not count lines ({exc})") if expected_requests and line_counts: - census = max(line_counts.values()) + census = _census(line_counts, leading_lines_to_skip) if census < expected_requests: warnings.append( f"the largest perf_metrics file holds {census} complete record(s) but the " @@ -870,6 +925,16 @@ def main() -> int: help="discard the client's un-measured warmup request (perf-sanity runs one for " "e2e and ctx_only); pass this to match what the harness uploads", ) + parser.add_argument( + "--leading-lines-to-skip", + type=json.loads, + default=None, + metavar="JSON", + help="per-file leading lines to ignore, as a JSON object of basename -> count; for " + "a disagg e2e run, pass the dict the harness logged after 'Time breakdown: excluded " + "connection-warmup records from' (with --drop-warmup-request) to match what it " + "uploads", + ) args = parser.parse_args() paths = list(args.input) @@ -879,7 +944,10 @@ def main() -> int: parser.error("no input: pass --input and/or --output-dir") metrics, info = compute_time_breakdown_metrics( - paths, args.mode, drop_warmup_request=args.drop_warmup_request + paths, + args.mode, + drop_warmup_request=args.drop_warmup_request, + leading_lines_to_skip=args.leading_lines_to_skip, ) if args.log_lines: