diff --git a/docs_new/docs/developer_guide/bench_serving.mdx b/docs_new/docs/developer_guide/bench_serving.mdx index e0dd504f7..ec9e1db6a 100644 --- a/docs_new/docs/developer_guide/bench_serving.mdx +++ b/docs_new/docs/developer_guide/bench_serving.mdx @@ -66,6 +66,7 @@ Select with `--dataset-name`: - `generated-shared-prefix`: synthetic dataset with shared long system prompts and short questions - `mmmu`: samples from MMMU (Math split) and includes images - `speed-bench`: [SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) (**SPEculative Evaluation Dataset**) — a unified benchmark for evaluating [Speculative Decoding (SD)](https://arxiv.org/abs/2604.09557) algorithms. Uses the Throughput split, which provides fixed-length input sequences (1K–32K tokens) grouped into three output-entropy categories (`low_entropy`, `mixed`, `high_entropy`). Requires a pre-downloaded JSONL file passed via `--dataset-path`. +- `agentic-trace`: replays pre-built multi-turn agentic traces (e.g. OpenHands / SWE-smith). Each conversation is replayed round by round, feeding the server's real assistant reply back into the next round's history. Requires a chat backend (`--backend sglang-oai-chat`) and a trace JSON passed via `--dataset-path`. Common dataset flags: @@ -93,6 +94,13 @@ Image dataset flags (for `image`): - `--image-format`: Image format (jpeg or png) - `--image-content`: Image content type (random or blank) +Agentic trace flags (for `agentic-trace`): + +- `--dataset-path`: path to the pre-built trace JSON +- `--sharegpt-output-len`: per-turn output length (default: 220) +- `--dataset-offset`: rotate the conversation list by this many entries before sampling, so successive sweep steps start on fresh conversations +- `--agentic-max-turns`: cap each conversation to at most this many turns (useful for small, fast profiling runs) + SPEED-Bench flags (for `speed-bench`): - `--dataset-path PATH`: path to the pre-downloaded SPEED-Bench Throughput JSONL (e.g., `throughput_1k.jsonl`). Use the [SPEED-Bench measurement framework](https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/specdec_bench) to generate it. diff --git a/python/sglang/benchmark/datasets/__init__.py b/python/sglang/benchmark/datasets/__init__.py index 36a437bfa..edc3a19b4 100644 --- a/python/sglang/benchmark/datasets/__init__.py +++ b/python/sglang/benchmark/datasets/__init__.py @@ -1,5 +1,6 @@ from typing import Dict, Type +from sglang.benchmark.datasets.agentic_trace import AgenticTraceDataset from sglang.benchmark.datasets.autobench import AutoBenchmarkDataset from sglang.benchmark.datasets.common import BaseDataset, DatasetRow from sglang.benchmark.datasets.custom import CustomDataset @@ -16,6 +17,7 @@ from sglang.benchmark.datasets.sharegpt import ShareGPTDataset from sglang.benchmark.datasets.speed_bench import SpeedBenchDataset DATASET_MAPPING: Dict[str, Type[BaseDataset]] = { + "agentic-trace": AgenticTraceDataset, "autobench": AutoBenchmarkDataset, "sharegpt": ShareGPTDataset, "custom": CustomDataset, diff --git a/python/sglang/benchmark/datasets/agentic_trace.py b/python/sglang/benchmark/datasets/agentic_trace.py new file mode 100644 index 000000000..78b11a735 --- /dev/null +++ b/python/sglang/benchmark/datasets/agentic_trace.py @@ -0,0 +1,114 @@ +import json +import os +from argparse import Namespace +from dataclasses import dataclass +from typing import List, Optional + +import numpy as np +from transformers import PreTrainedTokenizerBase + +from sglang.benchmark.datasets.common import BaseDataset, DatasetRow + +# Per-turn output length when --sharegpt-output-len is not given; matches the +# ~220-token average assistant reply of OpenHands-style agentic traces. +DEFAULT_AGENTIC_OUTPUT_LEN = 220 + + +@dataclass +class AgenticTraceDataset(BaseDataset): + """Multi-turn agentic trace loader (e.g. OpenHands / SWE-smith traces). + + Expects a trace JSON of the shape:: + + { + "metadata": {...}, + "conversations": [ + [ # one conversation == a list of turns + {"messages": [{"role": "system", ...}, {"role": "user", ...}], + "prompt_tokens": 73821}, + {"messages": [{"role": "user", ...}], "prompt_tokens": 74894}, + ... + ], + ... + ] + } + + Each turn's ``messages`` holds only the new non-assistant messages for that + turn. One conversation becomes one :class:`DatasetRow` whose ``prompt`` is + the list of per-turn message deltas; ``bench_serving`` detects this shape as + multi-turn and replays each conversation round by round, feeding the + server's real assistant reply back into the next round's history. + + Use with a chat backend (``--backend sglang-oai-chat``). + """ + + dataset_path: str + num_requests: int + fixed_output_len: Optional[int] + offset: int + max_turns: Optional[int] + + @classmethod + def from_args(cls, args: Namespace) -> "AgenticTraceDataset": + return cls( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + fixed_output_len=args.sharegpt_output_len, + offset=args.dataset_offset, + max_turns=args.agentic_max_turns, + ) + + def load( + self, tokenizer: PreTrainedTokenizerBase, model_id=None + ) -> List[DatasetRow]: + if not os.path.isfile(self.dataset_path): + raise FileNotFoundError(f"Dataset not found at {self.dataset_path}") + + with open(self.dataset_path, "r", encoding="utf-8") as f: + data = json.load(f) + + conversations = data.get("conversations", []) + if not conversations: + raise ValueError(f"No 'conversations' found in {self.dataset_path}.") + + offset = self.offset % len(conversations) + if offset: + conversations = conversations[offset:] + conversations[:offset] + + output_len = self.fixed_output_len or DEFAULT_AGENTIC_OUTPUT_LEN + + filtered_dataset: List[DatasetRow] = [] + for conversation in conversations: + if self.num_requests > 0 and len(filtered_dataset) >= self.num_requests: + break + + prompt = [turn["messages"] for turn in conversation if turn.get("messages")] + if self.max_turns: + prompt = prompt[: self.max_turns] + if not prompt: + continue + + # Informational only: multi-turn replay ignores per-row prompt_len. + prompt_len = int(conversation[0].get("prompt_tokens", 0)) + + filtered_dataset.append( + DatasetRow( + prompt=prompt, + prompt_len=prompt_len, + output_len=output_len, + ) + ) + + if not filtered_dataset: + raise ValueError( + f"No usable conversations loaded from {self.dataset_path}." + ) + + num_turns = [len(row.prompt) for row in filtered_dataset] + print( + f"#Conversations: {len(filtered_dataset)} " + f"(offset={offset}, turns/conv min={min(num_turns)} " + f"max={max(num_turns)} avg={np.mean(num_turns):.1f})" + ) + print(f"#Output tokens per turn: {output_len}") + return filtered_dataset diff --git a/python/sglang/benchmark/serving.py b/python/sglang/benchmark/serving.py index 5fdb491de..869c190ef 100644 --- a/python/sglang/benchmark/serving.py +++ b/python/sglang/benchmark/serving.py @@ -2158,6 +2158,7 @@ def cli_main(): type=str, default="sharegpt", choices=[ + "agentic-trace", "autobench", "sharegpt", "custom", @@ -2176,6 +2177,21 @@ def cli_main(): parser.add_argument( "--dataset-path", type=str, default="", help="Path to the dataset." ) + parser.add_argument( + "--dataset-offset", + type=int, + default=0, + help="Rotate the conversation list by this many entries before sampling " + "(agentic-trace dataset), so successive sweep steps start on fresh " + "conversations.", + ) + parser.add_argument( + "--agentic-max-turns", + type=int, + default=None, + help="Cap each conversation to at most this many turns (agentic-trace " + "dataset). Default: use all turns in the trace.", + ) parser.add_argument( "--speed-bench-category", type=str, diff --git a/test/registered/bench_fn/test_benchmark_datasets_api.py b/test/registered/bench_fn/test_benchmark_datasets_api.py index a85310594..0594ae465 100644 --- a/test/registered/bench_fn/test_benchmark_datasets_api.py +++ b/test/registered/bench_fn/test_benchmark_datasets_api.py @@ -19,6 +19,10 @@ from tokenizers.pre_tokenizers import Whitespace from transformers import PreTrainedTokenizerFast from sglang.benchmark.datasets import DATASET_MAPPING, get_dataset +from sglang.benchmark.datasets.agentic_trace import ( + DEFAULT_AGENTIC_OUTPUT_LEN, + AgenticTraceDataset, +) from sglang.benchmark.datasets.common import DatasetRow, gen_mm_prompt from sglang.benchmark.datasets.custom import sample_custom_requests from sglang.benchmark.datasets.generated_shared_prefix import ( @@ -148,6 +152,8 @@ def make_args(**overrides): "mooncake_workload": "conversation", "speed_bench_category": None, "speed_bench_output_len": 512, + "dataset_offset": 0, + "agentic_max_turns": None, } args.update(overrides) return SimpleNamespace(**args) @@ -284,6 +290,37 @@ class TestBenchmarkDatasetsAPI(unittest.TestCase): f.write(json.dumps(row) + "\n") return str(path) + def _write_agentic_trace_json(self): + trace = { + "metadata": {"source": "test"}, + "conversations": [ + [ + { + "messages": [ + {"role": "system", "content": "You are an agent."}, + {"role": "user", "content": "Fix the bug."}, + ], + "prompt_tokens": 100, + }, + { + "messages": [{"role": "user", "content": "Tool output: ok."}], + "prompt_tokens": 200, + }, + {"messages": []}, + ], + [ + { + "messages": [{"role": "user", "content": "Run the tests."}], + "prompt_tokens": 50, + }, + ], + ], + } + path = self.tmpdir_path / "agentic_trace.json" + with open(path, "w") as f: + json.dump(trace, f) + return str(path) + async def _collect_mooncake_rows(self, records): out = [] async for row in get_mooncake_request_over_time( @@ -517,8 +554,69 @@ class TestBenchmarkDatasetsAPI(unittest.TestCase): with self.assertRaises(ValueError): SpeedBenchDataset.from_args(args) + def test_agentic_trace_sampler(self): + dataset_path = self._write_agentic_trace_json() + args = make_args( + dataset_name="agentic-trace", + dataset_path=dataset_path, + num_prompts=10, + ) + dataset = AgenticTraceDataset.from_args(args) + rows = dataset.load(self.tokenizer) + self.assertEqual(len(rows), 2) + self.assertTrue(all(isinstance(row, DatasetRow) for row in rows)) + self.assertTrue( + all(row.output_len == DEFAULT_AGENTIC_OUTPUT_LEN for row in rows) + ) + # Multi-turn shape: prompt is a list of per-turn message lists, with + # the empty third turn of the first conversation dropped. + self.assertEqual(len(rows[0].prompt), 2) + self.assertEqual(len(rows[1].prompt), 1) + self.assertEqual(rows[0].prompt[0][0]["role"], "system") + self.assertEqual(rows[0].prompt_len, 100) + self.assertEqual(rows[1].prompt_len, 50) + + def test_agentic_trace_offset_and_max_turns(self): + dataset_path = self._write_agentic_trace_json() + args = make_args( + dataset_name="agentic-trace", + dataset_path=dataset_path, + num_prompts=10, + sharegpt_output_len=64, + dataset_offset=1, + agentic_max_turns=1, + ) + dataset = AgenticTraceDataset.from_args(args) + rows = dataset.load(self.tokenizer) + self.assertEqual(len(rows), 2) + # offset=1 rotates the second (single-turn) conversation to the front. + self.assertEqual(rows[0].prompt_len, 50) + self.assertTrue(all(len(row.prompt) == 1 for row in rows)) + self.assertTrue(all(row.output_len == 64 for row in rows)) + + def test_agentic_trace_invalid_input_raises(self): + args = make_args( + dataset_name="agentic-trace", + dataset_path=str(self.tmpdir_path / "missing.json"), + num_prompts=1, + ) + with self.assertRaises(FileNotFoundError): + AgenticTraceDataset.from_args(args).load(self.tokenizer) + + empty_path = self.tmpdir_path / "empty_trace.json" + with open(empty_path, "w") as f: + json.dump({"metadata": {}, "conversations": []}, f) + args = make_args( + dataset_name="agentic-trace", + dataset_path=str(empty_path), + num_prompts=1, + ) + with self.assertRaises(ValueError): + AgenticTraceDataset.from_args(args).load(self.tokenizer) + def test_dataset_mapping_and_dispatch(self): expected = { + "agentic-trace", "sharegpt", "custom", "openai", @@ -603,6 +701,15 @@ class TestBenchmarkDatasetsAPI(unittest.TestCase): self.assertEqual(len(speed_bench_rows), 2) self.assertTrue(all(isinstance(row, DatasetRow) for row in speed_bench_rows)) + agentic_args = make_args( + dataset_name="agentic-trace", + dataset_path=self._write_agentic_trace_json(), + num_prompts=2, + ) + agentic_rows = get_dataset(agentic_args, self.tokenizer, model_id="dummy-model") + self.assertEqual(len(agentic_rows), 2) + self.assertTrue(all(isinstance(row, DatasetRow) for row in agentic_rows)) + def test_get_dataset_unknown_dataset(self): args = make_args(dataset_name="not-a-dataset") with self.assertRaises(ValueError):