[Presets] Benchmark dataset support

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.
 
-In this case, the baseline trial reproduces the best comparable previous result to confirm it still holds before optimizing further.
+    ```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.
+
+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).
@@ -268,16 +278,31 @@ $ dstack preset delete c83375b4
 
 </div>
 
-For command options and agent settings, see the [`dstack preset` CLI reference](../reference/cli/dstack/preset.md).
+!!! info "Reference"
+    For command options and agent settings, see the [`dstack preset` CLI reference](../reference/cli/dstack/preset.md).
+
+## Troubleshooting
+
+To trace the agent's activity, pass `--debug` to `dstack preset create`:
+
+<div class="termy">
+
+```shell
+$ dstack preset create -f preset.dstack.yml --debug
+```
+
+</div>
+
+The trace is written to `~/.dstack/presets/<id>/trace.jsonl` while the session runs. It contains the agent's messages and every tool call with its result.
+
+## Limitations
 
-!!! info "Limitations"
-    * 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`
+* 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 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).
+> Report bugs and request features on [GitHub](https://github.com/dstackai/dstack/issues), and ask questions on [Discord](https://discord.gg/u8SmfwPpMd).
 
 !!! info "What's next?"
     1. Learn how dstack [services](services.md) work

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)