chore: import upstream snapshot with attribution

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