refactor(daser): remove dead code, merge duplication, split large files

benchmarks/README.md

@@ -148,6 +148,37 @@ interleaves samples by dataset, reducing queue-order bias. `--max-samples`
 limits each selected JSONL file independently. LongBench is intended for
 long-text performance validation rather than fast correctness smoke tests.
 
+### Adding a dataset
+
+Datasets live behind a registry in `utils/datasets.py`, so the runner and load
+generator never branch on a dataset name. To add one, subclass
+`BenchmarkDataset`, set `name`, implement `load()` (return `BenchmarkSample`
+objects), declare your own CLI options in `add_cli_args`, build yourself in
+`from_args`, and decorate the class with `@register_dataset`:
+
+```python
+@register_dataset
+class MyDataset(BenchmarkDataset):
+    name = "mydata"
+
+    @classmethod
+    def add_cli_args(cls, parser):
+        parser.add_argument("--mydata-path")
+
+    @classmethod
+    def from_args(cls, args):
+        return cls(args.mydata_path, max_samples=args.max_samples)
+
+    def load(self):
+        return [BenchmarkSample(...), ...]
+```
+
+`--dataset mydata` then works in both `run_bench.py` and `bench_load.py` with no
+edits to either: `add_dataset_cli_args` registers the options and
+`build_dataset` dispatches by name. Prompt construction, sizing, load
+generation, and metrics already consume only the normalized `BenchmarkSample`
+shape.
+
 ## Prompt Construction
 
 All service benchmarks use the same RAG chat prompt shape:

benchmarks/bench_load.py

@@ -20,9 +20,8 @@
     slot_size_for_block_tokens,
 )
 from benchmarks.utils.datasets import (
-    BenchmarkSample,
-    ImdbDataset,
-    LongBenchDataset,
+    add_dataset_cli_args,
+    build_dataset,
     dedup_by_context,
     interleave_samples,
 )
@@ -65,10 +64,7 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
     )
     parser.add_argument("--manifest", default=None)
     parser.add_argument("--prepared-config", default=None)
-    parser.add_argument("--dataset", choices=("imdb", "longbench"), required=True)
-    parser.add_argument("--imdb")
-    parser.add_argument("--longbench-dir")
-    parser.add_argument("--datasets", default=None)
+    add_dataset_cli_args(parser)
     parser.add_argument("--max-samples", type=int, default=20)
     parser.add_argument("--block-size", type=int, default=BLOCK_TOKENS)
     parser.add_argument("--max-context-tokens", type=int, default=0)
@@ -104,7 +100,7 @@ async def main_async(args: argparse.Namespace) -> None:
         raise ValueError("--model is required with --prepare-only")
     if store_dir is None:
         raise ValueError("--store-dir is required with --prepare-only")
-    samples = _load_samples(args)
+    samples = build_dataset(args).load()
 
     from transformers import AutoConfig, AutoTokenizer
 
@@ -260,21 +256,6 @@ async def main_async(args: argparse.Namespace) -> None:
     print(f"results={args.out}")
 
 
-def _load_samples(args: argparse.Namespace) -> list[BenchmarkSample]:
-    if args.dataset == "imdb":
-        if not args.imdb:
-            raise ValueError("--imdb is required for --dataset imdb")
-        return ImdbDataset(args.imdb, max_samples=args.max_samples).load()
-    if not args.longbench_dir:
-        raise ValueError("--longbench-dir is required for --dataset longbench")
-    datasets = None
-    if args.datasets:
-        datasets = [item.strip() for item in args.datasets.split(",") if item.strip()]
-    return LongBenchDataset(
-        args.longbench_dir, datasets=datasets, max_samples=args.max_samples
-    ).load()
-
-
 def _should_dedup_context(args: argparse.Namespace) -> bool:
     """Return whether duplicate contexts should be removed from the workload.