[Presets] Add PD Disaggregation Support

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`
 

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 <fleet name> --json` to inspect a specific fleet.
+
 ## Check Serving Sources
 
 Check serving-framework sources early enough to choose the image, command,
@@ -98,3 +100,20 @@ If service verification fails because the image, install, model download,
 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.
+
+## 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`.

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 <run name> --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 <run name>`, `dstack attach <run name>`, and `ssh <run name>` 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 <run name> --job 1` or `dstack attach <run name> --job 1`.
+- In a multi-replica service, replicas are numbered from 0 in order across replica groups. Target a replica via `dstack logs <run name> --replica 1` or `dstack attach <run name> --replica 1`.
+- Attaching with a non-zero `--job` or `--replica` creates the SSH alias `ssh <run name>-<job num>-<replica num>`.
+
 ## Configuration types
 
 `dstack` supports run configurations (dev environments, tasks, and services) and infrastructure configurations (fleets, volumes, and gateways). Configuration files can be named `<name>.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 <run name>` to reconnect and restore port forwarding. The run name becomes an SSH alias (e.g., `ssh <run name>`) 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 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":"<model name>","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

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.
 
-<!--?if prompt-->
-  Do not use P/D disaggregation setups,
-  unless `## Additional instructions` explicitly allows it.
-<!--?else-->
-  Do not use P/D disaggregation setups.
-<!--?end-->
-<!--!TODO: allow P/D disaggregation and multi-node once tasks support node
-groups.-->
-
 <!--?if prompt-->
   ## Additional instructions
 
@@ -260,12 +251,23 @@ 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/<n>/trial.json` is one JSON object with these fields and no others:
+`trials/<n>/trial.json` is a JSON object.
+
+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 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": ...}
+```
+
 - `resources`: the exact resources of the instance the task ran on, in
   `dstack` resources syntax, e.g. `{"cpu": "9", "memory": "50GB", "disk":
   "200GB", "gpu": {"name": "A40", "memory": "48GB", "count": 1}}`. Read the
@@ -274,6 +276,23 @@ mindful of which specific change was the root cause.
   `dstack run get <run name> --json`, converting MiB values to GB and the
   `gpus` list into one `gpu` object with the GPU `name`, per-GPU `memory`,
   and `count`.
+- `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 <run name> --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:
+
+  ```
+  [
+    [{"cpu": "16", "memory": "64GB", "disk": "100GB"}],
+    [{"cpu": "192", "memory": "2048GB", "disk": "1000GB", "gpu": {"name": "H200", "memory": "141GB", "count": 8}},
+     {"cpu": "192", "memory": "2048GB", "disk": "1000GB", "gpu": {"name": "H200", "memory": "141GB", "count": 8}}],
+    [{"cpu": "192", "memory": "2048GB", "disk": "1000GB", "gpu": {"name": "H200", "memory": "141GB", "count": 8}}]
+  ]
+  ```
 - `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 +320,11 @@ During trials, run benchmarks via SSH inside the task, directly against the
 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.
+
+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.
+
 <!--?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
@@ -444,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.
 
@@ -540,6 +565,10 @@ trial benchmarks so that the results are comparable with each other. Attach to
 the service with `dstack attach <run name>`, which enables `ssh <run name>`
 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 <run name> --replica <replica num> --job <job num>`.
+
 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

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 = []
@@ -642,8 +643,44 @@ def _trial_entry(
 
 
 def _format_trial_gpu(record: dict[str, Any]) -> Optional[str]:
+    """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"))
+        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"])

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

@@ -12,13 +12,11 @@
     GetCoreSchemaHandler,
     PositiveInt,
     RootModel,
-    SerializerFunctionWrapHandler,
     ValidationError,
     ValidationInfo,
     conint,
     constr,
     field_validator,
-    model_serializer,
     model_validator,
 )
 from pydantic_core import CoreSchema, core_schema
@@ -1256,20 +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
-    ) -> Dict[str, Any]:
-        res = handler(self)
-        groups = res.pop("groups", None)
-        if groups is None:
-            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:

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"})

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
 

src/tests/_internal/cli/services/presets/test_prompt.py

@@ -27,19 +27,18 @@ 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,
             previous=(),
             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 "<!--?" not in text
 
     def test_renders_only_the_custom_dataset_branch(self):