diff --git a/src/srtctl/benchmarks/scripts/sa-bench/rollup.py b/src/srtctl/benchmarks/scripts/sa-bench/rollup.py index b0271d814..c250351a7 100644 --- a/src/srtctl/benchmarks/scripts/sa-bench/rollup.py +++ b/src/srtctl/benchmarks/scripts/sa-bench/rollup.py @@ -21,6 +21,9 @@ "Config", "Total GPU Count", "Decode GPU Count", + "Total Working GPU Count", + "Decode Working GPU Count", + "Prefill Working GPU Count", "Concurrency", "Total Token Throughput", "Output Token Throughput", @@ -73,7 +76,6 @@ def _as_int(value: Any) -> int: except (TypeError, ValueError): return 0 - def _compute_gpu_counts(resources: dict[str, Any]) -> tuple[int | None, int | None]: """Compute total and decode-serving GPU counts from resource settings.""" gpus_per_node = _as_int(resources.get("gpus_per_node")) @@ -151,6 +153,28 @@ def _extract_p90_decode_running_requests(log_dir: Path, metadata: dict[str, Any] return None +def _compute_working_gpu_counts(resources: dict[str, Any]) -> tuple[int | None, int | None, int | None]: + """Compute prefill, decode, and total working GPU counts from worker-level resource settings.""" + prefill_workers = _as_int(resources.get("prefill_workers")) + gpus_per_prefill = _as_int(resources.get("gpus_per_prefill")) + decode_workers = _as_int(resources.get("decode_workers")) + gpus_per_decode = _as_int(resources.get("gpus_per_decode")) + + prefill_working = prefill_workers * gpus_per_prefill if prefill_workers > 0 and gpus_per_prefill > 0 else None + decode_working = decode_workers * gpus_per_decode if decode_workers > 0 and gpus_per_decode > 0 else None + + if prefill_working is not None and decode_working is not None: + total_working = prefill_working + decode_working + elif prefill_working is not None: + total_working = prefill_working + elif decode_working is not None: + total_working = decode_working + else: + total_working = None + + return prefill_working, decode_working, total_working + + def _safe_ratio(numerator: float | int | None, denominator: float | int | None) -> float | None: """Return numerator / denominator when both values are valid and denominator != 0.""" if numerator is None or denominator in (None, 0): @@ -175,14 +199,24 @@ def _build_csv_row( gpu_num: int | None, decode_gpu_count: int | None, p90_decode_running_requests: int | None, + prefill_working_gpu_count: int | None, + decode_working_gpu_count: int | None, + total_working_gpu_count: int | None, ) -> dict[str, object]: """Build one CSV row from a parsed sa-bench result.""" total_token_throughput = data.get("total_token_throughput") median_tpot = data.get("median_tpot_ms") + # Fall back to node-based GPU counts when worker-level counts are unavailable (e.g. agg mode). + effective_total_working = total_working_gpu_count if total_working_gpu_count is not None else gpu_num + effective_decode_working = decode_working_gpu_count if decode_working_gpu_count is not None else decode_gpu_count + effective_prefill_working = prefill_working_gpu_count if prefill_working_gpu_count is not None else (0 if gpu_num is not None else None) row = { "Config": config_name, "Total GPU Count": gpu_num, "Decode GPU Count": decode_gpu_count, + "Total Working GPU Count": effective_total_working, + "Decode Working GPU Count": effective_decode_working, + "Prefill Working GPU Count": effective_prefill_working, "Concurrency": data.get("max_concurrency"), "Total Token Throughput": total_token_throughput, "Output Token Throughput": data.get("output_throughput"), @@ -191,7 +225,7 @@ def _build_csv_row( "Median ITL": data.get("median_itl_ms"), "P90 Decode Running Requests": p90_decode_running_requests, "Output Token Throughput per User": _safe_ratio(1000.0, median_tpot), - "Total Token Throughput per GPU": _safe_ratio(total_token_throughput, gpu_num), + "Total Token Throughput per GPU": _safe_ratio(total_token_throughput, effective_total_working), } return {key: _format_csv_value(value) for key, value in row.items()} @@ -210,6 +244,9 @@ def main(log_dir: Path) -> None: config_name = metadata.get("job_name") if metadata else None resources = metadata.get("resources") if metadata else None total_gpu_count, decode_gpu_count = _compute_gpu_counts(resources) if resources else (None, None) + prefill_working_gpu_count, decode_working_gpu_count, total_working_gpu_count = ( + _compute_working_gpu_counts(resources) if resources else (None, None, None) + ) p90_decode_running_requests = _extract_p90_decode_running_requests(log_dir, metadata) for result_file in result_files: @@ -249,6 +286,9 @@ def main(log_dir: Path) -> None: gpu_num=total_gpu_count, decode_gpu_count=decode_gpu_count, p90_decode_running_requests=p90_decode_running_requests, + prefill_working_gpu_count=prefill_working_gpu_count, + decode_working_gpu_count=decode_working_gpu_count, + total_working_gpu_count=total_working_gpu_count, ) ) diff --git a/tests/test_sa_bench_rollup.py b/tests/test_sa_bench_rollup.py index 0cf9abef1..52ab05044 100644 --- a/tests/test_sa_bench_rollup.py +++ b/tests/test_sa_bench_rollup.py @@ -99,6 +99,9 @@ def test_sa_bench_rollup_generates_json_and_csv_without_metadata(tmp_path): assert first["Config"] == "GLM-5-FP8" assert first["Total GPU Count"] == "" assert first["Decode GPU Count"] == "" + assert first["Total Working GPU Count"] == "" + assert first["Decode Working GPU Count"] == "" + assert first["Prefill Working GPU Count"] == "" assert first["Total Token Throughput"] == "9876" assert first["Output Token Throughput"] == "1234.5" assert first["Median TTFT"] == "95" @@ -125,8 +128,11 @@ def test_sa_bench_rollup_uses_metadata_for_name_gpu_counts_and_p90(tmp_path): "resources": { "gpus_per_node": 8, "prefill_nodes": 1, + "prefill_workers": 1, + "gpus_per_prefill": 8, "decode_nodes": 2, "decode_workers": 2, + "gpus_per_decode": 8, "agg_workers": 0, }, } @@ -164,6 +170,9 @@ def test_sa_bench_rollup_uses_metadata_for_name_gpu_counts_and_p90(tmp_path): assert row["Config"] == "job-metadata-name" assert row["Total GPU Count"] == "24" assert row["Decode GPU Count"] == "16" + assert row["Total Working GPU Count"] == "24" + assert row["Decode Working GPU Count"] == "16" + assert row["Prefill Working GPU Count"] == "8" assert row["P90 Decode Running Requests"] == "32" assert row["Output Token Throughput per User"] == "40" assert row["Total Token Throughput per GPU"] == "33.333" @@ -215,6 +224,9 @@ def test_sa_bench_rollup_aggregated_deployment_reports_all_gpus(tmp_path): assert row["Config"] == "glm47flash-agg-tp4-baseline" assert row["Total GPU Count"] == "4" assert row["Decode GPU Count"] == "4" + assert row["Total Working GPU Count"] == "4" + assert row["Decode Working GPU Count"] == "4" + assert row["Prefill Working GPU Count"] == "0" assert row["Total Token Throughput per GPU"] == "6307.5" @@ -266,6 +278,62 @@ def test_sa_bench_rollup_tolerates_null_resource_fields(tmp_path): assert row["Total GPU Count"] == "4" +def test_sa_bench_rollup_sub_node_prefill_working_gpu_counts(tmp_path): + """When gpus_per_prefill < gpus_per_node, working GPU counts differ from node-based counts.""" + rollup = _load_rollup_module() + + logs_dir = tmp_path / "logs" + result_dir = logs_dir / "sa-bench_isl_4096_osl_512" + result_dir.mkdir(parents=True) + + # 1 prefill node (8 GPUs) but only 4 GPUs used per prefill worker → 1 worker × 4 = 4 prefill working GPUs + # 2 decode nodes, 2 workers × 8 GPUs = 16 decode working GPUs + # Total GPU Count = (1 + 2) × 8 = 24 (node-based) + # Total Working GPU Count = 4 + 16 = 20 (worker-based, diverges from node count) + (tmp_path / "5001.json").write_text( + json.dumps( + { + "job_name": "sub-node-prefill-job", + "backend_type": "sglang", + "resources": { + "gpus_per_node": 8, + "prefill_nodes": 1, + "prefill_workers": 1, + "gpus_per_prefill": 4, + "decode_nodes": 2, + "decode_workers": 2, + "gpus_per_decode": 8, + "agg_workers": 0, + }, + } + ) + ) + + result = { + "model_id": "test-model", + "max_concurrency": 64, + "output_throughput": 500.0, + "total_token_throughput": 1000.0, + "median_ttft_ms": 50.0, + "median_tpot_ms": 20.0, + "median_itl_ms": 10.0, + } + (result_dir / "results_concurrency_64_gpus_24_ctx_4_gen_16.json").write_text(json.dumps(result)) + + rollup.main(logs_dir) + + rows = _read_csv_rows(logs_dir / "benchmark-rollup.csv") + assert len(rows) == 1 + row = rows[0] + assert row["Total GPU Count"] == "24" + assert row["Decode GPU Count"] == "16" + assert row["Prefill Working GPU Count"] == "4" + assert row["Decode Working GPU Count"] == "16" + assert row["Total Working GPU Count"] == "20" + # 1000 / 20 = 50 (uses working count, not node count) + assert row["Total Token Throughput per GPU"] == "50" + + def test_compute_gpu_counts_handles_none_values(): """_compute_gpu_counts should never raise on null-valued keys.""" rollup = _load_rollup_module()