[Presets] Add PD Disaggregation Support

skills/dstack-prototyping/SKILL.md

@@ -98,3 +98,17 @@ 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.
+
+## 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`.

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

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():

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,7 +251,16 @@ 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. This object differs for a trial with
+node groups and a trial without node groups.
+
+1. For a trial with node groups the fields are these and no others:
+
+```
+{"groups": [[{...}], [{...}, {...}], [{...}]], "context_length": ..., "benchmark": {...}, "learned": ..., "failed": ...}
+```
+
+2. For a trial without node groups the fields are these and no others:
 
 ```
 {"resources": {...}, "context_length": ..., "benchmark": {...}, "learned": ..., "failed": ...}
@@ -274,6 +274,20 @@ 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`: 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:
+
+  ```
+  [
+    [{"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 +315,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.
+
+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.
 <!--?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

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

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}.",

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

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):

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())

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