From de78dfa4d292d0b3a4ad55867baacfb8b074e7fd Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Mon, 20 Jul 2026 19:13:01 -0500 Subject: [PATCH] feat(benchmark): expose logical worker endpoints Signed-off-by: Cam Quilici --- docs/config-reference.md | 40 +++++++ src/srtctl/benchmarks/custom.py | 5 + src/srtctl/cli/mixins/benchmark_stage.py | 121 +++++++++++++++----- tests/test_benchmarks.py | 135 +++++++++++++++++++++++ 4 files changed, 272 insertions(+), 29 deletions(-) diff --git a/docs/config-reference.md b/docs/config-reference.md index 2b89a5242..2d751cf4c 100644 --- a/docs/config-reference.md +++ b/docs/config-reference.md @@ -528,6 +528,7 @@ benchmark: | Type | Description | | ----------------- | ---------------------------------------------- | | `manual` | No benchmark (default), manual testing mode | +| `custom` | Arbitrary command with runtime endpoint metadata | | `sa-bench` | Throughput/latency serving benchmark | | `sglang-bench` | SGLang bench_serving benchmark | | `mmlu` | MMLU accuracy evaluation | @@ -545,6 +546,45 @@ benchmark: type: "manual" ``` +### custom + +Run an arbitrary command with `bash -lc`. The command is passed verbatim; srt-slurm does not +expand `{placeholder}` expressions. Use environment variables for runtime-discovered values: + +```yaml +benchmark: + type: custom + command: >- + ./run-benchmark.sh "$SRT_FRONTEND_HOST:$SRT_FRONTEND_PORT" + env: + MY_BENCHMARK_OPTION: "value" +``` + +Every custom benchmark command receives frontend metadata plus mode-specific metadata for each +logical worker leader: + +| Variable | Format | Description | +| ------------------------------- | ------------------------------ | ----------- | +| `SRT_FRONTEND_HOST` | IP | Frontend/orchestrator IP | +| `SRT_FRONTEND_PORT` | port | Frontend public port | +| `SRT_PREFILL_IPS` | comma-separated IPs | Prefill worker leader IPs | +| `SRT_PREFILL_ENDPOINTS` | comma-separated `IP:port` | Prefill worker endpoints | +| `SRT_DECODE_IPS` | comma-separated IPs | Decode worker leader IPs | +| `SRT_DECODE_ENDPOINTS` | comma-separated `IP:port` | Decode worker endpoints | +| `SRT_AGG_IPS` | comma-separated IPs | Aggregated worker leader IPs | +| `SRT_AGG_ENDPOINTS` | comma-separated `IP:port` | Aggregated worker endpoints | +| `AIPERF_SERVER_METRICS_URLS` | comma-separated HTTP URLs | AIPerf-compatible `/metrics` URLs for all logical workers | + +Only variables for modes present in the recipe are emitted. Entries follow logical topology order +(prefill index, decode index, or aggregated index). Multi-node follower ranks are excluded because +they do not own separate engines; co-located logical workers retain repeated IPs and distinct ports +so list positions remain aligned. With a Dynamo frontend, endpoint and metrics URLs use each +leader's `DYN_SYSTEM_PORT`; other frontends use the worker HTTP port. If KVBM metrics are configured, +their URLs are appended to `AIPERF_SERVER_METRICS_URLS` after the logical worker URLs. + +Values in `benchmark.env` are applied last and can explicitly override any automatically injected +variable. + ### sa-bench (Serving Accuracy) Throughput and latency benchmark at various concurrency levels. diff --git a/src/srtctl/benchmarks/custom.py b/src/srtctl/benchmarks/custom.py index 8dcd0bd29..8c9d6d337 100644 --- a/src/srtctl/benchmarks/custom.py +++ b/src/srtctl/benchmarks/custom.py @@ -36,6 +36,11 @@ class CustomBenchmarkRunner(BenchmarkRunner): * If you need to parameterize the command, render it yourself when you generate the recipe and paste the final string into ``benchmark.command``. + * Runtime-discovered frontend and logical worker endpoints are injected + through ``SRT_*`` environment variables. Custom AIPerf commands also + receive ``AIPERF_SERVER_METRICS_URLS``. Multi-node follower ranks are + intentionally excluded; see ``docs/config-reference.md`` for the full + contract. """ @property diff --git a/src/srtctl/cli/mixins/benchmark_stage.py b/src/srtctl/cli/mixins/benchmark_stage.py index b95e3e5f2..1a36fba24 100644 --- a/src/srtctl/cli/mixins/benchmark_stage.py +++ b/src/srtctl/cli/mixins/benchmark_stage.py @@ -145,6 +145,44 @@ def _benchmark_node(self) -> str: self.backend_processes, placement, self.runtime.nodes.head, kind="benchmark.client_placement" ) + def _logical_worker_endpoints(self) -> list[tuple[str, str, int]]: + """Return ``(mode, IP, port)`` for every logical worker leader. + + ``backend_processes`` contains one process per physical node for + multi-node workers. Only rank zero owns the logical worker endpoint, + so follower ranks must not be advertised to benchmark clients. + + Dynamo exposes worker metrics on each leader's system port. Other + frontends expose them on the worker HTTP port, matching the endpoint + selection already used by the profiling integration. + """ + use_sys_port = self.config.frontend.type == "dynamo" + endpoints: list[tuple[str, str, int]] = [] + for process in self.backend_processes: + if not process.is_leader: + continue + port = process.sys_port if use_sys_port else process.http_port + if port <= 0: + continue + host = get_hostname_ip(process.node, self.runtime.network_interface) + endpoints.append((process.endpoint_mode, host, port)) + return endpoints + + @staticmethod + def _get_worker_endpoint_env(endpoints: list[tuple[str, str, int]]) -> dict[str, str]: + """Build mode-specific benchmark environment from logical endpoints.""" + env: dict[str, str] = {} + prefixes = {"prefill": "PREFILL", "decode": "DECODE", "agg": "AGG"} + for mode, prefix in prefixes.items(): + mode_endpoints = [(host, port) for endpoint_mode, host, port in endpoints if endpoint_mode == mode] + if not mode_endpoints: + continue + # Keep one IP per logical endpoint, including repeated IPs for + # co-located workers, so IP and endpoint positions stay aligned. + env[f"SRT_{prefix}_IPS"] = ",".join(host for host, _ in mode_endpoints) + env[f"SRT_{prefix}_ENDPOINTS"] = ",".join(f"{host}:{port}" for host, port in mode_endpoints) + return env + def run_benchmark( self, registry: "ProcessRegistry", stop_event: threading.Event, reporter: StatusReporter | None = None ) -> int: @@ -293,7 +331,11 @@ def _run_benchmark_script( if snapshotter is not None: snapshotter.stop() - def _get_benchmark_profiling_env(self, runner: "BenchmarkRunner") -> dict[str, str]: + def _get_benchmark_profiling_env( + self, + runner: "BenchmarkRunner", + logical_endpoints: list[tuple[str, str, int]] | None = None, + ) -> dict[str, str]: """Get environment variables for the benchmark script.""" env: dict[str, str] = {} @@ -337,20 +379,17 @@ def _get_benchmark_profiling_env(self, runner: "BenchmarkRunner") -> dict[str, s decode_endpoints = [] agg_endpoints = [] - use_sys_port = self.config.frontend.type == "dynamo" - for process in self.backend_processes: - if not process.is_leader: - continue - leader_ip = get_hostname_ip(process.node, self.runtime.network_interface) - port = process.sys_port if use_sys_port else process.http_port + if logical_endpoints is None: + logical_endpoints = self._logical_worker_endpoints() + for mode, leader_ip, port in logical_endpoints: leader_endpoint = f"{leader_ip}:{port}" - if process.endpoint_mode == "prefill": + if mode == "prefill": prefill_ips.append(leader_ip) prefill_endpoints.append(leader_endpoint) - elif process.endpoint_mode == "decode": + elif mode == "decode": decode_ips.append(leader_ip) decode_endpoints.append(leader_endpoint) - elif process.endpoint_mode == "agg": + elif mode == "agg": agg_ips.append(leader_ip) agg_endpoints.append(leader_endpoint) @@ -414,26 +453,37 @@ def _get_sa_bench_slow_down_env(self) -> dict[str, str]: "SA_BENCH_SLOW_DOWN_WAIT_TIME": str(b.slow_down_wait_time), } - def _get_aiperf_server_metrics_env(self) -> dict[str, str]: + def _get_aiperf_server_metrics_env( + self, + logical_endpoints: list[tuple[str, str, int]] | None = None, + *, + logical_workers_only: bool = False, + ) -> dict[str, str]: """Build server metrics URLs for AIPerf benchmarks. - Collects metrics endpoints from all backend processes that expose - a sys_port (vLLM workers with AIPerf metrics enabled), plus KVBM - metrics endpoints if DYN_KVBM_METRICS_PORT is configured. + Built-in AIPerf runners retain their existing physical-process metrics + behavior, which is required by vLLM data-parallel layouts. Custom + benchmarks use logical worker leaders so distributed SGLang follower + ranks are not advertised as separate engines. """ urls: list[str] = [] - if self.config.frontend.type == "vllm": + if logical_workers_only: + if logical_endpoints is None: + logical_endpoints = self._logical_worker_endpoints() + urls = [f"http://{host}:{port}/metrics" for _, host, port in logical_endpoints] + else: + if self.config.frontend.type == "vllm": + for process in self.backend_processes: + if process.endpoint_mode == "agg" and process.is_leader: + host = get_hostname_ip(process.node, self.runtime.network_interface) + urls.append(f"http://{host}:{FRONTEND_PUBLIC_PORT}/metrics") + if urls: + return {"AIPERF_SERVER_METRICS_URLS": ",".join(sorted(set(urls)))} + for process in self.backend_processes: - if process.endpoint_mode == "agg" and process.is_leader: + if process.sys_port > 0: host = get_hostname_ip(process.node, self.runtime.network_interface) - urls.append(f"http://{host}:{FRONTEND_PUBLIC_PORT}/metrics") - if urls: - return {"AIPERF_SERVER_METRICS_URLS": ",".join(sorted(set(urls)))} - - for process in self.backend_processes: - if process.sys_port > 0: - host = get_hostname_ip(process.node, self.runtime.network_interface) - urls.append(f"http://{host}:{process.sys_port}/metrics") + urls.append(f"http://{host}:{process.sys_port}/metrics") # Add KVBM metrics endpoints for prefill processes with DYN_KVBM_METRICS_PORT prefill_env = getattr(self.config.backend, "prefill_environment", {}) @@ -447,13 +497,21 @@ def _get_aiperf_server_metrics_env(self) -> dict[str, str]: if not urls: return {} - return {"AIPERF_SERVER_METRICS_URLS": ",".join(sorted(set(urls)))} + # Custom commands preserve logical topology order; built-in AIPerf + # runners retain their historical sorted physical-process list. + urls = list(dict.fromkeys(urls)) if logical_workers_only else sorted(set(urls)) + return {"AIPERF_SERVER_METRICS_URLS": ",".join(urls)} def _get_benchmark_env(self, runner: "BenchmarkRunner") -> dict[str, str]: """Get environment variables for the benchmark script.""" from srtctl.benchmarks.base import AIPerfBenchmarkRunner - env = self._get_benchmark_profiling_env(runner) + is_custom = self.config.benchmark.type == "custom" + logical_endpoints = self._logical_worker_endpoints() if self.config.profiling.enabled or is_custom else None + env = self._get_benchmark_profiling_env(runner, logical_endpoints) + if is_custom: + assert logical_endpoints is not None + env.update(self._get_worker_endpoint_env(logical_endpoints)) env["SRTCTL_FRONTEND_TYPE"] = self.config.frontend.type # Orchestrator endpoint for the benchmark command. When the client runs on @@ -473,10 +531,15 @@ def _get_benchmark_env(self, runner: "BenchmarkRunner") -> dict[str, str]: if runner.name == "SA-Bench": env.update(self._get_sa_bench_slow_down_env()) - # Add AIPerf-specific env vars for AIPerf-driven benchmarks only + # Built-in AIPerf runners retain physical-process metrics for vLLM DP. + # Custom commands commonly wrap AIPerf but do not inherit from its base + # class, so give them the logical-worker view needed by SGLang TP. if isinstance(runner, AIPerfBenchmarkRunner): env.update(self._get_aiperf_server_metrics_env()) - if self.config.benchmark.aiperf_package: - env["AIPERF_PACKAGE"] = self.config.benchmark.aiperf_package + elif is_custom: + assert logical_endpoints is not None + env.update(self._get_aiperf_server_metrics_env(logical_endpoints, logical_workers_only=True)) + if isinstance(runner, AIPerfBenchmarkRunner) and self.config.benchmark.aiperf_package: + env["AIPERF_PACKAGE"] = self.config.benchmark.aiperf_package return env diff --git a/tests/test_benchmarks.py b/tests/test_benchmarks.py index 098b2779f..386168331 100644 --- a/tests/test_benchmarks.py +++ b/tests/test_benchmarks.py @@ -227,6 +227,42 @@ def test_http_connection_reuse_schema_default_and_roundtrip(self): class TestCustomBenchmarkRunner: """Test custom benchmark runner.""" + @staticmethod + def _benchmark_stage( + frontend_type, + processes, + *, + benchmark_type="custom", + prefill_environment=None, + aggregated_environment=None, + ): + from types import SimpleNamespace + + from srtctl.cli.mixins.benchmark_stage import BenchmarkStageMixin + + class Stage(BenchmarkStageMixin): + @property + def backend_processes(self): + return processes + + stage = Stage() + stage.config = SimpleNamespace( + benchmark=SimpleNamespace(type=benchmark_type, aiperf_package=None), + backend=SimpleNamespace( + prefill_environment=prefill_environment or {}, + aggregated_environment=aggregated_environment or {}, + ), + frontend=SimpleNamespace(type=frontend_type), + profiling=SimpleNamespace(enabled=False), + ) + stage.runtime = SimpleNamespace( + environment={}, + frontend_port=8000, + network_interface="ibp1s0", + nodes=SimpleNamespace(head="head-node"), + ) + return stage + def test_validate_config_requires_command(self): from srtctl.benchmarks.custom import CustomBenchmarkRunner from srtctl.core.schema import BenchmarkConfig, ModelConfig, ResourceConfig, SrtConfig @@ -267,6 +303,105 @@ def test_build_command_uses_custom_container_and_env(self): assert runner.get_container_image(config, runtime) == "nvcr.io/nvidia/python:3.11" assert runner.get_environment(config, runtime) == {"FOO": "bar"} + def test_disaggregated_worker_endpoints_use_logical_leaders(self): + from unittest.mock import patch + + from srtctl.benchmarks.custom import CustomBenchmarkRunner + from srtctl.core.topology import Process + + processes = [ + Process("node-a", frozenset(range(4)), 7500, 6100, "prefill", 0, node_rank=0), + Process("node-b", frozenset(range(4)), 7501, 0, "prefill", 0, node_rank=1), + Process("node-c", frozenset(range(4)), 7502, 6100, "prefill", 1, node_rank=0), + Process("node-d", frozenset(range(4)), 7503, 0, "prefill", 1, node_rank=1), + Process("node-e", frozenset(range(4)), 7504, 6100, "decode", 0, node_rank=0), + Process("node-f", frozenset(range(4)), 7505, 0, "decode", 0, node_rank=1), + ] + stage = self._benchmark_stage("dynamo", processes) + + with patch( + "srtctl.cli.mixins.benchmark_stage.get_hostname_ip", + side_effect=lambda node, interface: f"ip-{node}", + ): + env = stage._get_benchmark_env(CustomBenchmarkRunner()) + + assert env["SRT_PREFILL_IPS"] == "ip-node-a,ip-node-c" + assert env["SRT_PREFILL_ENDPOINTS"] == "ip-node-a:7500,ip-node-c:7502" + assert env["SRT_DECODE_IPS"] == "ip-node-e" + assert env["SRT_DECODE_ENDPOINTS"] == "ip-node-e:7504" + assert "SRT_AGG_IPS" not in env + assert env["AIPERF_SERVER_METRICS_URLS"] == ( + "http://ip-node-a:7500/metrics,http://ip-node-c:7502/metrics,http://ip-node-e:7504/metrics" + ) + + def test_aggregated_worker_endpoint_uses_http_port_without_dynamo(self): + from unittest.mock import patch + + from srtctl.benchmarks.custom import CustomBenchmarkRunner + from srtctl.core.topology import Process + + processes = [ + Process("node-a", frozenset(range(4)), 7500, 6100, "agg", 0, node_rank=0), + Process("node-b", frozenset(range(4)), 7501, 0, "agg", 0, node_rank=1), + ] + stage = self._benchmark_stage("sglang", processes) + + with patch( + "srtctl.cli.mixins.benchmark_stage.get_hostname_ip", + side_effect=lambda node, interface: f"ip-{node}", + ): + env = stage._get_benchmark_env(CustomBenchmarkRunner()) + + assert env["SRT_AGG_IPS"] == "ip-node-a" + assert env["SRT_AGG_ENDPOINTS"] == "ip-node-a:6100" + assert "SRT_PREFILL_ENDPOINTS" not in env + assert "SRT_DECODE_ENDPOINTS" not in env + assert env["AIPERF_SERVER_METRICS_URLS"] == "http://ip-node-a:6100/metrics" + + def test_worker_endpoint_order_keeps_colocated_logical_workers_aligned(self): + from unittest.mock import patch + + from srtctl.benchmarks.custom import CustomBenchmarkRunner + from srtctl.core.topology import Process + + processes = [ + Process("node-a", frozenset({0, 1}), 7500, 6100, "decode", 0), + Process("node-a", frozenset({2, 3}), 7501, 6132, "decode", 1), + ] + stage = self._benchmark_stage("dynamo", processes) + + with patch( + "srtctl.cli.mixins.benchmark_stage.get_hostname_ip", + side_effect=lambda node, interface: "10.0.0.1", + ): + env = stage._get_benchmark_env(CustomBenchmarkRunner()) + + assert env["SRT_DECODE_IPS"] == "10.0.0.1,10.0.0.1" + assert env["SRT_DECODE_ENDPOINTS"] == "10.0.0.1:7500,10.0.0.1:7501" + + def test_builtin_aiperf_retains_physical_process_metrics(self): + from unittest.mock import patch + + from srtctl.benchmarks.trace_replay import TraceReplayRunner + from srtctl.core.topology import Process + + processes = [ + Process("node-a", frozenset(range(4)), 7500, 6100, "prefill", 0, node_rank=0), + Process("node-b", frozenset(range(4)), 7501, 0, "prefill", 0, node_rank=1), + Process("node-c", frozenset(range(4)), 7502, 6100, "decode", 0, node_rank=0), + ] + stage = self._benchmark_stage("dynamo", processes, benchmark_type="trace-replay") + + with patch( + "srtctl.cli.mixins.benchmark_stage.get_hostname_ip", + side_effect=lambda node, interface: f"ip-{node}", + ): + env = stage._get_benchmark_env(TraceReplayRunner()) + + assert env["AIPERF_SERVER_METRICS_URLS"] == ( + "http://ip-node-a:7500/metrics,http://ip-node-b:7501/metrics,http://ip-node-c:7502/metrics" + ) + class TestSGLangBenchRunner: """Test SGLang-Bench runner."""