diff --git a/tests/integration/defs/perf/README.md b/tests/integration/defs/perf/README.md index 342b6cad5860..d857be62bcbf 100644 --- a/tests/integration/defs/perf/README.md +++ b/tests/integration/defs/perf/README.md @@ -32,6 +32,9 @@ If a CI run fails `sanity_perf_check.py`, it will upload a patch file as an arti This patch will only update the metrics that had a difference which was over the accepted threshold. The patch will also remove/add metrics according to the removed or added tests. +## Warmup pass +Before measuring, `test_perf.py` warms up with one un-measured wave of the benchmark's own traffic (its concurrency, or `min(reqs, maxbs)` without one), so JIT kernels for the real prefill shapes compile outside the measured window: `bench` cases raise `trtllm-bench --warmup` to that wave, and `serve` cases send it with other prompts before each client. `TRTLLM_PERF_WARMUP_PASS=0` turns it off. + ## Running locally Given a `target_perf_csv_path` you can compare it to another perf csv file. First make sure you install the dependencies: diff --git a/tests/integration/defs/perf/README_test_perf_sanity.md b/tests/integration/defs/perf/README_test_perf_sanity.md index 35e615a6b6c8..324d71200619 100644 --- a/tests/integration/defs/perf/README_test_perf_sanity.md +++ b/tests/integration/defs/perf/README_test_perf_sanity.md @@ -338,6 +338,17 @@ queue (nvbugs/6564113). Notes: connection-warmup iterations; start the range later to profile only measured iterations. +### Warmup pass + +Before each measured client of a normal aggregated, `ctx_only` or disaggregated `e2e` +test, the harness sends one un-measured wave (`concurrency` requests) of the client's +own `benchmark_serving` command with `--seed` + 1, so JIT kernels for the real prefill +shapes compile outside the measured window. `gen_only` modes, AgentX and accuracy-only +clients get none. Its log (`trtllm-benchmark-warmup-pass.*.log`) is not parsed, and its +records stay out of `d_tb_*` and the device step time; with the time breakdown on, only +the first measured client gets one. `TRTLLM_PERF_WARMUP_PASS=0` turns it off (locally, +via `EXTRA_CONTAINER_EXPORTS`). + ### The optional instrumentation modifier The three shapes that read a disaggregated config (2, 3 and 4 above) take one @@ -399,8 +410,9 @@ to run after the writers are done. Only the **gen** workers announce that workers and the disaggregated server do not. `wait_for_perf_metrics_files` therefore polls the discovered set until its total size holds still for `PERF_METRICS_SETTLE_SECONDS`, bounded by `PERF_METRICS_SETTLE_TIMEOUT`, before -anything is read, and then compares the largest file's complete-record count against -the client's `--num-prompts`. Both a still-growing set at the timeout and a census +anything is read, and then compares the census (the complete records of one server +role, summed over its files) against the client's `--num-prompts`. Both a +still-growing set at the timeout and a census shortfall are logged as warnings, not failures — a nearly-complete file still yields usable statistics, and the numbers are not regression-gated. This matters because a truncated read is otherwise **invisible**: all 108 fields are still populated and the @@ -409,13 +421,14 @@ 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 +The servers also append records of the e2e connection warmup's and the warmup pass's +requests to these files. The harness therefore waits until the census 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 +`Time breakdown: excluded 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 ''`. diff --git a/tests/integration/defs/perf/test_perf.py b/tests/integration/defs/perf/test_perf.py index 02419708bd7f..e2ee8cbd1f48 100644 --- a/tests/integration/defs/perf/test_perf.py +++ b/tests/integration/defs/perf/test_perf.py @@ -24,7 +24,7 @@ import os import re import shutil -from typing import Dict, List, NamedTuple +from typing import Dict, List, NamedTuple, Optional import pytest import yaml @@ -40,6 +40,8 @@ from .utils import (AbstractPerfScriptTestClass, PerfBenchScriptTestCmds, PerfMetricType, PerfServeScriptTestCmds, generate_test_nodes) +from .warmup_pass import (WarmupPass, build_warmup_pass_cmd, one_wave_requests, + warmup_pass_enabled) if not hasattr(re, "Pattern"): re.Pattern = type(re.compile("")) @@ -49,6 +51,14 @@ TIMING_CACHE_DIR = os.environ.get("TIMING_CACHE_DIR", "") +# trtllm-bench throughput's --warmup default. +TRTLLM_BENCH_DEFAULT_WARMUP = 2 + + +class NotEnoughSamplesError(ValueError): + """A real-text source holds fewer usable articles than a dataset file asks for.""" + + # Models served with the Nemotron reasoning/tool parsers, a long max model len # and a longer server-start timeout. "nano-v3" and "nemotron-v3" are two # registered names for the same reasoning parser, so Lightning belongs here too. @@ -1115,6 +1125,29 @@ def get_prepare_data_command(self, engine_dir, input_len, return data_cmd + def _get_bench_warmup_requests(self) -> int: + """Return trtllm-bench's --warmup for this case, or 0 for its default. + + trtllm-bench runs these requests in its own process before the measured + ones, at the same concurrency, and reports none of them. With the warmup + pass on, they are one wave of the measured traffic (see warmup_pass.py), + and never fewer than without it. They are drawn from the measured + dataset, though, so a case whose config reuses KV cache blocks keeps the + fewer requests it has without the pass: otherwise its measured requests + would hit the warmup's cached prefixes. + """ + without_pass = 2 * self._config.num_gpus if self._config.num_gpus > 1 else 0 + if not warmup_pass_enabled(): + return without_pass + kv_cache_config = self._get_model_yaml_config().get( + 'kv_cache_config') or {} + if kv_cache_config.get('enable_block_reuse'): + return without_pass + wave = one_wave_requests(self._config.concurrency, + self._config.num_reqs, + self._config.max_batch_size) + return max(without_pass or TRTLLM_BENCH_DEFAULT_WARMUP, wave) + def get_trtllm_bench_command(self, engine_dir): model_dir = self.get_trtllm_bench_model() model_name = self._config.model_name @@ -1149,8 +1182,9 @@ def get_trtllm_bench_command(self, engine_dir): benchmark_cmd += [f"--pp={self._config.pp_size}"] if self._config.streaming == "streaming": benchmark_cmd += [f"--streaming"] - if self._config.num_gpus > 1: - benchmark_cmd += [f"--warmup={2 * self._config.num_gpus}"] + warmup_requests = self._get_bench_warmup_requests() + if warmup_requests: + benchmark_cmd += [f"--warmup={warmup_requests}"] # Add extra-llm-api-config.yml for the PyTorch backend. if self._config.backend == "pytorch": @@ -1227,9 +1261,20 @@ def get_trtllm_serve_server_command(self, engine_dir): return server_cmd - def generate_trtllm_custom_dataset(self, dst_dataset_path: str, - input_len: int, output_len: int, - dataset_source: str): + def generate_trtllm_custom_dataset(self, + dst_dataset_path: str, + input_len: int, + output_len: int, + dataset_source: str, + num_requests: Optional[int] = None, + skip_requests: int = 0): + """Write a real-text dataset of fixed-length prompts. + + Args: + num_requests: Samples to write; the case's request count if None. + skip_requests: Leading usable articles to skip, so that another + file can hold prompts this one does not. + """ # Currently only support cnn_dailymail dataset source. if dataset_source != "cnn_dailymail": raise ValueError( @@ -1260,8 +1305,9 @@ def generate_trtllm_custom_dataset(self, dst_dataset_path: str, if not os.path.exists(os.path.dirname(dst_dataset_path)): os.makedirs(os.path.dirname(dst_dataset_path), exist_ok=True) - num_reqs = self._config.num_reqs + num_reqs = self._config.num_reqs if num_requests is None else num_requests req_count = 0 + num_skipped = 0 with open(dst_dataset_path, "w", encoding="utf-8") as f: for req in dataset: article = req.get("article") @@ -1281,6 +1327,9 @@ def generate_trtllm_custom_dataset(self, dst_dataset_path: str, prompt_ids = (prompt_ids * repeats)[:input_len] elif len(prompt_ids) > input_len: prompt_ids = prompt_ids[:input_len] + if num_skipped < skip_requests: + num_skipped += 1 + continue prompt_text = tokenizer.decode(prompt_ids, skip_special_tokens=False) if self._config.model_name in DEEPSEEK_V41_MODELS: @@ -1322,7 +1371,7 @@ def generate_trtllm_custom_dataset(self, dst_dataset_path: str, break if req_count < num_reqs: - raise ValueError( + raise NotEnoughSamplesError( f"Cannot sample enough requests from cnn_dailymail: requested={num_reqs}, sampled={req_count}" ) print_info(f"Generated {req_count} samples from {dataset_source} to " @@ -1416,6 +1465,52 @@ def get_trtllm_serve_client_command(self, client_cmd += ["--trust-remote-code"] return client_cmd + def _get_serve_warmup_pass(self, engine_dir: str, input_len: int, + output_len: int, + client_cmd: List[str]) -> Optional[WarmupPass]: + """Return the warmup pass before a serve client, or None; see warmup_pass.py. + + A real-text client gets a prompt file of its own, with the articles + after the measured ones. The measured file holds exactly the measured + prompts, so another shuffle of it would resend them, and KV cache block + reuse would then serve their prefixes to the measured requests. If the + source has no articles to spare, the client gets no pass. + """ + if not warmup_pass_enabled(): + return None + num_requests = one_wave_requests(self._config.concurrency, + self._config.num_reqs, + self._config.max_batch_size) + if num_requests < 1: + return None + dataset_source = SPEC_DEC_REAL_DATASET_MODELS.get( + self._config.model_name) + if dataset_source is not None: + warmup_dataset_path = os.path.join( + engine_dir, + f"dataset_custom_{input_len}_{output_len}_warmup_pass.jsonl") + try: + self.generate_trtllm_custom_dataset( + warmup_dataset_path, + input_len, + output_len, + dataset_source=dataset_source, + num_requests=num_requests, + skip_requests=self._config.num_reqs) + except NotEnoughSamplesError as error: + print_warning( + f"Warmup pass: none for this client, since {dataset_source} " + f"has no {num_requests} articles to spare after the " + f"{self._config.num_reqs} measured ones: {error}") + return None + client_cmd = self.get_trtllm_serve_client_command( + engine_dir, + input_len, + output_len, + real_dataset_path=warmup_dataset_path) + return WarmupPass(build_warmup_pass_cmd(client_cmd, num_requests), + num_requests) + def get_commands(self): num_gpus = self._config.num_gpus @@ -1428,6 +1523,7 @@ def get_commands(self): if self._config.runtime == "serve": server_cmd = self.get_trtllm_serve_server_command(engine_dir) client_cmds = [] + warmup_passes = [] data_cmds = [] for bs in self._config.batch_sizes: for len_idx, input_len in enumerate(self._config.input_lens): @@ -1453,6 +1549,9 @@ def get_commands(self): output_len, real_dataset_path=real_dataset_path) client_cmds.append(client_cmd) + warmup_passes.append( + self._get_serve_warmup_pass(engine_dir, input_len, + output_len, client_cmd)) server_env = os.environ.copy() if self._config.model_name.startswith("qwen3.8_flash_next_"): server_env["TRTLLM_QWEN4_EXP_PLE_HOST_OFFLOAD"] = "1" @@ -1490,7 +1589,8 @@ def get_commands(self): client_cmds=client_cmds, data_cmds=data_cmds, server_env=server_env, - server_timeout=server_timeout) + server_timeout=server_timeout, + warmup_passes=warmup_passes) # PyTorch loads model checkpoints directly. build_cmd = [] diff --git a/tests/integration/defs/perf/test_perf_sanity.py b/tests/integration/defs/perf/test_perf_sanity.py index 08b4c5fb6539..556021dd2644 100644 --- a/tests/integration/defs/perf/test_perf_sanity.py +++ b/tests/integration/defs/perf/test_perf_sanity.py @@ -54,6 +54,14 @@ format_metric_log_lines, wait_for_perf_metrics_files, ) +from .warmup_pass import ( + build_warmup_pass_cmd, + count_failed_requests, + send_warmup_pass, + warmup_pass_applies, + warmup_pass_enabled, + warmup_pass_plan, +) SUPPORTED_GPU_MAPPING = { "GB200": "gb200", @@ -70,7 +78,7 @@ # 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). +# (DisaggTestCmds._run_connection_warmup) and to the warmup pass (warmup_pass.py). WARMUP_BENCHMARK_MODES = ("e2e", "ctx_only") @@ -217,9 +225,6 @@ def ensure_bench_serving_repo() -> str: # 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: @@ -1776,6 +1781,158 @@ def _run_benchmark_with_log( return raw_output.decode("utf-8", errors="replace") +def warmup_pass_log_path(test_output_dir: str, server_idx: int, client_idx: int) -> str: + """Log of the warmup pass that precedes client ``client_idx`` of server ``server_idx``.""" + return os.path.join( + test_output_dir, f"trtllm-benchmark-warmup-pass.{server_idx}.{client_idx}.log" + ) + + +def warmup_pass_logs(test_output_dir: str, server_idx: int) -> List[str]: + """Every warmup pass log present for server ``server_idx``.""" + return sorted( + glob.glob(os.path.join(test_output_dir, f"trtllm-benchmark-warmup-pass.{server_idx}.*.log")) + ) + + +def warmup_pass_requests_for( + client_config: Optional["ClientConfig"], benchmark_mode: Optional[str] +) -> int: + """Requests in the warmup pass before this client, or 0 when it gets none. + + Args: + client_config: The measured client's config; None gets no pass. + benchmark_mode: The case's benchmark mode; None or "" for a plain + aggregated case. + """ + if ( + client_config is None + or not warmup_pass_enabled() + or not warmup_pass_applies(benchmark_mode) + ): + return 0 + return client_config.warmup_pass_requests + + +def warmup_pass_plan_for( + client_configs: List[Optional["ClientConfig"]], + benchmark_mode: Optional[str], + time_breakdown: bool, +) -> List[int]: + """Warmup pass size for each client of one server, in run order; 0 for none. + + The run loops send these passes and the uploaded rows report them. While the + time breakdown is collected, only the first measured client gets one; see + warmup_pass_plan. + + Args: + client_configs: The server's client configs in run order; None entries are + measured clients without a config. + benchmark_mode: The case's benchmark mode; None or "" for a plain + aggregated case. + time_breakdown: Whether the time breakdown is collected for the server. + """ + return warmup_pass_plan( + [warmup_pass_requests_for(config, benchmark_mode) for config in client_configs], + [config is None or not config.only_run_accuracy for config in client_configs], + time_breakdown, + ) + + +def warn_skipped_warmup_pass(client_idx: int) -> None: + """Log that a client gets no warmup pass because the time breakdown is on.""" + print_warning( + f"Warmup pass: none before client {client_idx}. With the time breakdown on, only " + "the first measured client gets one: the breakdown cannot leave out the records " + "of a later pass, so they would count as measured." + ) + + +def run_warmup_pass( + client_cmd: List[str], + client_config: "ClientConfig", + num_requests: int, + hostname: str, + port: int, + log_path: str, + timeout: float, +) -> int: + """Send the warmup pass for one measured client; see warmup_pass.py. + + Args: + client_cmd: The measured client's command, without host and port. + client_config: The measured client's config; its env is reused. + num_requests: Requests to send, from warmup_pass_requests_for. + hostname: Server host. + port: Server port. + log_path: File that receives the pass's output. + timeout: Seconds before the pass is killed. + + Returns: + The number of warmup requests that succeeded. + + Raises: + RuntimeError: benchmark_serving exited non-zero, timed out, or had every + request fail. + """ + env = copy.deepcopy(os.environ) + env.update(client_config.to_env()) + return send_warmup_pass( + add_host_port_to_cmd(build_warmup_pass_cmd(client_cmd, num_requests), hostname, port), + num_requests, + run_benchmark=lambda cmd: _run_benchmark_with_log(cmd, env, log_path, timeout=timeout), + log_path=log_path, + timeout=timeout, + log_info=print_info, + log_warning=print_warning, + ) + + +def wait_for_unmeasured_line_counts( + perf_metrics_output_dir: str, num_unmeasured_requests: int +) -> Dict[str, int]: + """Wait for the un-measured requests' perf_metrics records; 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, one server role must hold num_unmeasured_requests records, + summed over its files: each role (the disagg server, the ctx workers, the gen + workers, or an aggregated server) sees every request, but writes one file per + instance and per frontend process. Then the files must hold still for + PERF_METRICS_SETTLE_SECONDS, so the other roles' queued records are written + too. + + Args: + perf_metrics_output_dir: Directory the servers write perf_metrics JSONLs to. + num_unmeasured_requests: Requests that succeeded before the measured client + started (connection warmup and warmup pass). + + 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( + perf_metrics_output_dir, + expected_requests=num_unmeasured_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"Expected {num_unmeasured_requests} warmup records from one server role " + f"under {perf_metrics_output_dir}, but the perf_metrics files did not reach " + f"that count and stop growing within {PERF_METRICS_SETTLE_TIMEOUT:.0f}s " + f"(line counts: {settle_info['line_counts']}). Failing so that late warmup " + "records are not counted as measured in d_tb_*." + ) + return settle_info["line_counts"] + + class ServerConfig: """Configurations of trtllm-server.""" @@ -2254,6 +2411,20 @@ def num_requests(self) -> int: """ return self.concurrency * self.iterations + @property + def warmup_pass_requests(self) -> int: + """Requests in the warmup pass before this client, or 0 when it gets none. + + One wave at the client's concurrency, so the pass prefills the same batches + the measured run starts with. Only benchmark_serving clients get one: + AgentX runs its own warmup phase, and an accuracy-only client measures no + performance. Whether the case's mode gets a pass at all is decided by + warmup_pass_applies and warmup_pass_enabled. + """ + if self.only_run_accuracy or self.benchmark_client or self.use_nv_sa_benchmark: + return 0 + return self.concurrency + def to_cmd(self) -> List[str]: """Generate benchmark command.""" model_dir = get_model_dir(self.model_name) @@ -2506,8 +2677,8 @@ def append_time_breakdown_metrics( 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. + before the measured clients started (the connection warmup's and the warmup + pass's records); those lines are left out of both the census and the breakdown. """ if not pending_time_breakdown: return @@ -2567,7 +2738,7 @@ def append_time_breakdown_metrics( 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 " + "Time breakdown: excluded warmup records from " f"{json.dumps(info['leading_lines_skipped'])}" ) @@ -2604,7 +2775,11 @@ def get_server_logs(self, server_idx) -> List[str]: benchmark_logs = sorted( glob.glob(os.path.join(self.test_output_dir, f"trtllm-benchmark.{server_idx}.*.log")) ) - return [server_file_path, *benchmark_logs] + return [ + server_file_path, + *benchmark_logs, + *warmup_pass_logs(self.test_output_dir, server_idx), + ] def run_cmd(self, server_idx: int) -> List[str]: """Run all clients for a server and return outputs. @@ -2624,6 +2799,13 @@ def run_cmd(self, server_idx: int) -> List[str]: # after the server exits -- i.e. after the finally below, not before it. pending_time_breakdown: List[dict] = [] collect_time_breakdown = bool(self.perf_metrics_output_dir) + # Requests the warmup passes completed before the first measured client. + num_unmeasured_requests = 0 + # Filled only when a warmup pass ran with the time breakdown on. + unmeasured_line_counts: Dict[str, int] = {} + unmeasured_lines_counted = False + for stale_log in warmup_pass_logs(self.test_output_dir, server_idx): + os.remove(stale_log) try: server_hostname = "localhost" @@ -2684,14 +2866,46 @@ def run_cmd(self, server_idx: int) -> List[str]: ) write_startup_observations(self.test_output_dir, server_idx, [observation]) + client_cmds = self.client_cmds[server_idx] + loop_configs = [ + client_configs[client_idx] if client_idx < len(client_configs) else None + for client_idx in range(len(client_cmds)) + ] + warmup_plan = warmup_pass_plan_for( + loop_configs, self.benchmark_mode, collect_time_breakdown + ) # Run all clients for this server - for client_idx, client_cmd in enumerate(self.client_cmds[server_idx]): - client_config = ( - client_configs[client_idx] if client_idx < len(client_configs) else None - ) + for client_idx, client_cmd in enumerate(client_cmds): + client_config = loop_configs[client_idx] only_run_accuracy = bool(client_config and client_config.only_run_accuracy) if not only_run_accuracy: + num_warmup_requests = warmup_plan[client_idx] + if num_warmup_requests: + num_unmeasured_requests += run_warmup_pass( + client_cmd, + client_config, + num_warmup_requests, + server_hostname, + server_port, + warmup_pass_log_path(self.test_output_dir, server_idx, client_idx), + timeout=min(self.timeout, DEFAULT_TIMEOUT), + ) + elif warmup_pass_requests_for(client_config, self.benchmark_mode): + warn_skipped_warmup_pass(client_idx) + # The time breakdown skips only leading lines, so count them + # once, before the first measured client. warmup_plan gives + # the later clients no pass while it is collected. + if ( + collect_time_breakdown + and num_unmeasured_requests + and not unmeasured_lines_counted + ): + unmeasured_line_counts = wait_for_unmeasured_line_counts( + self.perf_metrics_output_dir, num_unmeasured_requests + ) + unmeasured_lines_counted = True + client_file_path = os.path.join( self.test_output_dir, f"trtllm-benchmark.{server_idx}.{client_idx}.log" ) @@ -2749,7 +2963,12 @@ def run_cmd(self, server_idx: int) -> List[str]: # complete. terminate() is SIGTERM, which trtllm-serve handles as a # graceful shutdown, so PerfMetricsJsonlWriter.close() has run and the # tail of the queue is on disk. - 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=unmeasured_line_counts, + ) return outputs @@ -3179,6 +3398,7 @@ def get_server_logs(self, server_idx: int) -> List[str]: connection_warmup_log = self._connection_warmup_log_path(server_idx) if os.path.exists(connection_warmup_log): server_logs.append(connection_warmup_log) + server_logs.extend(warmup_pass_logs(self.test_output_dir, server_idx)) return server_logs @staticmethod @@ -3344,8 +3564,7 @@ def _run_connection_warmup( 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 + failed_count = count_failed_requests(output) if failed_count >= num_requests: raise RuntimeError( f"Connection warmup: all {num_requests} requests failed. See {log_path}" @@ -3360,47 +3579,6 @@ def _run_connection_warmup( 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"] - @staticmethod def _dwdp_ctx_rank(saved_env: Dict[str, str]) -> Tuple[int, int]: """(rank, world size) of this task inside the DWDP ctx MPI step.""" @@ -3665,10 +3843,16 @@ 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] = {} + # Requests that succeeded before the measured client started: the + # connection warmup's and the warmup pass's. + num_unmeasured_requests = 0 + # Filled only when un-measured requests ran with the time breakdown on. + unmeasured_line_counts: Dict[str, int] = {} + unmeasured_lines_counted = False try: self._remove_stale_connection_warmup_log(server_idx) + for stale_log in warmup_pass_logs(self.test_output_dir, server_idx): + os.remove(stale_log) disagg_server_hostname, disagg_server_port = ( self._get_disagg_server_hostname_and_port(server_idx) ) @@ -3685,26 +3869,62 @@ def run_cmd(self, server_idx: int) -> List[str]: 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( + num_unmeasured_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, []) + client_cmds = self.client_cmds[server_idx] + loop_configs = [ + client_configs[client_idx] if client_idx < len(client_configs) else None + for client_idx in range(len(client_cmds)) + ] + warmup_plan = warmup_pass_plan_for( + loop_configs, benchmark_mode_for_idx, collect_time_breakdown + ) # Run all clients for this server - for client_idx, client_cmd in enumerate(self.client_cmds[server_idx]): - client_config = ( - client_configs[client_idx] if client_idx < len(client_configs) else None - ) + for client_idx, client_cmd in enumerate(client_cmds): + client_config = loop_configs[client_idx] only_run_accuracy = bool(client_config and client_config.only_run_accuracy) if not only_run_accuracy: + num_warmup_requests = warmup_plan[client_idx] + if not num_warmup_requests and warmup_pass_requests_for( + client_config, benchmark_mode_for_idx + ): + warn_skipped_warmup_pass(client_idx) + if num_warmup_requests: + # The pass runs between two clients' gen-log windows, + # so neither one's device step time counts its + # iterations: close the previous client's window + # here, and open this client's after the pass. + if collect_device_step_time and pending_device_step_time: + pending_device_step_time[-1]["end_offsets"] = gen_worker_log_sizes( + self.test_output_dir, self.num_gen_servers + ) + num_unmeasured_requests += run_warmup_pass( + client_cmd, + client_config, + num_warmup_requests, + disagg_server_hostname, + disagg_server_port, + warmup_pass_log_path(self.test_output_dir, server_idx, client_idx), + timeout=min(self.timeout, DEFAULT_TIMEOUT), + ) + # The time breakdown skips only leading lines, so count + # them once, before the first measured client. warmup_plan + # gives the later clients no pass while it is collected. + if ( + collect_time_breakdown + and num_unmeasured_requests + and not unmeasured_lines_counted + ): + unmeasured_line_counts = wait_for_unmeasured_line_counts( + self.perf_metrics_output_dir, num_unmeasured_requests + ) + unmeasured_lines_counted = True + benchmark_file_path = os.path.join( self.test_output_dir, f"trtllm-benchmark.{server_idx}.{client_idx}.log" ) @@ -3714,11 +3934,12 @@ def run_cmd(self, server_idx: int) -> List[str]: print_info(f"Starting benchmark. cmd is {client_cmd_with_port}") # Snapshot gen_server log sizes so each client's stats - # cover only iterations driven by that client. This is - # also the *end* bound of the previous client's window - # (see the fixup below): taken here, it is necessarily - # after that client returned, so it absorbs whatever - # the gen workers flushed late. Modes outside + # cover only iterations driven by that client. Unless a + # warmup pass closed it above, this is also the *end* + # bound of the previous client's window (see the fixup + # below): taken here, it is necessarily after that + # client returned, so it absorbs whatever the gen + # workers flushed late. Modes outside # DEVICE_STEP_TIME_MODES skip this and must not wait for # the GEN teardown sentinel. gen_log_start_offsets = None @@ -3726,7 +3947,10 @@ def run_cmd(self, server_idx: int) -> List[str]: gen_log_start_offsets = gen_worker_log_sizes( self.test_output_dir, self.num_gen_servers ) - if pending_device_step_time: + if ( + pending_device_step_time + and pending_device_step_time[-1]["end_offsets"] is None + ): # Close the previous client's window here rather # than at its own return: this snapshot is the # first byte of the current client's segment, so @@ -3839,7 +4063,7 @@ def run_cmd(self, server_idx: int) -> List[str]: self._append_time_breakdown_metrics( pending_time_breakdown, outputs, - leading_lines_to_skip=connection_warmup_line_counts, + leading_lines_to_skip=unmeasured_line_counts, ) return outputs @@ -5202,6 +5426,9 @@ def add_dict_prefix(config_dict: dict, prefix_name: str) -> dict: if len(server_perf_results) != len(client_configs): cmd_idx += len(client_configs) continue + warmup_plan = warmup_pass_plan_for( + client_configs, self.benchmark_mode, self.time_breakdown + ) for client_idx, client_config in enumerate(client_configs): client_config_dict = client_config.to_db_data() @@ -5222,6 +5449,7 @@ def add_dict_prefix(config_dict: dict, prefix_name: str) -> dict: } new_data.update(server_config_dict) new_data.update(client_config_dict) + new_data["l_warmup_pass_requests"] = warmup_plan[client_idx] # Add test_case_name for convenient filtering on OpenSearch new_data["s_test_case_name"] = f"{server_config.name}-{client_config.name}" @@ -5275,6 +5503,9 @@ def add_dict_prefix(config_dict: dict, prefix_name: str) -> dict: if len(server_perf_results) != len(client_configs): cmd_idx += len(client_configs) continue + warmup_plan = warmup_pass_plan_for( + client_configs, disagg_config.benchmark_mode, self.time_breakdown + ) for client_idx, client_config in enumerate(client_configs): # Skip if metrics missing @@ -5311,6 +5542,7 @@ def add_dict_prefix(config_dict: dict, prefix_name: str) -> dict: if num_gen_servers > 0: new_data.update(gen_server_config_dict) new_data.update(client_config_dict) + new_data["l_warmup_pass_requests"] = warmup_plan[client_idx] # Add test_case_name for convenient filtering on OpenSearch new_data["s_test_case_name"] = f"{disagg_config.name}-{client_config.name}" diff --git a/tests/integration/defs/perf/time_breakdown_metrics.py b/tests/integration/defs/perf/time_breakdown_metrics.py index 5bbe26532bb1..7da698af72f6 100644 --- a/tests/integration/defs/perf/time_breakdown_metrics.py +++ b/tests/integration/defs/perf/time_breakdown_metrics.py @@ -510,8 +510,8 @@ def compute_time_breakdown_metrics( :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 + connection warmup and warmup pass). The servers append every request to the + same files, so those records would otherwise join the population. Reported as ``info["leading_lines_skipped"]``. Returns: @@ -668,7 +668,7 @@ def compute_time_breakdown_metrics( # 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). + # connection warmup and the warmup pass). "leading_lines_skipped": leading_lines_skipped, "sample_counts": { name: len( @@ -742,12 +742,53 @@ 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, - ) +# Records _census_role parses to tell a file's role. A file holds one role's records, so a +# few suffice, and a gen worker's records carry an entry per decode step, so they can be large. +_CENSUS_ROLE_SAMPLE_RECORDS = 8 + + +def _census_role(path: str) -> str: + """``_classify`` of the file's first complete records; ``empty`` until it has one.""" + records: List[Dict[str, Any]] = [] + with open(path, "rb") as handle: + for line in handle: + if not line.endswith(b"\n"): + break # A partial write. + try: + record = json.loads(line) + except ValueError: + continue + if isinstance(record, dict): + records.append(record) + if len(records) == _CENSUS_ROLE_SAMPLE_RECORDS: + break + return _classify(path, records) + + +def _census( + line_counts: Mapping[str, int], + leading_lines_to_skip: Mapping[str, int], + roles: Mapping[str, str], +) -> int: + """Complete records of the role that wrote the most, after each file's skipped lines. + + The disagg server, the ctx workers and the gen workers each write one record per + request, and so does an aggregated server, so any one role's records are the run's + census. A role writes several files when it has several instances, or when an instance + runs several frontend processes (``num_serve_frontends``): each process writes its own + file, which holds only its share of the requests. So the census sums each role's files. + + Args: + line_counts: Complete lines per file basename. + leading_lines_to_skip: Lines per file basename to leave out of the count. + roles: ``_census_role`` per file basename; ``empty`` files are left out. + """ + totals: Dict[str, int] = defaultdict(int) + for name, count in line_counts.items(): + role = roles.get(name, "empty") + if role != "empty": + totals[role] += max(0, count - leading_lines_to_skip.get(name, 0)) + return max(totals.values(), default=0) def wait_for_perf_metrics_files( @@ -772,18 +813,19 @@ def wait_for_perf_metrics_files( Poll the discovered set's total byte size until it holds still for ``stable_seconds`` (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. 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). + ``expected_requests`` is known, compare it against the census -- the complete records + of the role that wrote the most, summed over that role's files (see ``_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. + decides. The harness snapshots its warmup traffic (the connection warmup and the + warmup pass) in this mode, because the snapshot starts as soon as the last 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. @@ -811,12 +853,25 @@ def snapshot() -> Tuple[List[str], Tuple[Tuple[str, Optional[int]], ...]]: sizes.append((os.path.basename(path), None)) return paths, tuple(sizes) + # A file's role never changes once it has a record, so each file is read once for it. + roles: Dict[str, str] = {} + + def census(paths: List[str], line_counts: Mapping[str, int]) -> int: + for path in paths: + name = os.path.basename(path) + if roles.get(name, "empty") == "empty": + try: + roles[name] = _census_role(path) + except OSError: + roles[name] = "empty" + return _census(line_counts, leading_lines_to_skip, roles) + 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 + return census(paths, line_counts) >= expected_requests started = monotonic() paths, fingerprint = snapshot() @@ -863,12 +918,12 @@ def holds_expected_requests(paths: List[str]) -> bool: warnings.append(f"{os.path.basename(path)}: could not count lines ({exc})") if expected_requests and line_counts: - census = _census(line_counts, leading_lines_to_skip) - if census < expected_requests: + num_recorded = census(paths, line_counts) + if num_recorded < expected_requests: warnings.append( - f"the largest perf_metrics file holds {census} complete record(s) but the " - f"client issued {expected_requests} request(s); the breakdown covers a " - "subset of the run" + f"the server role with the most perf_metrics records holds {num_recorded} " + f"complete record(s) but the client issued {expected_requests} request(s); " + "the breakdown covers a subset of the run" ) info = { @@ -931,8 +986,8 @@ def main() -> int: 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 " + "a run with warmup traffic, pass the dict the harness logged after 'Time breakdown: " + "excluded warmup records from' (with --drop-warmup-request) to match what it " "uploads", ) args = parser.parse_args() diff --git a/tests/integration/defs/perf/utils.py b/tests/integration/defs/perf/utils.py index d5b3cab6e4fc..6f6b37ec9759 100644 --- a/tests/integration/defs/perf/utils.py +++ b/tests/integration/defs/perf/utils.py @@ -32,7 +32,7 @@ import pytest from defs.trt_test_alternative import (Popen, cleanup_process_tree, print_error, - print_info) + print_info, print_warning) from ..common import get_trt_llm_lib_dir from ..local_venv import PythonVenvRunnerImpl @@ -41,6 +41,7 @@ write_gpu_monitoring_no_test_results, write_yaml) from .gpu_clock_lock import GPUClockLock, InvalidGPUMonitoringResultError from .session_data_writer import SessionDataWriter +from .warmup_pass import WarmupPass, send_warmup_pass def clean_myelin_time(log: str): @@ -371,14 +372,20 @@ def _wait_for_server_ready(url: str, class PerfServeScriptTestCmds: """Commands for serve runtime perf tests (server-client model).""" - def __init__(self, server_cmd: List[str], client_cmds: List[List[str]], - data_cmds: List[List[str]], server_env: Dict[str, str], - server_timeout: int): + def __init__(self, + server_cmd: List[str], + client_cmds: List[List[str]], + data_cmds: List[List[str]], + server_env: Dict[str, str], + server_timeout: int, + warmup_passes: Optional[List[Optional[WarmupPass]]] = None): self.server_cmd = server_cmd self.client_cmds = client_cmds self.data_cmds = data_cmds self.server_env = server_env self.server_timeout = server_timeout + # Per client, the warmup pass to send before it, or None. + self.warmup_passes = warmup_passes or [None] * len(client_cmds) self._server_proc = None self._server_log_path = None @@ -462,11 +469,12 @@ def run_cmd(self, cmd_idx: int, venv) -> str: # instead of an empty string. output = self.get_server_log_content() else: - client_cmd = self.client_cmds[cmd_idx - 1 - len(self.data_cmds)] - client_cmd_with_port = client_cmd + [ - "--host", self._host, "--port", - str(self._port) - ] + client_idx = cmd_idx - 1 - len(self.data_cmds) + host_port = ["--host", self._host, "--port", str(self._port)] + warmup_pass = self.warmup_passes[client_idx] + if warmup_pass is not None: + self._send_warmup_pass(warmup_pass, host_port) + client_cmd_with_port = self.client_cmds[client_idx] + host_port print_info( f"Running benchmark client: {' '.join(client_cmd_with_port)}") output = _run_command_with_captured_output(client_cmd_with_port, @@ -474,6 +482,38 @@ def run_cmd(self, cmd_idx: int, venv) -> str: os.environ)) return output + @staticmethod + def _send_warmup_pass(warmup_pass: WarmupPass, + host_port: List[str]) -> None: + """Send a client's warmup pass; its output goes to a log of its own.""" + # Unique per pass: cluster jobs can share a working directory. + with tempfile.NamedTemporaryFile(mode="w", + prefix="benchmark-warmup-pass-", + suffix=".log", + dir=os.getcwd(), + delete=False) as log_file: + log_path = log_file.name + + def run_benchmark(cmd: List[str]) -> str: + try: + output = _run_command_with_captured_output(cmd, + env=copy.deepcopy( + os.environ)) + except subprocess.CalledProcessError as e: + with open(log_path, "wb") as log: + log.write(e.stdout or b"") + raise + with open(log_path, "w", encoding="utf-8") as log: + log.write(output) + return output + + send_warmup_pass(warmup_pass.cmd + host_port, + warmup_pass.num_requests, + run_benchmark, + log_path=log_path, + log_info=print_info, + log_warning=print_warning) + def get_cmd_str(self, cmd_idx) -> str: if cmd_idx <= len(self.data_cmds) - 1: return " ".join(self.data_cmds[cmd_idx]) diff --git a/tests/integration/defs/perf/warmup_pass.py b/tests/integration/defs/perf/warmup_pass.py new file mode 100644 index 000000000000..614827b93958 --- /dev/null +++ b/tests/integration/defs/perf/warmup_pass.py @@ -0,0 +1,263 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The un-measured warmup pass the perf tests send before each measured client. + +The server's own startup warmup cannot reach every prefill shape a benchmark +produces: how many prompts arrive together, how they are chunked against +``max_num_tokens`` and which decode tokens share a step with them are only known +once real traffic arrives. JIT kernels keyed on those shapes would otherwise +compile inside the first measured requests. The warmup pass sends that traffic +once beforehand: the measured client's own ``benchmark_serving`` command, cut to +one wave of requests with different prompts. Nothing it prints is parsed or +uploaded. + +perf-sanity (``test_perf_sanity.py``) runs it for plain aggregated cases, +``ctx_only`` and disagg ``e2e``, whose measured window contains prefill. The +``gen_only`` modes never get one: they measure decode only, and their generation +server's benchmark fill gate opens once, so warmup requests would open it before +the measured burst arrives. + +The single-node perf tests (``test_perf.py``) run it for both runtimes. A +``serve`` case sends it before each measured client. A ``bench`` case measures +inside the ``trtllm-bench throughput`` process that loads the model, so its pass +is that command's own ``--warmup`` requests, raised to one wave. + +``TRTLLM_PERF_WARMUP_PASS=0`` turns it off in both, e.g. to reproduce a JIT stall +the pass would otherwise absorb. +""" + +import os +import re +import subprocess +import time +from typing import Callable, List, Mapping, NamedTuple, Optional, Sequence + +WARMUP_PASS_ENV = "TRTLLM_PERF_WARMUP_PASS" + +# Benchmark modes that get a warmup pass. A plain aggregated case has no +# benchmark mode, which the runner carries as "". +WARMUP_PASS_MODES = ("", "ctx_only", "e2e") + +_DISABLED_VALUES = ("0", "false", "no", "off") + +_NUM_PROMPTS_FLAG = "--num-prompts" +_SEED_FLAG = "--seed" +_NO_TEST_INPUT_FLAG = "--no-test-input" +# Flags whose value the warmup pass replaces or drops: the request count, the seed, +# and where the measured client writes its time breakdown and result file. +_FLAGS_WITH_REPLACED_VALUE = ( + _NUM_PROMPTS_FLAG, + _SEED_FLAG, + "--save-request-time-breakdown", + "--result-dir", + "--result-filename", +) +# Flags without a value that the warmup pass drops; it adds --no-test-input back. +_DROPPED_FLAGS = (_NO_TEST_INPUT_FLAG, "--save-result", "--save-detailed") +# benchmark_serving's --seed default. +_BENCHMARK_SERVING_DEFAULT_SEED = 0 +# 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 warmup_pass_enabled(environ: Optional[Mapping[str, str]] = None) -> bool: + """Return whether the warmup pass is on. + + It is on unless ``TRTLLM_PERF_WARMUP_PASS`` is ``0``, ``false``, ``no`` or + ``off`` (any case). + + Args: + environ: Environment to read; ``os.environ`` when None. + """ + env = os.environ if environ is None else environ + return env.get(WARMUP_PASS_ENV, "1").strip().lower() not in _DISABLED_VALUES + + +def warmup_pass_applies(benchmark_mode: Optional[str]) -> bool: + """Return whether a case in this benchmark mode gets a warmup pass. + + Args: + benchmark_mode: The case's benchmark mode; None or "" for a plain + aggregated case. + """ + return (benchmark_mode or "") in WARMUP_PASS_MODES + + +def warmup_pass_plan( + eligible_requests: Sequence[int], measured: Sequence[bool], time_breakdown: bool +) -> List[int]: + """Return each client's warmup pass size for one server, in run order (0: none). + + The time breakdown leaves out only the records written before the server's + first measured client, and pools every later record. So while it is collected, + only that first measured client gets a pass: a later pass's records would count + as measured for every client of the server. Without it, each client keeps its + own pass. + + Args: + eligible_requests: Each client's pass size on its own; 0 if it gets none. + measured: Whether each client is measured; False for accuracy-only clients. + time_breakdown: Whether the time breakdown is collected for the server. + """ + plan: List[int] = [] + after_first_measured = False + for num_requests, is_measured in zip(eligible_requests, measured): + plan.append(0 if time_breakdown and after_first_measured else num_requests) + after_first_measured = after_first_measured or is_measured + return plan + + +def one_wave_requests(concurrency: int, num_requests: int, max_batch_size: int) -> int: + """Return the requests in one wave of a benchmark: those it keeps in flight at once. + + That is the client's concurrency limit. A client without one sends every + request at once, so the server's batch size caps the wave. A wave never holds + more requests than the benchmark sends. + + Args: + concurrency: The client's concurrency limit; 0 or less for none. + num_requests: Requests the benchmark sends; 0 or less if unknown. + max_batch_size: The server's maximum batch size. + """ + wave = concurrency if concurrency > 0 else max_batch_size + return min(wave, num_requests) if num_requests > 0 else wave + + +class WarmupPass(NamedTuple): + """The warmup pass to send before one measured client.""" + + # The build_warmup_pass_cmd command, without host and port. + cmd: List[str] + num_requests: int + + +def build_warmup_pass_cmd(measured_cmd: Sequence[str], num_requests: int) -> List[str]: + """Derive the warmup pass's ``benchmark_serving`` command from the measured one. + + Every flag that shapes the traffic is kept: model, tokenizer, dataset and + sequence lengths, maximum concurrency, backend, streaming, chat template. Only + these change: + + * ``--num-prompts`` becomes ``num_requests``. + * ``--seed`` becomes the measured seed plus one, so the pass draws different + prompts with the same length distribution: other random token ids, or + another shuffle of the same dataset file. + * ``--no-test-input`` is added, so no single test request goes first. + * ``--save-request-time-breakdown``, ``--save-result``, ``--save-detailed``, + ``--result-dir`` and ``--result-filename`` are dropped with their values, + so the pass writes none of the measured client's output files. + + The measured command must pass each flag and its value as separate items, as + the runners build it. The caller adds host and port. + + Args: + measured_cmd: The measured client's ``benchmark_serving`` command. + num_requests: Requests the pass sends; at least 1. + + Returns: + The warmup pass command. + + Raises: + ValueError: ``num_requests`` is below 1. + """ + if num_requests < 1: + raise ValueError(f"A warmup pass needs at least one request, got {num_requests}") + warmup_cmd: List[str] = [] + measured_seed = _BENCHMARK_SERVING_DEFAULT_SEED + args = iter(measured_cmd) + for arg in args: + if arg in _FLAGS_WITH_REPLACED_VALUE: + value = next(args, None) + if arg == _SEED_FLAG and value is not None: + measured_seed = int(value) + continue + if arg in _DROPPED_FLAGS: + continue + warmup_cmd.append(arg) + return warmup_cmd + [ + _NUM_PROMPTS_FLAG, + str(num_requests), + _SEED_FLAG, + str(measured_seed + 1), + _NO_TEST_INPUT_FLAG, + ] + + +def count_failed_requests(output: str) -> int: + """Return the failed requests a ``benchmark_serving`` output reports, or 0.""" + match = _FAILED_REQUESTS_RE.search(output) + return int(match.group(1)) if match else 0 + + +def send_warmup_pass( + warmup_cmd: List[str], + num_requests: int, + run_benchmark: Callable[[List[str]], str], + log_path: str, + timeout: Optional[float] = None, + log_info: Callable[[str], None] = print, + log_warning: Callable[[str], None] = print, +) -> int: + """Send the warmup pass; return how many of its requests succeeded. + + If only some requests fail, the pass still warmed the shapes it reached, so it + logs a warning and the measured client decides whether the test passes. + + Args: + warmup_cmd: The ``build_warmup_pass_cmd`` command, with the server's host + and port. + num_requests: Requests the command sends. + run_benchmark: Runs a command and returns its output. Raises + ``subprocess.CalledProcessError`` if the command exits non-zero, and + ``subprocess.TimeoutExpired`` if it applies a time limit that runs out. + log_path: Where ``run_benchmark`` writes the output; named in messages. + timeout: The time limit ``run_benchmark`` applies, in seconds; None if it + applies none. + log_info: Receives progress messages. + log_warning: Receives the message about a partial failure. + + Returns: + The number of warmup requests that succeeded. + + Raises: + RuntimeError: The command exited non-zero, timed out, or had every + request fail. + """ + log_info(f"Warmup pass: sending {num_requests} requests. cmd is {warmup_cmd}") + start = time.monotonic() + try: + output = run_benchmark(warmup_cmd) + except subprocess.CalledProcessError as e: + raise RuntimeError( + f"Warmup pass: benchmark_serving exited with code {e.returncode}. See {log_path}" + ) from e + except subprocess.TimeoutExpired as e: + limit = f" within {timeout:.0f}s" if timeout is not None else "" + raise RuntimeError( + f"Warmup pass: {num_requests} requests did not finish{limit}. See {log_path}" + ) from e + elapsed = time.monotonic() - start + failed_count = count_failed_requests(output) + if failed_count >= num_requests: + raise RuntimeError(f"Warmup pass: all {num_requests} requests failed. See {log_path}") + if failed_count: + log_warning( + f"Warmup pass: {failed_count} of {num_requests} requests failed ({elapsed:.0f}s). " + f"See {log_path}" + ) + else: + log_info(f"Warmup pass: {num_requests} requests done in {elapsed:.0f}s") + return num_requests - failed_count