[bench] Add agentic-trace multi-turn dataset to bench_serving (#29215)

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

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,

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

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,