feat(benchmark): expose logical worker endpoints

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.

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

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

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."""