diff --git a/skills/dstack-prototyping/SKILL.md b/skills/dstack-prototyping/SKILL.md index c41cf2465..66201f07f 100644 --- a/skills/dstack-prototyping/SKILL.md +++ b/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`. 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/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/resources/system_prompt.md b/src/dstack/_internal/cli/services/presets/resources/system_prompt.md index 0ebc1e92b..2b2df7e20 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 @@ -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//trial.json` is one JSON object with these fields and no others: +`trials//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 --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 `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. + +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 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/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_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 "