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/skills/dstack-prototyping/SKILL.md b/skills/dstack-prototyping/SKILL.md index c41cf2465..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, @@ -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`. diff --git a/skills/dstack/SKILL.md b/skills/dstack/SKILL.md index bd840b774..d9bf1e7b7 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 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`. +- 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`) 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 0ebc1e92b..0e439a0d0 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,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//trial.json` is one JSON object with these fields and no others: +`trials//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 --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 --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 `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. + +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 @@ -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 `, 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 68bd44446..81df2c951 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 = [] @@ -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"]) diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py index 344d12aca..1b52d0140 100644 --- a/src/dstack/_internal/core/models/configurations.py +++ b/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: 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 "