Support a benchmark dataset in preset creation

`dataset` selects what every benchmark in a preset session measures: the
synthetic `random` prompts shaped by `input_tokens` and `output_tokens`, a
dataset the benchmark tool supports, or a Hugging Face dataset ID. A custom
dataset provides the requests, so the request-shape properties can't be set
with it, and the preset records the measured means instead.

The dataset is part of the contract: it is written to the session
constraints, the agent reports it with the benchmark, the preset records it,
and `dstack preset` shows it. A session that doesn't set it renders exactly
as before.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: Andrey Cheptsov

mkdocs/docs/concepts/presets.md

@@ -136,57 +136,67 @@ Alternatively, pass `--fleet` to `dstack preset create` or `dstack preset apply`
     repo: Qwen/Qwen2.5-7B-Instruct
     ```
 
-### Shared prefix
+### Previous sessions
 
-By default every request is unique, so the cache hit rate is near zero. Set `shared_prefix_tokens` to control how much of each request the serving framework can serve from its prefix cache.
+Set `previous` to a list of preset IDs to give the agent the results of earlier creation sessions. It analyzes what they tried and how it worked, and aims to improve on them instead of rediscovering it.
 
 <div editor-title="preset.dstack.yml">
 
 ```yaml
-input_tokens: 8192
-output_tokens: 1024
-
-# Roughly 90% of prompt tokens can be served from cache
-shared_prefix_tokens: 7360
+previous:
+  - c83375b4
 ```
 
 </div>
 
-The `shared_prefix_tokens` value is the part of `input_tokens` that is identical across requests, such as a system prompt or conversation history, and must be less than `input_tokens`.
+Alternatively, pass `--previous` (repeatable) to `dstack preset create`.
 
 ### Prompt
 
-The `prompt` property is optional. Set it to guide the agent with custom objectives, target metrics, or an experimentation approach. It accepts inline text or a file `path`.
+Set `prompt` to steer what the agent explores: which frameworks or model variants to try, or how deep to go before settling. It accepts inline text or a file `path`. Constraints such as `concurrency` and `max_ttft` can't be changed this way.
 
 <div editor-title="preset.dstack.yml">
 
 ```yaml
 prompt: |
-  Optimize for the lowest TTFT at concurrency 32. Consider FP8 quantization.
+  Profile the engine before each trial and report how far it is from the
+  memory-bandwidth roofline. While that gap is large, prefer patching the
+  serving framework over tuning flags.
 ```
 
 </div>
 
-### Baseline
+### Dataset
 
-By default, the first trial is a baseline: the agent serves the model the way the chosen serving framework recommends, without tuning it for performance. Later trials are optimization attempts. Set `baseline: false` to make every trial an optimization attempt.
+The requests every benchmark measures.
 
-### Previous sessions
+=== "Random"
 
-Set `previous` to a list of preset IDs to give the agent the results of earlier creation sessions. It analyzes what they tried and how it worked, and aims to improve on them instead of rediscovering it.
+    By default, benchmarks use synthetic prompts shaped by `input_tokens` and `output_tokens`. Set `shared_prefix_tokens` to make part of every request identical, such as a system prompt or conversation history, so the serving framework can serve it from its prefix cache. It must be less than `input_tokens`.
 
-<div editor-title="preset.dstack.yml">
+    ```yaml
+    input_tokens: 8192
+    output_tokens: 1024
 
-```yaml
-previous:
-  - c83375b4
-```
+    # Roughly 90% of prompt tokens can be served from cache
+    shared_prefix_tokens: 7360
+    ```
 
-</div>
+=== "Custom"
 
-Alternatively, pass `--previous` (repeatable) to `dstack preset create`.
+    Set `dataset` to benchmark on real text instead: a dataset the benchmark tool supports, or a Hugging Face dataset ID.
+
+    ```yaml
+    dataset: sharegpt
+    ```
+
+    The dataset provides the requests, so `input_tokens`, `output_tokens`, and `shared_prefix_tokens` can't be set with it, and the preset records the measured means. A gated dataset requires `HF_TOKEN` in `env`.
+
+### Baseline
+
+By default, the first trial is a baseline: the agent serves the model the way the chosen serving framework recommends, without tuning it for performance. Later trials are optimization attempts. Set `baseline: false` to make every trial an optimization attempt.
 
-In this case, the baseline trial reproduces the best comparable previous result to confirm it still holds before optimizing further.
+When the session builds on `previous`, the baseline trial reproduces the best comparable previous result instead, to confirm it still holds before optimizing further.
 
 !!! info "Reference"
     The `preset` configuration supports many more options. See the [`.dstack.yml` reference](../reference/dstack.yml/preset.md).
@@ -274,7 +284,6 @@ For command options and agent settings, see the [`dstack preset` CLI reference](
     * Currently, the agent doesn't upload compiled binaries anywhere; patches compile at runtime
     * Doesn't support PD disaggregation (coming soon)
     * Presets are saved locally (a preset registry is coming soon)
-    * Doesn't allow a custom dataset; always uses `random`
     * Doesn't support ranges for `concurrency`
 
     Report bugs and request features on [GitHub](https://github.com/dstackai/dstack/issues), and ask questions on [Discord](https://discord.gg/u8SmfwPpMd).

src/dstack/_internal/cli/models/configurations.py

@@ -18,6 +18,7 @@
 DEFAULT_INPUT_TOKENS = 1024
 DEFAULT_OUTPUT_TOKENS = 1024
 DEFAULT_BASELINE = True
+DEFAULT_DATASET = "random"
 
 
 class PresetModelRepo(CoreModel):
@@ -199,6 +200,17 @@ class PresetConfiguration(
             )
         ),
     ] = None
+    dataset: Annotated[
+        Optional[str],
+        Field(
+            description=(
+                "The benchmark dataset used during preset creation: `random` for synthetic"
+                " prompts shaped by `input_tokens` and `output_tokens`, a benchmark tool's"
+                " dataset name (e.g. `sharegpt`, `spec_bench`), or a Hugging Face dataset ID."
+                " Defaults to `random`"
+            )
+        ),
+    ] = None
     baseline: Annotated[
         Optional[bool],
         Field(
@@ -236,6 +248,38 @@ def effective_output_tokens(self) -> int:
     def effective_baseline(self) -> bool:
         return self.baseline if self.baseline is not None else DEFAULT_BASELINE
 
+    @property
+    def effective_dataset(self) -> str:
+        return self.dataset if self.dataset is not None else DEFAULT_DATASET
+
+    @field_validator("dataset")
+    @classmethod
+    def validate_dataset_name(cls, value: Optional[str]) -> Optional[str]:
+        if value is None:
+            return None
+        # Stripped because the agent reports the dataset it actually loaded, and
+        # the two are compared for equality when the preset is verified.
+        value = value.strip()
+        if not value:
+            raise ValueError("dataset must be a non-empty string")
+        return value
+
+    @model_validator(mode="after")
+    def validate_dataset(self) -> Self:
+        if self.dataset in (None, DEFAULT_DATASET):
+            return self
+        set_fields = [
+            name
+            for name in ("input_tokens", "output_tokens", "shared_prefix_tokens")
+            if getattr(self, name) is not None
+        ]
+        if set_fields:
+            raise ValueError(
+                f"{', '.join(set_fields)} can only be set with the `random` dataset;"
+                " a custom dataset defines its own request shape"
+            )
+        return self
+
     @model_validator(mode="after")
     def validate_shared_prefix_tokens(self) -> Self:
         # The prefix is carved out of the request, so something has to be left
@@ -294,9 +338,10 @@ class PresetConstraints(CoreModel):
     max_ttft: PositiveInt
     trials_num: PositiveInt
     concurrency: PositiveInt
-    input_tokens: PositiveInt
-    output_tokens: PositiveInt
-    shared_prefix_tokens: int = 0
+    input_tokens: Optional[PositiveInt] = None
+    output_tokens: Optional[PositiveInt] = None
+    shared_prefix_tokens: Optional[int] = None
+    dataset: Optional[str] = None
     baseline: bool = False
     fleets: list[str] = Field(min_length=1)
     env: list[str] = []

src/dstack/_internal/cli/models/preset_agent.py

@@ -36,14 +36,16 @@
                 "output_tokens": {"type": "integer", "minimum": 2},
                 "concurrency": {"type": "integer", "minimum": 1},
                 "shared_prefix_tokens": {"type": "integer", "minimum": 0},
+                "dataset": {"type": "string", "minLength": 1},
             },
+            # `shared_prefix_tokens` and `dataset` are not required: one schema
+            # serves both session modes, and each mode knows only its own field.
             "required": [
                 "api",
                 "num_requests",
                 "input_tokens",
                 "output_tokens",
                 "concurrency",
-                "shared_prefix_tokens",
             ],
             "additionalProperties": False,
         },

src/dstack/_internal/cli/models/presets.py

@@ -20,12 +20,16 @@
 class PresetBenchmarkWorkload(CoreModel):
     api: Literal["chat_completions", "completions"]
     num_requests: PositiveInt
+    # With a dataset other than `random`, the measured means rather than the
+    # configured request shape.
     input_tokens: PositiveInt
     output_tokens: Annotated[int, Field(ge=2)]
     concurrency: PositiveInt
-    # Defaulted rather than required: presets saved before this field existed
-    # must still load, and for them the benchmark was fully unique.
-    shared_prefix_tokens: Annotated[int, Field(ge=0)] = 0
+    # Absent for presets saved before the field existed, and with a dataset
+    # other than `random`, where the dataset decides prefix sharing.
+    shared_prefix_tokens: Annotated[Optional[int], Field(ge=0)] = None
+    # Absent means the synthetic `random` dataset.
+    dataset: Optional[str] = None
 
 
 class PresetBenchmarkLatency(CoreModel):

src/dstack/_internal/cli/services/presets/create.py

@@ -15,6 +15,7 @@
 from rich.text import Text
 
 from dstack._internal.cli.models.configurations import (
+    DEFAULT_DATASET,
     PresetConfiguration,
     PresetConstraints,
 )
@@ -589,6 +590,7 @@ async def _create_preset(
         user_prompt=setup.user_prompt,
         baseline=configuration.effective_baseline,
         previous=", ".join(setup.previous) if setup.previous else None,
+        custom_dataset=configuration.effective_dataset != DEFAULT_DATASET,
     )
     if setup.write_constraints:
         if setup.user_prompt:
@@ -880,6 +882,7 @@ def _build_constraints(
     build_name: str,
     allowed_fleets: Sequence[str],
 ) -> str:
+    dataset = configuration.effective_dataset
     constraints = PresetConstraints.model_validate(
         {
             "run_name_prefix": build_name,
@@ -888,16 +891,21 @@ def _build_constraints(
             "max_ttft": configuration.max_ttft,
             "trials_num": configuration.trials,
             "concurrency": configuration.concurrency,
-            "input_tokens": configuration.effective_input_tokens,
-            "output_tokens": configuration.effective_output_tokens,
-            "shared_prefix_tokens": configuration.shared_prefix_tokens or 0,
+            **(
+                {
+                    "input_tokens": configuration.effective_input_tokens,
+                    "output_tokens": configuration.effective_output_tokens,
+                    "shared_prefix_tokens": configuration.shared_prefix_tokens or 0,
+                }
+                if dataset == DEFAULT_DATASET
+                else {"dataset": dataset}
+            ),
             "baseline": configuration.effective_baseline,
             "fleets": list(allowed_fleets),
             "env": list(configuration.env),
         }
     )
-    # All fields are always present; unset optional constraints render as null.
-    return json.dumps(json.loads(constraints.model_dump_json()), indent=2) + "\n"
+    return json.dumps(json.loads(constraints.model_dump_json(exclude_none=True)), indent=2) + "\n"
 
 
 def _save_final_report_copy(

src/dstack/_internal/cli/services/presets/output.py

@@ -192,6 +192,8 @@ def _add_session(table: Table, session: dict[str, Any], *, verbose: bool = False
     constraints = session.get("constraints") or {}
     parts = []
     objective = []
+    if dataset := constraints.get("dataset"):
+        objective.append(f"data={dataset}")
     if constraints.get("input_tokens") and constraints.get("output_tokens"):
         objective.append(
             f"io={_format_token_count(constraints['input_tokens'])}"
@@ -302,12 +304,16 @@ def format_preset_objective(
     verbose: bool = False,
 ) -> str:
     workload = preset.validations[0].benchmark.workload
-    parts = [
-        f"io={_format_token_count(workload.input_tokens)}"
-        f"/{_format_token_count(workload.output_tokens)}",
-    ]
-    share = round(100 * workload.shared_prefix_tokens / workload.input_tokens)
-    parts.append(f"prefix={share}%")
+    parts = []
+    if workload.dataset:
+        parts.append(f"data={workload.dataset}")
+    else:
+        parts.append(
+            f"io={_format_token_count(workload.input_tokens)}"
+            f"/{_format_token_count(workload.output_tokens)}"
+        )
+        share = round(100 * (workload.shared_prefix_tokens or 0) / workload.input_tokens)
+        parts.append(f"prefix={share}%")
     parts.append(f"conc={workload.concurrency}")
     # Absent for presets saved before the creation record was consulted.
     if verbose and min_context_length is not None:

src/dstack/_internal/cli/services/presets/prompt.py

@@ -174,6 +174,7 @@ def get_preset_agent_system_prompt(
     user_prompt: Optional[str] = None,
     baseline: bool = False,
     previous: Optional[str] = None,
+    custom_dataset: bool = False,
 ) -> str:
     text = _SYSTEM_PROMPT_PATH.read_text(encoding="utf-8").strip()
     variables = {
@@ -182,6 +183,8 @@ def get_preset_agent_system_prompt(
         "baseline": "on" if baseline else None,
         # A comma-separated list of the previous session IDs.
         "previous": previous.strip() if previous else None,
+        # Rendered for its presence only; the dataset itself is in constraints.json.
+        "dataset": "on" if custom_dataset else None,
     }
     applied: set[str] = set()
     rendered = _render_branch(_parse_directives(text, variables), variables, applied, dedent=False)

src/dstack/_internal/cli/services/presets/resources/system_prompt.md

@@ -39,10 +39,15 @@ Field semantics:
   this session. It is fixed so that benchmark results are comparable.
 <!--!TODO: support a concurrency sweep, so that a trial is measured at several
 concurrencies instead of one.-->
+<!--?if dataset-->
+- `dataset`: the benchmark dataset for every benchmark in this session; see
+  `## Benchmark`.
+<!--?else-->
 - `input_tokens`, `output_tokens`: the request shape for every benchmark in
   this session. They are fixed for the same reason.
 - `shared_prefix_tokens`: how many of `input_tokens` are identical in every
   request. `0` means every request is fully unique.
+<!--?end-->
 - `baseline`: whether the first trial must be a baseline rather than an
   optimization attempt; see `# Trials (Main Section)`.
 - `fleets`: use these existing `dstack` fleets only. Do not create, delete,
@@ -293,9 +298,31 @@ no trials remain. In that case, log the failure to `final_report.json` (see
 ## Benchmark
 
 During trials, run benchmarks via SSH inside the task, directly against the
-serving engine: use `concurrency`, `input_tokens`, `output_tokens`, and
-`shared_prefix_tokens` from `constraints.json` and measure all trials the same
+serving engine: use <!--?if dataset-->`dataset` and `concurrency`<!--?else-->`concurrency`, `input_tokens`, `output_tokens`, and
+`shared_prefix_tokens`<!--?end--> from `constraints.json` and measure all trials the same
 way so that their results are comparable with each other.
+<!--?if dataset-->
+Before any benchmark, reset the serving engine's prefix cache, or restart the
+engine, so it does not reuse what a previous benchmark cached. Do not vary
+which samples the dataset provides between benchmarks.
+
+Use the `dataset` for every benchmark. Choose the benchmark tool's options
+that load exactly that dataset, and confirm from the tool's own
+documentation, for the version you run, how it loads the dataset. If the
+dataset fails to load, fix the loading; never fall back to another dataset or
+to synthetic prompts. Prefer the dataset's own output lengths; when the tool
+forces an output length instead, use the same value in every benchmark. For
+example, the dataset options are:
+
+| tool | dataset options |
+| --- | --- |
+| `vllm bench serve` | `--dataset-name <dataset>` when `dataset` is the tool's own dataset name, or `--dataset-name hf --dataset-path <dataset>` when it is a Hugging Face dataset ID |
+| `sglang.benchmark.serving` | `--dataset-name <dataset>` when `dataset` is the tool's own dataset name; the tool has no Hugging Face dataset option |
+
+The table is an example and not a full command: the remaining options still
+come from `concurrency`, option names and defaults differ between versions,
+and any other tool needs its own equivalent.
+<!--?else-->
 Before any benchmark, ensure it uses a different seed than the previous
 benchmark. Otherwise the benchmark will depend on what has been cached by the
 previous benchmark.
@@ -315,6 +342,7 @@ lengths. For example, the shared-prefix options are:
 The table is an example and not a full command: the remaining options still come
 from `concurrency` and `output_tokens`, option names and defaults differ between
 versions, and any other tool needs its own equivalent.
+<!--?end-->
 
 Before any benchmark — a trial one or the final one — verify that the model
 works as expected: send real requests and check the responses, including
@@ -331,8 +359,9 @@ trial benchmarks in `trials/<n>/trial.json`, the final benchmark as
 {
   "tool": "vllm bench serve",
   "tool_version": "0.11.0",
-  "command": "vllm bench serve ...",
-  "workload": {"api": "chat_completions", "num_requests": 16, "input_tokens": 1024, "output_tokens": 128, "concurrency": 8, "shared_prefix_tokens": 768},
+  "command": "vllm bench serve ...",<!--?if dataset-->
+  "workload": {"api": "chat_completions", "num_requests": 16, "input_tokens": 1024, "output_tokens": 128, "concurrency": 8, "dataset": "sharegpt"},<!--?else-->
+  "workload": {"api": "chat_completions", "num_requests": 16, "input_tokens": 1024, "output_tokens": 128, "concurrency": 8, "shared_prefix_tokens": 768},<!--?end-->
   "metrics": {
     "successful_requests": 16, "failed_requests": 0, "duration_seconds": 4.0,
     "total_input_tokens": 16384, "total_output_tokens": 2048,
@@ -343,6 +372,11 @@ trial benchmarks in `trials/<n>/trial.json`, the final benchmark as
 }
 ```
 
+<!--?if dataset-->
+Set `workload.dataset` to `dataset` from `constraints.json`, and compute
+`workload.input_tokens` and `workload.output_tokens` as the measured mean
+input and output token counts of the benchmark, rounded to whole tokens.
+<!--?end-->
 Compute `output_tok_per_s` as `total_output_tokens / duration_seconds` and
 `per_user_tok_per_s` as `output_tok_per_s / workload.concurrency`. These are
 the numbers used to compare trials (see `## Performance`).