From dc9444a640e03fc5f54a53c869025173ef08dfae Mon Sep 17 00:00:00 2001 From: Bihan Rana Date: Thu, 27 Aug 2026 13:26:50 +0545 Subject: [PATCH 1/6] [Presets] Allow PD disaggregation (1s iteration) --- skills/dstack/SKILL.md | 2 +- .../presets/resources/system_prompt.md | 19 ++++---- .../_internal/cli/services/presets/session.py | 28 ++++++++++- .../cli/services/presets/test_output.py | 47 +++++++++++++++++++ .../cli/services/presets/test_prompt.py | 5 +- 5 files changed, 86 insertions(+), 15 deletions(-) diff --git a/skills/dstack/SKILL.md b/skills/dstack/SKILL.md index bd840b774..f4ef32df9 100644 --- a/skills/dstack/SKILL.md +++ b/skills/dstack/SKILL.md @@ -222,7 +222,7 @@ resources: **Port forwarding:** When you specify `ports`, `dstack apply` forwards them to `localhost` while attached. Use `dstack attach ` to reconnect and restore port forwarding. The run name becomes an SSH alias (e.g., `ssh `) for direct access. -**Distributed training:** Multi-node tasks are supported (e.g., via `nodes`) and require fleets that support inter-node communication (see `placement: cluster` in fleets). +**Distributed training:** Multi-node tasks are supported (e.g., via `nodes`, or via `groups` for heterogeneous multi-node tasks) and require fleets that support inter-node communication (see `placement: cluster` in fleets). [Concept documentation](https://dstack.ai/docs/concepts/tasks.md) | [Configuration reference](https://dstack.ai/docs/reference/dstack.yml/task.md) diff --git a/src/dstack/_internal/cli/services/presets/resources/system_prompt.md b/src/dstack/_internal/cli/services/presets/resources/system_prompt.md index 0ebc1e92b..a36d49d0f 100644 --- a/src/dstack/_internal/cli/services/presets/resources/system_prompt.md +++ b/src/dstack/_internal/cli/services/presets/resources/system_prompt.md @@ -61,15 +61,6 @@ the model variant (only if `model` has `base`), the serving framework, the Docker image and dependencies, the serving framework parameters, patch the serving framework source code, generate custom kernels, and patch drivers. - - Do not use P/D disaggregation setups, - unless `## Additional instructions` explicitly allows it. - - Do not use P/D disaggregation setups. - - - ## Additional instructions @@ -263,7 +254,7 @@ mindful of which specific change was the root cause. `trials//trial.json` is one JSON object with these fields and no others: ``` -{"resources": {...}, "context_length": ..., "benchmark": {...}, "learned": ..., "failed": ...} +{"resources": {...} or "groups": [...], "context_length": ..., "benchmark": {...}, "learned": ..., "failed": ...} ``` - `resources`: the exact resources of the instance the task ran on, in @@ -274,6 +265,10 @@ mindful of which specific change was the root cause. `dstack run get --json`, converting MiB values to GB and the `gpus` list into one `gpu` object with the GPU `name`, per-GPU `memory`, and `count`. +- `groups`: replaces `resources` when the task used node groups. One entry per + group, in the order the groups appear in the task configuration; each entry + is the list of that group's nodes, one object per node, in the same syntax as + `resources` and read the same way. - `context_length`: the largest context the trial's configuration handles, found as described in `## Benchmark`; `null` only when the benchmark couldn't be done at all. @@ -301,6 +296,10 @@ During trials, run benchmarks via SSH inside the task, directly against the serving engine: use `dataset` and `concurrency``concurrency`, `input_tokens`, `output_tokens`, and `shared_prefix_tokens` from `constraints.json` and measure all trials the same way so that their results are comparable with each other. + +When the configuration serves through a router, as PD disaggregation does, the +router is the serving endpoint: run the benchmark there, against the router's +own API port, never against a prefill or decode worker. 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 diff --git a/src/dstack/_internal/cli/services/presets/session.py b/src/dstack/_internal/cli/services/presets/session.py index 68bd44446..ccc21651f 100644 --- a/src/dstack/_internal/cli/services/presets/session.py +++ b/src/dstack/_internal/cli/services/presets/session.py @@ -642,8 +642,34 @@ def _trial_entry( def _format_trial_gpu(record: dict[str, Any]) -> Optional[str]: + """A trial on one instance records a flat `resources` object; a trial that + used node groups records `groups` instead.""" + counts: dict[str, int] = {} + for node in _trial_nodes(record): + spec = _format_gpu(node.get("gpu")) + if spec: + # Insertion order is group order, so the roles read in the order they ran. + counts[spec] = counts.get(spec, 0) + 1 + if not counts: + return None + return " + ".join(spec if n == 1 else f"{spec} x{n}" for spec, n in counts.items()) + + +def _trial_nodes(record: dict[str, Any]) -> list[dict[str, Any]]: + groups = record.get("groups") + if isinstance(groups, list): + return [ + node + for group in groups + if isinstance(group, list) + for node in group + if isinstance(node, dict) + ] resources = record.get("resources") - gpu = resources.get("gpu") if isinstance(resources, dict) else None + return [resources] if isinstance(resources, dict) else [] + + +def _format_gpu(gpu: Any) -> Optional[str]: if not isinstance(gpu, dict) or not gpu.get("name"): return None text = str(gpu["name"]) diff --git a/src/tests/_internal/cli/services/presets/test_output.py b/src/tests/_internal/cli/services/presets/test_output.py index ce70c66f4..379a6472b 100644 --- a/src/tests/_internal/cli/services/presets/test_output.py +++ b/src/tests/_internal/cli/services/presets/test_output.py @@ -415,6 +415,53 @@ def test_reports_the_gpu_when_no_trial_produced_a_benchmark(self, tmp_path): assert summary["best_failed"] is None assert summary["gpu"] == "MI300X:192GB:1" + def test_reports_the_gpu_of_a_single_node_group(self, tmp_path): + from dstack._internal.cli.services.presets.session import _summarize_session_trials + + trials_dir = _write_trials( + tmp_path, + [{"groups": [[{"gpu": {"name": "MI300X", "memory": "192GB", "count": 1}}]]}], + ) + + summary = _summarize_session_trials(trials_dir) + + assert summary["gpu"] == "MI300X:192GB:1" + + def test_counts_the_worker_nodes_of_a_disaggregated_trial(self, tmp_path): + from dstack._internal.cli.services.presets.session import _summarize_session_trials + + # The CPU router has no GPU: it must neither blank the column nor split it. + h200 = {"gpu": {"name": "H200", "memory": "141GB", "count": 8}} + trials_dir = _write_trials( + tmp_path, + [{"groups": [[{"cpu": "16"}], [h200], [h200, h200]]}], + ) + + summary = _summarize_session_trials(trials_dir) + + assert summary["gpu"] == "H200:141GB:8 x3" + + def test_shows_the_split_when_roles_ran_different_gpus(self, tmp_path): + from dstack._internal.cli.services.presets.session import _summarize_session_trials + + h100 = {"gpu": {"name": "H100", "memory": "80GB", "count": 8}} + trials_dir = _write_trials( + tmp_path, + [ + { + "groups": [ + [{"gpu": {"name": "H200", "memory": "141GB", "count": 8}}], + [h100, h100], + ] + } + ], + ) + + summary = _summarize_session_trials(trials_dir) + + # Group order, not sorted: the roles read in the order they ran. + assert summary["gpu"] == "H200:141GB:8 + H100:80GB:8 x2" + def test_the_fastest_failed_trial_is_kept_when_nothing_passed(self, tmp_path): from dstack._internal.cli.services.presets.session import _summarize_session_trials diff --git a/src/tests/_internal/cli/services/presets/test_prompt.py b/src/tests/_internal/cli/services/presets/test_prompt.py index 1c163e627..160d6a642 100644 --- a/src/tests/_internal/cli/services/presets/test_prompt.py +++ b/src/tests/_internal/cli/services/presets/test_prompt.py @@ -27,7 +27,7 @@ def test_stays_byte_identical_without_user_prompt(self): assert "TODO" not in text assert "{prompt}" not in text - def test_injects_user_prompt_with_escape_clause(self): + def test_injects_user_prompt(self): text = get_preset_agent_system_prompt( user_prompt="Optimize for RAG traffic.", baseline=False, @@ -35,11 +35,10 @@ def test_injects_user_prompt_with_escape_clause(self): custom_dataset=False, ) - clause_at = text.index("unless `## Additional instructions` explicitly allows it.") section_at = text.index( "## Additional instructions\n\n```\nOptimize for RAG traffic.\n```" ) - assert clause_at < section_at < text.index("## CLI And Skills") + assert section_at < text.index("## CLI And Skills") assert "`dataset` and `concurrency``co `shared_prefix_tokens` from `constraints.json` and measure all trials the same way so that their results are comparable with each other. -When the configuration serves through a router, as PD disaggregation does, the -router is the serving endpoint: run the benchmark there, against the router's -own API port, never against a prefill or decode worker. +To benchmark a PD disaggregation setup, SSH into the job running the router and +run the benchmark directly against the router. Never benchmark a prefill or +decode worker — each handles only part of a request, so the result would not +describe the configuration. 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 From 13b922d0f4fb18fa8deb63f7faa9690f7f4e5822 Mon Sep 17 00:00:00 2001 From: Bihan Rana Date: Sun, 30 Aug 2026 09:39:44 +0545 Subject: [PATCH 3/6] [Presets] Use the groups syntax for grouped services in exported and stored YAML --- .../_internal/cli/services/presets/export.py | 7 +++++- .../_internal/cli/services/presets/store.py | 5 +++- .../_internal/core/models/configurations.py | 8 +++++- .../cli/services/presets/test_export.py | 24 ++++++++++++++++++ .../cli/services/presets/test_store.py | 25 ++++++++++++++++++- .../core/models/test_configurations.py | 17 +++++++++++++ 6 files changed, 82 insertions(+), 4 deletions(-) diff --git a/src/dstack/_internal/cli/services/presets/export.py b/src/dstack/_internal/cli/services/presets/export.py index e3c5d4313..f6247aa21 100644 --- a/src/dstack/_internal/cli/services/presets/export.py +++ b/src/dstack/_internal/cli/services/presets/export.py @@ -54,7 +54,12 @@ def export_preset( if target.exists(): raise CLIError(f"{target} already exists. Use --force to overwrite") destination.parent.mkdir(parents=True, exist_ok=True) - destination.write_text(yaml.safe_dump(service.model_dump(mode="json"), sort_keys=False)) + destination.write_text( + yaml.safe_dump( + service.model_dump(mode="json", context={"keep_groups": True}), + sort_keys=False, + ) + ) for source, target in copies: target.parent.mkdir(parents=True, exist_ok=True) if source.is_dir(): diff --git a/src/dstack/_internal/cli/services/presets/store.py b/src/dstack/_internal/cli/services/presets/store.py index 621168460..c5910c6be 100644 --- a/src/dstack/_internal/cli/services/presets/store.py +++ b/src/dstack/_internal/cli/services/presets/store.py @@ -89,7 +89,10 @@ def save(self, preset: AnyStoredPreset) -> Path: preset = preset.model_copy(deep=True) for mapping in preset.service.files: mapping.local_path = _relative_to_preset_dir(mapping.local_path, directory) - content = yaml.safe_dump(preset.model_dump(mode="json"), sort_keys=False) + content = yaml.safe_dump( + preset.model_dump(mode="json", context={"keep_groups": True}), + sort_keys=False, + ) fd, temporary_path = tempfile.mkstemp( dir=directory, prefix=f".{preset.id}.", diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py index 344d12aca..47945f7df 100644 --- a/src/dstack/_internal/core/models/configurations.py +++ b/src/dstack/_internal/core/models/configurations.py @@ -12,6 +12,7 @@ GetCoreSchemaHandler, PositiveInt, RootModel, + SerializationInfo, SerializerFunctionWrapHandler, ValidationError, ValidationInfo, @@ -1258,12 +1259,17 @@ def _normalize_legacy_replica_groups(cls, data: Any) -> Any: @model_serializer(mode="wrap") def _serialize_legacy_replica_groups( - self, handler: SerializerFunctionWrapHandler + self, handler: SerializerFunctionWrapHandler, info: SerializationInfo ) -> Dict[str, Any]: res = handler(self) groups = res.pop("groups", None) if groups is None: return res + # keep_groups=True: dump `groups:` for `dstack preset export` (`*.dstack.yml`) + # and PresetStore.save (`preset.yml`). + if info.context and info.context.get("keep_groups"): + res["groups"] = groups + return res for group in groups: if "replicas" in group: group["count"] = group.pop("replicas") diff --git a/src/tests/_internal/cli/services/presets/test_export.py b/src/tests/_internal/cli/services/presets/test_export.py index e8e17ba79..d1244252a 100644 --- a/src/tests/_internal/cli/services/presets/test_export.py +++ b/src/tests/_internal/cli/services/presets/test_export.py @@ -70,6 +70,30 @@ def test_exports_a_deployable_service_configuration_with_its_files(self, tmp_pat ) assert ServiceConfiguration.model_validate(data).model is not None + def test_exports_replica_groups_in_the_groups_syntax(self, tmp_path: Path): + store = PresetStore(tmp_path / "presets") + preset = get_preset() + preset.service = ServiceConfiguration.model_validate( + { + "port": 8000, + "model": "meta-llama/Llama-3.2-3B-Instruct", + "groups": [ + {"replicas": 1, "commands": ["smg launch"]}, + {"replicas": 1, "commands": ["python -m sglang.launch_server"]}, + ], + } + ) + preset_dir = store.save(preset).parent + destination = tmp_path / "llama.dstack.yml" + + export_preset(preset, preset_dir=preset_dir, destination=destination, force=False) + + data = yaml.safe_load(destination.read_text()) + assert "groups" in data + assert data.get("replicas") is None + assert "replicas" in data["groups"][0] + assert "count" not in data["groups"][0] + def test_names_the_service_after_the_preset(self, tmp_path: Path): store = PresetStore(tmp_path / "presets") preset = get_preset().model_copy(update={"name": "qwen-fast"}) diff --git a/src/tests/_internal/cli/services/presets/test_store.py b/src/tests/_internal/cli/services/presets/test_store.py index a2895aa79..523032eab 100644 --- a/src/tests/_internal/cli/services/presets/test_store.py +++ b/src/tests/_internal/cli/services/presets/test_store.py @@ -11,7 +11,7 @@ from dstack._internal.cli.services.presets.store import PresetStore from dstack._internal.compat import IS_WINDOWS from dstack._internal.core.errors import CLIError, ConfigurationError -from dstack._internal.core.models.configurations import PresetConfiguration +from dstack._internal.core.models.configurations import PresetConfiguration, ServiceConfiguration from dstack._internal.core.models.envs import EnvSentinel from dstack._internal.core.models.files import FilePathMapping from dstack._internal.core.models.presets import PortablePreset @@ -49,6 +49,29 @@ def test_saves_and_lists_self_contained_preset(self, tmp_path: Path): assert store.get(preset.id) == preset assert not list(path.parent.glob("*.tmp")) + def test_saves_replica_groups_in_the_groups_syntax(self, tmp_path: Path): + store = PresetStore(tmp_path / "presets") + preset = get_preset() + preset.service = ServiceConfiguration.model_validate( + { + "port": 8000, + "model": "meta-llama/Llama-3.2-3B-Instruct", + "groups": [ + {"replicas": 1, "commands": ["smg launch"]}, + {"replicas": 1, "commands": ["python -m sglang.launch_server"]}, + ], + } + ) + + path = store.save(preset) + + data = yaml.safe_load(path.read_text()) + service = data["service"] + assert "groups" in service + assert service.get("replicas") is None + assert "replicas" in service["groups"][0] + assert "count" not in service["groups"][0] + def test_a_verified_document_loads_as_a_verified_preset(self, tmp_path: Path): store = PresetStore(tmp_path / "presets") store.save(get_preset()) diff --git a/src/tests/_internal/core/models/test_configurations.py b/src/tests/_internal/core/models/test_configurations.py index e0d91c718..64c5e714e 100644 --- a/src/tests/_internal/core/models/test_configurations.py +++ b/src/tests/_internal/core/models/test_configurations.py @@ -1196,6 +1196,20 @@ def test_dumped_json_parses_as_0_21_client(self): dumped = parsed.model_dump() validate_extra_ignore(_Legacy021Service, dumped) + def test_dump_keep_groups_context_keeps_groups(self): + parsed = parse_run_configuration( + { + "type": "service", + "port": 8000, + "groups": [{"replicas": 1, "commands": ["x"]}], + } + ) + dumped = parsed.model_dump(mode="json", context={"keep_groups": True}) + assert "groups" in dumped + assert dumped.get("replicas") is None + assert "replicas" in dumped["groups"][0] + assert "count" not in dumped["groups"][0] + def test_homogeneous_dump_has_no_groups_key(self): parsed = parse_run_configuration( { @@ -1208,6 +1222,9 @@ def test_homogeneous_dump_has_no_groups_key(self): dumped = parsed.model_dump() assert "groups" not in dumped assert dumped["replicas"] == {"min": 2, "max": 2} + kept = parsed.model_dump(mode="json", context={"keep_groups": True}) + assert "groups" not in kept + assert kept["replicas"] == {"min": 2, "max": 2} def test_replicas_and_groups_rejected(self): with pytest.raises(ConfigurationError, match="mutually exclusive"): From 354c25f4f701a7ce35fbed21d189626320dd3c68 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Mon, 31 Aug 2026 16:24:49 +0200 Subject: [PATCH 4/6] [Presets] Document PD disaggregation across skills and the preset system prompt Cover node groups (tasks) and replica groups (services) in the dstack and dstack-prototyping skills and the preset system prompt: replica/job targeting for logs/attach/ssh, SSH alias naming, cluster placement for PD, per-group sleep-infinity for prototyping, and the groups-based trial.json format. Co-Authored-By: Claude Fable 5 --- skills/dstack-prototyping/SKILL.md | 31 ++++++++------ skills/dstack/SKILL.md | 12 +++++- .../presets/resources/system_prompt.md | 42 ++++++++++++------- .../_internal/cli/services/presets/session.py | 1 + 4 files changed, 56 insertions(+), 30 deletions(-) diff --git a/skills/dstack-prototyping/SKILL.md b/skills/dstack-prototyping/SKILL.md index 66201f07f..552d00dad 100644 --- a/skills/dstack-prototyping/SKILL.md +++ b/skills/dstack-prototyping/SKILL.md @@ -27,6 +27,8 @@ Pick the offer whose hardware best fits the goal at hand. Only when several offe Fetch `https://dstack.ai/docs/concepts/backends.md` and classify backends from the fetched document, not from memory. +If the intention is to use PD disaggregation, the fleet must use `placement: cluster`. Since PD disaggregation implies running a router, unlike workers that must run on GPUs, the router normally should run on a CPU instance. Use `dstack fleet` to see existing fleets and `dstack fleet get --json` to inspect a specific fleet. + ## Check Serving Sources Check serving-framework sources early enough to choose the image, command, @@ -99,16 +101,19 @@ command, resources, cache, or model behavior needs to change, go back to a task. If the tested serving setup is still right and only the dstack service configuration is wrong, fix the configuration and submit the service again. -## Aggregated Vs Disaggregated - -There are two types of inference service, aggregated and disaggregated. - -For aggregated inference, prototype and convert as described in -`## Use A Task Before Service` and `## Verify As A Service`. - -Disaggregated inference differs only in shape: the task uses node groups — one -group each for the router, the prefill workers and the decode workers — and the -service uses replica groups, one per role. - -See `https://dstack.ai/docs/concepts/tasks.md#node-groups` and -`https://dstack.ai/docs/concepts/services.md#pd-disaggregation`. +## PD disaggregation + +If the intention is to use PD disaggregation: + +- Use node groups for the task and replica groups for the service: tasks' node + groups are the equivalent of services' replica groups. +- In both cases, you run a router and prefill/decode workers separately, and + you need to use a fleet with an interconnect (`placement: cluster`). +- With tasks, still use `sleep infinity` even when using `groups` (set it in + each group's `commands`; top-level `commands` is not allowed with `groups`), + and run the actual commands on each node interactively over SSH. +- When testing inference, call the router endpoint, not the workers directly + (unless you want to test if they are alive). +- Look for "Node groups" and "PD disaggregation" in + `https://dstack.ai/docs/concepts/tasks.md` and "Replica groups" and + "PD disaggregation" in `https://dstack.ai/docs/concepts/services.md`. diff --git a/skills/dstack/SKILL.md b/skills/dstack/SKILL.md index f4ef32df9..234a5e54b 100644 --- a/skills/dstack/SKILL.md +++ b/skills/dstack/SKILL.md @@ -148,6 +148,14 @@ If background attach fails in the sandbox (permissions writing `~/.dstack` or `~ **"Connect to" or "open" a dev environment:** If a dev environment is already running, use `dstack attach --logs` (agent runs it in the background by default) to surface the IDE URL (`cursor://`, `vscode://`, etc.) and SSH alias. If sandboxed attach fails, request escalation or ask the user to run attach locally and share the link. +### Distributed tasks and multi-replica services + +Unless you use **Distributed tasks** (see `### 2. Tasks`) or **Multi-replica services** (see `### 3. Services`), both tasks and services run on a single node. That's why `dstack logs `, `dstack attach `, and `ssh ` default to replica 0/job 0. + +- In a distributed task, each node runs its own job, numbered from 0 in order across node groups. Target a node via `dstack logs --job 1` or `dstack attach --job 1`. +- In a multi-replica service, replicas are numbered from 0 in order across replica groups. Target a replica via `dstack logs --replica 1` or `dstack attach --replica 1`. +- Attaching with a non-zero `--job` or `--replica` creates the SSH alias `ssh --`. + ## Configuration types `dstack` supports run configurations (dev environments, tasks, and services) and infrastructure configurations (fleets, volumes, and gateways). Configuration files can be named `.dstack.yml` or simply `.dstack.yml`. @@ -222,7 +230,7 @@ resources: **Port forwarding:** When you specify `ports`, `dstack apply` forwards them to `localhost` while attached. Use `dstack attach ` to reconnect and restore port forwarding. The run name becomes an SSH alias (e.g., `ssh `) for direct access. -**Distributed training:** Multi-node tasks are supported (e.g., via `nodes`, or via `groups` for heterogeneous multi-node tasks) and require fleets that support inter-node communication (see `placement: cluster` in fleets). +**Distributed tasks:** Set `nodes` to run a task across multiple nodes, or use `groups` to define node groups, each with its own `nodes` count, `resources`, `commands`, and `ports` (`groups` and top-level `nodes` are mutually exclusive). Requires a fleet that supports inter-node communication (see `placement: cluster` in fleets). [Concept documentation](https://dstack.ai/docs/concepts/tasks.md) | [Configuration reference](https://dstack.ai/docs/reference/dstack.yml/task.md) @@ -262,6 +270,8 @@ resources: -d '{"model":"","messages":[{"role":"user","content":"Hello"}],"max_tokens":64}' ``` +**Multi-replica services:** Set `replicas` to run multiple replicas, or use `groups` to define replica groups, each with its own `replicas` count, `resources`, and `commands` (`groups` and top-level `replicas` are mutually exclusive). If replicas require an interconnect (e.g., PD disaggregation), the service must run on a fleet with `placement: cluster`. + [Concept documentation](https://dstack.ai/docs/concepts/services.md) | [Configuration reference](https://dstack.ai/docs/reference/dstack.yml/service.md) ### 4. Fleets diff --git a/src/dstack/_internal/cli/services/presets/resources/system_prompt.md b/src/dstack/_internal/cli/services/presets/resources/system_prompt.md index 2b2df7e20..b5af21f72 100644 --- a/src/dstack/_internal/cli/services/presets/resources/system_prompt.md +++ b/src/dstack/_internal/cli/services/presets/resources/system_prompt.md @@ -251,19 +251,21 @@ also failed when its benchmark does not meet the constraints (see `# Constraints`). When a trial that changed several things fails, be mindful of which specific change was the root cause. -`trials//trial.json` is a JSON object. This object differs for a trial with -node groups and a trial without node groups. +`trials//trial.json` is a JSON object. -1. For a trial with node groups the fields are these and no others: +1. In case the task is not using PD disaggregation (and thus no node groups), + the fields are these and no others: ``` -{"groups": [[{...}], [{...}, {...}], [{...}]], "context_length": ..., "benchmark": {...}, "learned": ..., "failed": ...} +{"resources": {...}, "context_length": ..., "benchmark": {...}, "learned": ..., "failed": ...} ``` -2. For a trial without node groups the fields are these and no others: +2. In case the task is using PD disaggregation (and thus node groups), instead + of a single `resources` it includes `groups`, and the fields are these and + no others: ``` -{"resources": {...}, "context_length": ..., "benchmark": {...}, "learned": ..., "failed": ...} +{"groups": [[{...}], [{...}, {...}], [{...}]], "context_length": ..., "benchmark": {...}, "learned": ..., "failed": ...} ``` - `resources`: the exact resources of the instance the task ran on, in @@ -274,11 +276,14 @@ node groups and a trial without node groups. `dstack run get --json`, converting MiB values to GB and the `gpus` list into one `gpu` object with the GPU `name`, per-GPU `memory`, and `count`. -- `groups`: a list of node groups, in the order they appear in the task - configuration. Each node group is a list of its nodes. Each node is the exact - resources of the instance that node ran on, e.g. a PD disaggregation task - with a one-node router group, a two-node prefill group and a one-node decode - group records: +- `groups`: the exact resources of each instance the task ran on, per node + group; groups are in the order they appear in the task configuration, and + each group lists the resources of its instances, in the same format as + `resources`. Read the actual values from each job's latest submission's + `job_runtime_data.offer.instance.resources` in + `dstack run get --json`, converting them as for `resources`. + E.g. a one-node router group, a two-node prefill group, and a one-node + decode group record: ``` [ @@ -316,10 +321,10 @@ serving engine: use `dataset` and `concurrency``co `shared_prefix_tokens` from `constraints.json` and measure all trials the same way so that their results are comparable with each other. -To benchmark a PD disaggregation setup, SSH into the job running the router and -run the benchmark directly against the router. Never benchmark a prefill or -decode worker — each handles only part of a request, so the result would not -describe the configuration. +In case the task is using PD disaggregation, run benchmarks via SSH inside +the router node, directly against the router engine. Never benchmark prefill +or decode workers — each handles only part of a request. + 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 @@ -463,7 +468,8 @@ patches are correct (and will exactly replicate the result). # Task Usage Trials are done entirely using `dstack` tasks. For maximum efficiency, it is a -requirement that you always set the task `commands` to `sleep infinity` and +requirement that you always set the task `commands` to `sleep infinity` (for +a task with `groups`, in each group's `commands`) and run commands inside the task interactively, via SSH. It is important that you follow the `/dstack-prototyping` skill when working with tasks. @@ -559,6 +565,10 @@ trial benchmarks so that the results are comparable with each other. Attach to the service with `dstack attach `, which enables `ssh ` into the replica. +In case the service is using PD disaggregation, run the final benchmark +inside the router replica, directly against the router engine. Attach to it +via `dstack attach --replica --job `. + If the service or its benchmark cannot be completed, stop that service, pick the next-best trial, and repeat, until a service is verified or there are no unverified trials left. Report the result accordingly (see diff --git a/src/dstack/_internal/cli/services/presets/session.py b/src/dstack/_internal/cli/services/presets/session.py index ccc21651f..a4fb69d59 100644 --- a/src/dstack/_internal/cli/services/presets/session.py +++ b/src/dstack/_internal/cli/services/presets/session.py @@ -562,6 +562,7 @@ def _read_last_session_verification(path: Path) -> Optional[dict[str, Any]]: def _summarize_session_trials(path: Path) -> Optional[dict[str, Any]]: + # TODO: Refactor this crap - must be explicit what is this and where is it used; also dicts are prohibited in dstack repo """A trial directory without `trial.json` is still in flight and is not counted.""" records = [] From 76e11a5c641c21366eb40a30d2d298efb53c01dd Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Mon, 31 Aug 2026 16:37:44 +0200 Subject: [PATCH 5/6] [Presets] Key trial.json format on node groups and fix replica/job default wording Co-Authored-By: Claude Fable 5 --- skills/dstack/SKILL.md | 2 +- .../cli/services/presets/resources/system_prompt.md | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/skills/dstack/SKILL.md b/skills/dstack/SKILL.md index 234a5e54b..d9bf1e7b7 100644 --- a/skills/dstack/SKILL.md +++ b/skills/dstack/SKILL.md @@ -150,7 +150,7 @@ If background attach fails in the sandbox (permissions writing `~/.dstack` or `~ ### Distributed tasks and multi-replica services -Unless you use **Distributed tasks** (see `### 2. Tasks`) or **Multi-replica services** (see `### 3. Services`), both tasks and services run on a single node. That's why `dstack logs `, `dstack attach `, and `ssh ` default to replica 0/job 0. +Unless you use **Distributed tasks** (see `### 2. Tasks`) or **Multi-replica services** (see `### 3. Services`), both tasks and services run on a single node. That's why `dstack logs `, `dstack attach `, and `ssh ` default to the first replica/job. - In a distributed task, each node runs its own job, numbered from 0 in order across node groups. Target a node via `dstack logs --job 1` or `dstack attach --job 1`. - In a multi-replica service, replicas are numbered from 0 in order across replica groups. Target a replica via `dstack logs --replica 1` or `dstack attach --replica 1`. diff --git a/src/dstack/_internal/cli/services/presets/resources/system_prompt.md b/src/dstack/_internal/cli/services/presets/resources/system_prompt.md index b5af21f72..0e439a0d0 100644 --- a/src/dstack/_internal/cli/services/presets/resources/system_prompt.md +++ b/src/dstack/_internal/cli/services/presets/resources/system_prompt.md @@ -253,16 +253,16 @@ mindful of which specific change was the root cause. `trials//trial.json` is a JSON object. -1. In case the task is not using PD disaggregation (and thus no node groups), - the fields are these and no others: +1. In case the task is not using node groups, the fields are these and no + others: ``` {"resources": {...}, "context_length": ..., "benchmark": {...}, "learned": ..., "failed": ...} ``` -2. In case the task is using PD disaggregation (and thus node groups), instead - of a single `resources` it includes `groups`, and the fields are these and - no others: +2. In case the task is using node groups (e.g. for PD disaggregation), + instead of a single `resources` it includes `groups`, and the fields are + these and no others: ``` {"groups": [[{...}], [{...}, {...}], [{...}]], "context_length": ..., "benchmark": {...}, "learned": ..., "failed": ...} From ea5cb683e0c7b778bdc67245d1e867cc87bef951 Mon Sep 17 00:00:00 2001 From: Bihan Rana Date: Wed, 2 Sep 2026 12:51:47 +0545 Subject: [PATCH 6/6] Drop groups serializer and keep groups --- mkdocs/docs/concepts/presets.md | 1 - .../_internal/cli/services/presets/export.py | 7 +- .../_internal/cli/services/presets/session.py | 14 +++- .../_internal/cli/services/presets/store.py | 5 +- .../_internal/core/models/configurations.py | 22 ----- .../core/models/test_configurations.py | 80 ++++--------------- 6 files changed, 29 insertions(+), 100 deletions(-) diff --git a/mkdocs/docs/concepts/presets.md b/mkdocs/docs/concepts/presets.md index b48b7d064..8f631b1a3 100644 --- a/mkdocs/docs/concepts/presets.md +++ b/mkdocs/docs/concepts/presets.md @@ -362,7 +362,6 @@ At the same time, it's recommended to create presets using your own agent — ei ## Limitations * Currently, the agent doesn't upload compiled binaries anywhere; patches compile at runtime -* Doesn't support PD disaggregation (coming soon) * The registry doesn't support public presets (coming soon) * Doesn't support ranges for `concurrency` diff --git a/src/dstack/_internal/cli/services/presets/export.py b/src/dstack/_internal/cli/services/presets/export.py index f6247aa21..e3c5d4313 100644 --- a/src/dstack/_internal/cli/services/presets/export.py +++ b/src/dstack/_internal/cli/services/presets/export.py @@ -54,12 +54,7 @@ def export_preset( if target.exists(): raise CLIError(f"{target} already exists. Use --force to overwrite") destination.parent.mkdir(parents=True, exist_ok=True) - destination.write_text( - yaml.safe_dump( - service.model_dump(mode="json", context={"keep_groups": True}), - sort_keys=False, - ) - ) + destination.write_text(yaml.safe_dump(service.model_dump(mode="json"), sort_keys=False)) for source, target in copies: target.parent.mkdir(parents=True, exist_ok=True) if source.is_dir(): diff --git a/src/dstack/_internal/cli/services/presets/session.py b/src/dstack/_internal/cli/services/presets/session.py index a4fb69d59..81df2c951 100644 --- a/src/dstack/_internal/cli/services/presets/session.py +++ b/src/dstack/_internal/cli/services/presets/session.py @@ -643,8 +643,18 @@ def _trial_entry( def _format_trial_gpu(record: dict[str, Any]) -> Optional[str]: - """A trial on one instance records a flat `resources` object; a trial that - used node groups records `groups` instead.""" + """A trial records its hardware in one of two formats, as described in + `system_prompt.md`: `{"resources": {...}}` without node groups, and + `{"groups": [[...], ...]}` with them. + + Without node groups it returns that one instance's GPU, e.g. `H200:141GB:1`. + With node groups it returns the GPUs of every node, e.g. `H200:141GB:1 x5` + for a router plus 2 prefill and 3 decode nodes, or + `H200:141GB:1 x2 + H100:80GB:1 x3` when the GPU models differ. + + The value fills the `RESOURCES` column of a session row in + `dstack preset list -v`. + """ counts: dict[str, int] = {} for node in _trial_nodes(record): spec = _format_gpu(node.get("gpu")) diff --git a/src/dstack/_internal/cli/services/presets/store.py b/src/dstack/_internal/cli/services/presets/store.py index c5910c6be..621168460 100644 --- a/src/dstack/_internal/cli/services/presets/store.py +++ b/src/dstack/_internal/cli/services/presets/store.py @@ -89,10 +89,7 @@ def save(self, preset: AnyStoredPreset) -> Path: preset = preset.model_copy(deep=True) for mapping in preset.service.files: mapping.local_path = _relative_to_preset_dir(mapping.local_path, directory) - content = yaml.safe_dump( - preset.model_dump(mode="json", context={"keep_groups": True}), - sort_keys=False, - ) + content = yaml.safe_dump(preset.model_dump(mode="json"), sort_keys=False) fd, temporary_path = tempfile.mkstemp( dir=directory, prefix=f".{preset.id}.", diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py index 47945f7df..1b52d0140 100644 --- a/src/dstack/_internal/core/models/configurations.py +++ b/src/dstack/_internal/core/models/configurations.py @@ -12,14 +12,11 @@ GetCoreSchemaHandler, PositiveInt, RootModel, - SerializationInfo, - SerializerFunctionWrapHandler, ValidationError, ValidationInfo, conint, constr, field_validator, - model_serializer, model_validator, ) from pydantic_core import CoreSchema, core_schema @@ -1257,25 +1254,6 @@ def _normalize_legacy_replica_groups(cls, data: Any) -> Any: raise ValueError("`replicas` and `groups` are mutually exclusive") return data - @model_serializer(mode="wrap") - def _serialize_legacy_replica_groups( - self, handler: SerializerFunctionWrapHandler, info: SerializationInfo - ) -> Dict[str, Any]: - res = handler(self) - groups = res.pop("groups", None) - if groups is None: - return res - # keep_groups=True: dump `groups:` for `dstack preset export` (`*.dstack.yml`) - # and PresetStore.save (`preset.yml`). - if info.context and info.context.get("keep_groups"): - res["groups"] = groups - return res - for group in groups: - if "replicas" in group: - group["count"] = group.pop("replicas") - res["replicas"] = groups - return res - @field_validator("port") @classmethod def convert_port(cls, v) -> PortMapping: diff --git a/src/tests/_internal/core/models/test_configurations.py b/src/tests/_internal/core/models/test_configurations.py index 64c5e714e..d72b5ba9d 100644 --- a/src/tests/_internal/core/models/test_configurations.py +++ b/src/tests/_internal/core/models/test_configurations.py @@ -1,12 +1,11 @@ from copy import deepcopy -from typing import Any, Optional, Union +from typing import Any, Optional import pytest -from pydantic import ValidationError, model_validator -from typing_extensions import Self +from pydantic import ValidationError from dstack._internal.core.errors import ConfigurationError -from dstack._internal.core.models.common import CoreModel, RegistryAuth, validate_extra_ignore +from dstack._internal.core.models.common import RegistryAuth from dstack._internal.core.models.configurations import ( DevEnvironmentConfigurationParams, PresetConfiguration, @@ -1096,29 +1095,6 @@ def test_accepts_top_level_resources_with_groups(self): assert parsed.resources.gpu.name == ["H100"] -class _Legacy021ReplicaGroup(CoreModel): - """0.21-shaped group: size is `count`, no `groups` parent field.""" - - count: Range[int] - commands: list[str] = [] - - -class _Legacy021Service(CoreModel): - """Stand-in for a 0.21 client that does not know `groups`.""" - - commands: list[str] = [] - image: Optional[str] = None - replicas: Optional[Union[list[_Legacy021ReplicaGroup], Range[int]]] = None - - @model_validator(mode="after") - def check_image_or_commands_present(self) -> Self: - if isinstance(self.replicas, list): - return self - if not self.commands and self.image is None: - raise ValueError("Either `commands` or `image` must be set") - return self - - class TestServiceGroupsPhase1: def test_legacy_replicas_list_parses_to_groups(self): parsed = parse_run_configuration( @@ -1154,7 +1130,7 @@ def test_new_groups_syntax_parses_identically(self): assert new.replicas is None assert legacy.groups == new.groups - def test_dump_is_legacy_canonical(self): + def test_dump_is_groups_canonical(self): parsed = parse_run_configuration( { "type": "service", @@ -1162,15 +1138,16 @@ def test_dump_is_legacy_canonical(self): "groups": [{"replicas": 1, "commands": ["x"]}], } ) + # The legacy `replicas: [{count: ...}]` shape is produced only for old + # clients, by `server/compatibility/runs.py`, not by the model. dumped = parsed.model_dump() - assert "groups" not in dumped - assert isinstance(dumped["replicas"], list) - assert "count" in dumped["replicas"][0] - assert "replicas" not in dumped["replicas"][0] + assert dumped["replicas"] is None + assert isinstance(dumped["groups"], list) + assert "replicas" in dumped["groups"][0] + assert "count" not in dumped["groups"][0] dumped_json = parsed.model_dump(mode="json") - assert "groups" not in dumped_json - assert "count" in dumped_json["replicas"][0] - assert "replicas" not in dumped_json["replicas"][0] + assert dumped_json["replicas"] is None + assert "replicas" in dumped_json["groups"][0] def test_dump_validate_is_fixed_point(self): parsed = parse_run_configuration( @@ -1185,32 +1162,7 @@ def test_dump_validate_is_fixed_point(self): twice = ServiceConfiguration.model_validate(once.model_dump()) assert once.model_dump() == twice.model_dump() == parsed.model_dump() - def test_dumped_json_parses_as_0_21_client(self): - parsed = parse_run_configuration( - { - "type": "service", - "port": 8000, - "groups": [{"replicas": 1, "commands": ["x"]}], - } - ) - dumped = parsed.model_dump() - validate_extra_ignore(_Legacy021Service, dumped) - - def test_dump_keep_groups_context_keeps_groups(self): - parsed = parse_run_configuration( - { - "type": "service", - "port": 8000, - "groups": [{"replicas": 1, "commands": ["x"]}], - } - ) - dumped = parsed.model_dump(mode="json", context={"keep_groups": True}) - assert "groups" in dumped - assert dumped.get("replicas") is None - assert "replicas" in dumped["groups"][0] - assert "count" not in dumped["groups"][0] - - def test_homogeneous_dump_has_no_groups_key(self): + def test_homogeneous_dump_has_null_groups(self): parsed = parse_run_configuration( { "type": "service", @@ -1219,12 +1171,10 @@ def test_homogeneous_dump_has_no_groups_key(self): "replicas": 2, } ) + # Nothing strips the key now that the model no longer rewrites groups. dumped = parsed.model_dump() - assert "groups" not in dumped + assert dumped["groups"] is None assert dumped["replicas"] == {"min": 2, "max": 2} - kept = parsed.model_dump(mode="json", context={"keep_groups": True}) - assert "groups" not in kept - assert kept["replicas"] == {"min": 2, "max": 2} def test_replicas_and_groups_rejected(self): with pytest.raises(ConfigurationError, match="mutually exclusive"):