From dc9444a640e03fc5f54a53c869025173ef08dfae Mon Sep 17 00:00:00 2001 From: Bihan Rana Date: Thu, 27 Aug 2026 13:26:50 +0545 Subject: [PATCH 1/3] [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/3] [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"):