|
| 1 | +diff --git a/src/srtctl/cli/mixins/benchmark_stage.py b/src/srtctl/cli/mixins/benchmark_stage.py |
| 2 | +index fe6cbf1..f282fe6 100644 |
| 3 | +--- a/src/srtctl/cli/mixins/benchmark_stage.py |
| 4 | ++++ b/src/srtctl/cli/mixins/benchmark_stage.py |
| 5 | +@@ -46,12 +46,32 @@ def _vllm_data_parallel_size(config: "SrtConfig", mode: str) -> int: |
| 6 | + return int(mode_config.get("data-parallel-size") or mode_config.get("data_parallel_size") or 1) |
| 7 | + |
| 8 | + |
| 9 | +-def _get_health_expectations(config: "SrtConfig") -> tuple[int, int, str, int]: |
| 10 | ++def _vllm_health_entries( |
| 11 | ++ config: "SrtConfig", |
| 12 | ++ mode: str, |
| 13 | ++ logical_workers: int, |
| 14 | ++ backend_processes: list["Process"] | None, |
| 15 | ++) -> int: |
| 16 | ++ """Return expected Dynamo generate registrations for a vLLM worker mode.""" |
| 17 | ++ dp_size = _vllm_data_parallel_size(config, mode) |
| 18 | ++ if dp_size > 1 and getattr(config.backend, "dp_launch_mode", "per_gpu") == "per_node": |
| 19 | ++ if backend_processes is None: |
| 20 | ++ raise ValueError("backend_processes are required for per-node DP health expectations") |
| 21 | ++ endpoint_mode = "agg" if mode == "aggregated" else mode |
| 22 | ++ return sum(process.endpoint_mode == endpoint_mode for process in backend_processes) |
| 23 | ++ |
| 24 | ++ return logical_workers * dp_size |
| 25 | ++ |
| 26 | ++ |
| 27 | ++def _get_health_expectations( |
| 28 | ++ config: "SrtConfig", backend_processes: list["Process"] | None = None |
| 29 | ++) -> tuple[int, int, str, int]: |
| 30 | + """Compute expected health counts in the units reported by the frontend. |
| 31 | + |
| 32 | + Dynamo's /health endpoint reports registered generate instances. For vLLM |
| 33 | +- DP workers, that means one entry per DP rank, not one entry per logical |
| 34 | +- srt-slurm worker. Other frontends keep using logical worker counts. |
| 35 | ++ DP workers, per-GPU launch registers one entry per DP rank, while per-node |
| 36 | ++ launch registers one entry per node-local process. Other frontends keep |
| 37 | ++ using logical worker counts. |
| 38 | + """ |
| 39 | + r = config.resources |
| 40 | + |
| 41 | +@@ -67,10 +87,10 @@ def _get_health_expectations(config: "SrtConfig") -> tuple[int, int, str, int]: |
| 42 | + if config.frontend.type == "dynamo" and getattr(config.backend, "type", None) == "vllm": |
| 43 | + if r.num_agg > 0: |
| 44 | + n_prefill = 0 |
| 45 | +- n_decode = logical_decode * _vllm_data_parallel_size(config, "aggregated") |
| 46 | ++ n_decode = _vllm_health_entries(config, "aggregated", logical_decode, backend_processes) |
| 47 | + else: |
| 48 | +- n_prefill = logical_prefill * _vllm_data_parallel_size(config, "prefill") |
| 49 | +- n_decode = logical_decode * _vllm_data_parallel_size(config, "decode") |
| 50 | ++ n_prefill = _vllm_health_entries(config, "prefill", logical_prefill, backend_processes) |
| 51 | ++ n_decode = _vllm_health_entries(config, "decode", logical_decode, backend_processes) |
| 52 | + |
| 53 | + count_desc = f"{n_prefill}P + {n_decode}D Dynamo generate instances; logical workers: {worker_desc}" |
| 54 | + return n_prefill, n_decode, count_desc, n_prefill + n_decode |
| 55 | +@@ -131,7 +151,7 @@ class BenchmarkStageMixin: |
| 56 | + """Run the benchmark.""" |
| 57 | + logger.info("Waiting for workers to be ready...") |
| 58 | + |
| 59 | +- n_prefill, n_decode, count_desc, num_workers = _get_health_expectations(self.config) |
| 60 | ++ n_prefill, n_decode, count_desc, num_workers = _get_health_expectations(self.config, self.backend_processes) |
| 61 | + logger.info("Waiting for server health (expecting %d health entries: %s)...", num_workers, count_desc) |
| 62 | + |
| 63 | + hc = self.config.health_check |
0 commit comments