diff --git a/mkdocs.yml b/mkdocs.yml index 64aac83c6..84ac0f3b4 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -297,7 +297,7 @@ nav: - Services: docs/concepts/services.md - Volumes: docs/concepts/volumes.md - More: - - Endpoints: docs/concepts/endpoints.md + - Presets: docs/concepts/presets.md - Gateways: docs/concepts/gateways.md - Secrets: docs/concepts/secrets.md - Projects: docs/concepts/projects.md @@ -348,7 +348,7 @@ nav: - fleet: docs/reference/dstack.yml/fleet.md - gateway: docs/reference/dstack.yml/gateway.md - volume: docs/reference/dstack.yml/volume.md - - endpoint: docs/reference/dstack.yml/endpoint.md + - preset: docs/reference/dstack.yml/preset.md - server/config.yml: docs/reference/server/config.yml.md - CLI: - dstack server: docs/reference/cli/dstack/server.md @@ -368,7 +368,7 @@ nav: - dstack volume: docs/reference/cli/dstack/volume.md - dstack gateway: docs/reference/cli/dstack/gateway.md - dstack secret: docs/reference/cli/dstack/secret.md - - dstack endpoint: docs/reference/cli/dstack/endpoint.md + - dstack preset: docs/reference/cli/dstack/preset.md - dstack export: docs/reference/cli/dstack/export.md - dstack import: docs/reference/cli/dstack/import.md - HTTP API: diff --git a/mkdocs/docs/concepts/endpoints.md b/mkdocs/docs/concepts/endpoints.md deleted file mode 100644 index 700756d4c..000000000 --- a/mkdocs/docs/concepts/endpoints.md +++ /dev/null @@ -1,180 +0,0 @@ ---- -title: Endpoints -description: Creating and reusing optimized model inference endpoint configurations ---- - -# Endpoints - -An endpoint configuration lets you use an agent to create a preset: a validated and optimized model inference configuration. Once created, the preset can be reused to deploy model inference on validated hardware without an agent. - -The value of presets comes from combining two fundamental features: agent-driven model inference optimization and the `dstack` [service](services.md) primitive, which can deploy model inference to any cloud, Kubernetes, or on-prem cluster. - -> The endpoints feature is experimental and may change. - -??? info "Prerequisites" - Before using endpoint presets, make sure you’ve [installed](../installation.md) the server and CLI, and created a [fleet](fleets.md). - - Creating an endpoint preset requires the `claude` CLI to be installed on the machine where you create a preset. - -## Define an endpoint - -Before you can create or reuse an endpoint preset, you first have to define an endpoint configuration. The filename must end with `.dstack.yml`. - -
- -```yaml -type: endpoint -name: qwen25-7b - -model: - base: Qwen/Qwen2.5-7B-Instruct - -env: - - HF_TOKEN -``` - -
- -Since `base` is specified, the preset can use any compatible variant of the base model, including a different precision, quantization, or trusted fork. - -If you want to deploy an exact model, set `model` directly to the repo of that model: - -
- -```yaml -type: endpoint -name: qwen25-7b - -model: Qwen/Qwen2.5-7B-Instruct - -env: - - HF_TOKEN -``` - -
- -Set `context_length` to require a minimum context length. Placement properties, -including `fleets`, `backends`, `max_price`, and `spot_policy`, constrain both -creation and reuse. Environment variables such as `HF_TOKEN` can be passed -through `env`. - -See the [reference](../reference/dstack.yml/endpoint.md) for all supported configuration options. - -## Create a preset - -To create a preset, pass the configuration file to the `dstack endpoint preset create` command: - -
- -```shell -$ dstack endpoint preset create -f endpoint.dstack.yml -[2026-07-15 11:32:01] Starting endpoint preset creation for Qwen/Qwen2.5-7B-Instruct. Allowed fleets: gpu-fleet. -[2026-07-15 11:41:06] Prototype task qwen25-7b-a1b2c3-2 verified vLLM on an L4:24GB. -[2026-07-15 11:52:06] Final service qwen25-7b-a1b2c3-3 verified with context length 32768. -[2026-07-15 11:52:18] Benchmark via guidellm 0.7.1: 32/32 requests succeeded. -[2026-07-15 11:52:18] Saved endpoint preset 8f3a12c4 for Qwen/Qwen2.5-7B-Instruct. -``` - -
- -This command executes entirely locally and uses the locally installed `claude` CLI along with `dstack`'s bundled skills. The agent uses a `dstack` task to find the best serving configuration for the available fleet offers. It then submits the configuration as a `dstack` service for a final benchmark. The validated preset is saved locally under `~/.dstack/presets`. - -??? info "Claude configuration" - Preset creation uses the existing `claude` login. To use an Anthropic API key instead, set: - - ```shell - export DSTACK_AGENT_ANTHROPIC_API_KEY=... - ``` - - By default, the agent uses `claude-opus-4-8` and the default `claude` CLI effort. To override them, set: - - ```shell - export DSTACK_AGENT_ANTHROPIC_MODEL=claude-fable-5 - export DSTACK_AGENT_CLAUDE_EFFORT=high - ``` - - Supported effort levels are `low`, `medium`, `high`, `xhigh`, and `max`. - -## List presets - -Use `dstack endpoint preset` to list existing presets: - -
- -```shell -$ dstack endpoint preset list - MODEL GPU CONTEXT BENCHMARK CREATED - Qwen/Qwen2.5-7B-Instruct - preset=8f3a12c4 nvidia:16GB..24GB:1.. 32K concurrency=1 464 tok/s TTFT 312ms 1 hour ago -``` - -
- -Presets are grouped by base model. Each preset contains an optimized serving configuration for a specific model variant, along with its hardware requirements, validation, and benchmark data. - -Pass `-v` to include validation resources and all benchmark metrics, or `--json` -to output complete preset objects. - -## Apply a preset - -To deploy a preset as a service, pass the endpoint configuration to the `dstack endpoint preset apply` command: - -
- -```shell -$ dstack endpoint preset apply -f endpoint.dstack.yml - Model Qwen/Qwen2.5-7B-Instruct (base) - Preset 8f3a12c4 (context=32K, concurrency=1 464 tok/s TTFT 312ms) - - # BACKEND RESOURCES INSTANCE TYPE PRICE - 1 runpod (CA-MTL-1) cpu=9 mem=50GB disk=200GB NVIDIA RTX A5000 $0.27 - gpu=A5000:24GB:1 - 2 runpod (CA-MTL-1) cpu=9 mem=50GB disk=200GB NVIDIA RTX A5000 $0.27 - gpu=A5000:24GB:1 (spot) - 3 runpod (US-IL-1) cpu=12 mem=25GB disk=200GB NVIDIA RTX A5000 $0.27 - gpu=A5000:24GB:1 - ... - Shown 3 of 4 offers, $0.27max - -Submit the run qwen25-7b? [y/n]: y -``` - -
- -If you don't pass `--preset ID` or specify `preset` in the endpoint configuration, `dstack` automatically selects a matching preset based on the available fleet offers. It then deploys the preset as a service. - -## Delete presets - -You can delete a specific preset by ID or all presets for a base model. - -
- -```shell -$ dstack endpoint preset delete 8f3a12c4 -``` - -
- -To delete all presets for a base model, pass `--model`: - -
- -```shell -$ dstack endpoint preset delete --model Qwen/Qwen2.5-7B-Instruct -``` - -
- -For command options and agent settings, see the -[`dstack endpoint` CLI reference](../reference/cli/dstack/endpoint.md). - -!!! info "Roadmap" - Here's what's coming soon in endpoint presets: - - 1. Support multiple trials, allowing the agent to improve benchmark results based on previous trials. - 2. Allow the endpoint configuration to define custom agent instructions, such as target metrics and the experimentation approach. - 3. Control which experiments the agent may perform, including modifying serving framework code and generating custom kernels. - -!!! info "What's next?" - 1. Learn how dstack [services](services.md) work - 2. Learn how to configure [fleets](fleets.md) diff --git a/mkdocs/docs/concepts/presets.md b/mkdocs/docs/concepts/presets.md new file mode 100644 index 000000000..1ea18073f --- /dev/null +++ b/mkdocs/docs/concepts/presets.md @@ -0,0 +1,203 @@ +--- +title: Presets +description: Creating and reusing optimized model inference configurations +--- + +# Presets + +A preset configuration lets you use an agent to create a preset: a validated and optimized model inference configuration. Once created, the preset can be reused to deploy model inference on validated hardware without an agent. + +The value of presets comes from combining two fundamental features: agent-driven model inference optimization and the `dstack` [service](services.md) primitive, which can deploy model inference to any cloud, Kubernetes, or on-prem cluster. + +> The presets feature is experimental and may change. + +??? info "Prerequisites" + Before using presets, make sure you’ve [installed](../installation.md) the server and CLI, and created a [fleet](fleets.md). + + Creating a preset requires the `claude` CLI to be installed on the machine where you create a preset. + +## Create a preset + +First, define a preset configuration as a YAML file in your project folder. +The filename must end with `.dstack.yml` (e.g. `.dstack.yml` or `preset.dstack.yml` are both acceptable). + +
+ +```yaml +type: preset +name: qwen25-7b + +# The agent picks a compatible variant of the base model +base: Qwen/Qwen2.5-7B-Instruct + +# The number of benchmarked trials +max_trials: 3 + +env: + - HF_TOKEN +``` + +
+ +To create the preset, pass the configuration to the `dstack preset create` command: + +
+ +```shell +$ dstack preset create -f preset.dstack.yml +Create the preset qwen25-7b? [y/n]: y +[2026-07-15 11:32:01] Starting preset creation for Qwen/Qwen2.5-7B-Instruct. Allowed fleets: gpu-fleet. +[2026-07-15 11:41:06] Prototype task qwen25-7b-a1b2c3-2 verified vLLM on an L4:24GB. +[2026-07-15 11:52:06] Final service qwen25-7b-a1b2c3-3 verified with context length 32768. +[2026-07-15 11:52:18] Benchmark via guidellm 0.7.1: 32/32 requests succeeded. +``` + +
+ +The command executes entirely locally and uses the locally installed `claude` CLI along with `dstack`'s bundled skills. The agent uses a `dstack` task to find the best serving configuration for the available fleet offers, then submits it as a `dstack` service for a final benchmark. The validated preset is saved locally under `~/.dstack/presets`. + +You can stop watching with Ctrl+C at any time. The agent keeps running, and `dstack preset logs -f` follows it again. Resume an interrupted creation with `dstack preset create --resume`: + +
+ +```shell +$ dstack preset create -f preset.dstack.yml --resume a1b2c3d4 +``` + +
+ +To stop a creation and its runs, use `dstack preset stop`. + +!!! info "Claude configuration" + By default, preset creation uses the existing `claude` login. To use an Anthropic API key instead, set: + + ```shell + export DSTACK_AGENT_ANTHROPIC_API_KEY=... + ``` + + By default, the agent uses `claude-opus-4-8` and the default `claude` CLI effort. To override them, set: + + ```shell + export DSTACK_AGENT_ANTHROPIC_MODEL=claude-fable-5 + export DSTACK_AGENT_CLAUDE_EFFORT=high + ``` + + Supported effort levels are `low`, `medium`, `high`, `xhigh`, and `max`. + +## Configuration options + +### Model + +=== "Base" + + Set `base` to let the creation agent select any compatible variant of the base model, including a different precision, quantization, or trusted fork. + + ```yaml + base: Qwen/Qwen2.5-7B-Instruct + ``` + +=== "Repo" + + Set `repo` to deploy an exact model. + + ```yaml + repo: Qwen/Qwen2.5-7B-Instruct + ``` + +### Trials + +`max_trials` is required and sets how many benchmarked trials the agent runs before promoting the best one. Set `concurrency` to control the benchmark concurrency. + +### Context length + +Set `context_length` to require a minimum supported context length. + +### Fleets + +Set `fleets` to restrict creation and reuse to specific [fleets](fleets.md). Placement properties such as `backends`, `max_price`, and `spot_policy` constrain both creation and reuse too. + +### Prompt + +Set `prompt` to guide the agent with custom objectives, target metrics, or an experimentation approach. It accepts inline text or a file `path`. + +
+ +```yaml +prompt: | + Optimize for the lowest TTFT at concurrency 32. Consider FP8 quantization. +``` + +
+ +!!! info "Reference" + The `preset` configuration supports many more options. See the [`.dstack.yml` reference](../reference/dstack.yml/preset.md). + +## Apply a preset + +To deploy a preset as a service, pass the preset configuration and the preset ID to the `dstack preset apply` command: + +
+ +```shell +$ dstack preset apply -f preset.dstack.yml --id 532f3f4b + Project main + User admin + Type service + Resources cpu=2.. mem=8GB.. disk=100GB.. gpu=RTXPRO4500:32GB:1.. + Spot policy on-demand + Max price off + Model Qwen/Qwen3.5-27B (base) + Preset 532f3f4b (ctx=8K con=8 387 tok/s TTFT 582ms) + + # BACKEND RESOURCES INSTANCE TYPE PRICE + 1 runpod (EU-RO-1) cpu=12 mem=54GB disk=100GB gpu=RTXPRO4500:32GB:1 NVIDIA RTX PRO 4500 Blackwell $0.74 + +Submit the run qwen35-27b? [y/n]: y +``` + +
+ +## Manage presets + +### List presets + +Use `dstack preset` to list presets: + +
+ +```shell +$ dstack preset list + BASE ID GPU BENCHMARK STATUS SUBMITTED NAME + Qwen/Qwen2.5-0.5B + bc592b38 clauding (0/3) 23 sec ago qwen05 + Qwen/Qwen3-32B + f91d6b60 RTX5090:32GB:1 con=8 576 tok/s TTFT 368ms verified (10/10) 2 days ago qwen3-32b + Qwen/Qwen3.5-27B + 3c4d5e6f verifying (3/3) 2 min ago qwen35-27b-2 + 532f3f4b RTXPRO4500:32GB:1.. con=8 387 tok/s TTFT 582ms verified (4/4) yesterday qwen35-27b + d1c2e12b RTX5090:32GB:1 con=8 266 tok/s TTFT 2.15s verified (7/10) yesterday +``` + +
+ +Presets are grouped by base model. In-progress creations appear too, with a live status like `clauding` or `verifying`. Pass `-w` to watch in realtime. Pass `-v` to include validation resources and all benchmark metrics, or `--json` for complete preset objects. Filter with `--base` or `--repo`. + +### Delete presets + +Delete a preset by ID or name, or all presets for a base model with `--base`: + +
+ +```shell +$ dstack preset delete 8f3a12c4 +``` + +
+ +For command options and agent settings, see the [`dstack preset` CLI reference](../reference/cli/dstack/preset.md). + +> Presets are experimental, and we’d love your feedback. Report bugs and request features on [GitHub](https://github.com/dstackai/dstack/issues), and ask questions on [Discord](https://discord.gg/u8SmfwPpMd). + +!!! info "What's next?" + 1. Learn how dstack [services](services.md) work + 2. Learn how to configure [fleets](fleets.md) diff --git a/mkdocs/docs/reference/cli/dstack/endpoint.md b/mkdocs/docs/reference/cli/dstack/endpoint.md deleted file mode 100644 index 60c7530bb..000000000 --- a/mkdocs/docs/reference/cli/dstack/endpoint.md +++ /dev/null @@ -1,100 +0,0 @@ -# dstack endpoint - -The `dstack endpoint` commands create, list, apply, and delete local -[endpoint presets](../../../concepts/endpoints.md). - -## dstack endpoint preset list - -The `dstack endpoint preset list` command lists locally stored presets. - -##### Usage - -
- -```shell -$ dstack endpoint preset list --help -#GENERATE# -``` - -
- -## dstack endpoint preset create - -The `dstack endpoint preset create` command uses an agent to create and save a -verified preset from an endpoint configuration. - -##### Usage - -
- -```shell -$ dstack endpoint preset create --help -#GENERATE# -``` - -
- -##### Agent settings - -Preset creation uses the existing `claude` login unless -`DSTACK_AGENT_ANTHROPIC_API_KEY` is set. - -| Variable | Description | -| --- | --- | -| `DSTACK_AGENT_ANTHROPIC_API_KEY` | Anthropic API key used by the agent. | -| `DSTACK_AGENT_CLAUDE_PATH` | `claude` executable name or path. Defaults to `claude` from `PATH`. | -| `DSTACK_AGENT_ANTHROPIC_MODEL` | Claude model used by the agent. Defaults to `claude-opus-4-8`. | -| `DSTACK_AGENT_CLAUDE_EFFORT` | Claude effort level: `low`, `medium`, `high`, `xhigh`, or `max`. If unset, the `claude` CLI default is used. | - -Agent progress is written to `agent.log` under -`~/.dstack/agent//-/`. Failed attempts use -the `-failed` suffix. Pass `--debug` to also save the effective endpoint -configuration (`endpoint.dstack.yml`), agent prompt (`prompt.md`), and raw trace -(`trace.jsonl`). - -## dstack endpoint preset get - -The `dstack endpoint preset get` command outputs one locally stored preset as JSON. - -##### Usage - -
- -```shell -$ dstack endpoint preset get --help -#GENERATE# -``` - -
- -## dstack endpoint preset apply - -The `dstack endpoint preset apply` command selects a matching local preset and -submits its service. - -##### Usage - -
- -```shell -$ dstack endpoint preset apply --help -#GENERATE# -``` - -
- -## dstack endpoint preset delete - -The `dstack endpoint preset delete` command deletes one local preset by ID or -all presets for a base model. - -##### Usage - -
- -```shell -$ dstack endpoint preset delete --help -#GENERATE# -``` - -
diff --git a/mkdocs/docs/reference/cli/dstack/preset.md b/mkdocs/docs/reference/cli/dstack/preset.md new file mode 100644 index 000000000..f899383a6 --- /dev/null +++ b/mkdocs/docs/reference/cli/dstack/preset.md @@ -0,0 +1,130 @@ +# dstack preset + +The `dstack preset` commands create, list, apply, and delete local +[presets](../../../concepts/presets.md). + +## dstack preset list + +The `dstack preset list` command lists locally stored presets. + +##### Usage + +
+ +```shell +$ dstack preset list --help +#GENERATE# +``` + +
+ +## dstack preset create + +The `dstack preset create` command uses an agent to create and save a +verified preset from a preset configuration. + +##### Usage + +
+ +```shell +$ dstack preset create --help +#GENERATE# +``` + +
+ +##### Agent settings + +Preset creation uses the existing `claude` login unless +`DSTACK_AGENT_ANTHROPIC_API_KEY` is set. + +| Variable | Description | +| --- | --- | +| `DSTACK_AGENT_ANTHROPIC_API_KEY` | Anthropic API key used by the agent. | +| `DSTACK_AGENT_CLAUDE_PATH` | `claude` executable name or path. Defaults to `claude` from `PATH`. | +| `DSTACK_AGENT_ANTHROPIC_MODEL` | Claude model used by the agent. Defaults to `claude-opus-4-8`. | +| `DSTACK_AGENT_CLAUDE_EFFORT` | Claude effort level: `low`, `medium`, `high`, `xhigh`, or `max`. If unset, the `claude` CLI default is used. | + +Agent progress is written to `agent.log` under `~/.dstack/presets//`, +alongside the effective configuration (`preset.dstack.yml`) and the recorded +trials (`trials.jsonl`). Pass `--debug` to also save the agent prompt +(`prompt.md`) and raw trace (`trace.jsonl`). + +## dstack preset logs + +The `dstack preset logs` command shows a preset creation's log. Pass `-f` to +re-follow a detached or running creation to completion. + +##### Usage + +
+ +```shell +$ dstack preset logs --help +#GENERATE# +``` + +
+ +## dstack preset stop + +The `dstack preset stop` command stops a running preset creation and its runs. + +##### Usage + +
+ +```shell +$ dstack preset stop --help +#GENERATE# +``` + +
+ +## dstack preset get + +The `dstack preset get` command outputs one locally stored preset as JSON. + +##### Usage + +
+ +```shell +$ dstack preset get --help +#GENERATE# +``` + +
+ +## dstack preset apply + +The `dstack preset apply` command selects a matching local preset and +submits its service. + +##### Usage + +
+ +```shell +$ dstack preset apply --help +#GENERATE# +``` + +
+ +## dstack preset delete + +The `dstack preset delete` command deletes one local preset by ID or name, or +all presets for a base model. + +##### Usage + +
+ +```shell +$ dstack preset delete --help +#GENERATE# +``` + +
diff --git a/mkdocs/docs/reference/dstack.yml/endpoint.md b/mkdocs/docs/reference/dstack.yml/preset.md similarity index 78% rename from mkdocs/docs/reference/dstack.yml/endpoint.md rename to mkdocs/docs/reference/dstack.yml/preset.md index 8a3461103..a3ac4176f 100644 --- a/mkdocs/docs/reference/dstack.yml/endpoint.md +++ b/mkdocs/docs/reference/dstack.yml/preset.md @@ -1,11 +1,11 @@ -# `endpoint` +# `preset` -The `endpoint` configuration type describes a model request and the constraints -used to create or apply an [endpoint preset](../../concepts/endpoints.md). +The `preset` configuration type describes a model request and the constraints +used to create or apply a [preset](../../concepts/presets.md). ## Root reference -#SCHEMA# dstack._internal.cli.models.endpoints.EndpointConfiguration +#SCHEMA# dstack._internal.cli.models.configurations.PresetConfiguration overrides: show_root_heading: false type: @@ -17,7 +17,7 @@ used to create or apply an [endpoint preset](../../concepts/endpoints.md). Allows the creation agent to select a compatible model variant. - #SCHEMA# dstack._internal.cli.models.endpoints.EndpointModelBase + #SCHEMA# dstack._internal.cli.models.configurations.PresetModelBase overrides: show_root_heading: false @@ -26,10 +26,18 @@ used to create or apply an [endpoint preset](../../concepts/endpoints.md). Requires an exact model repo or path and optionally sets another client-facing model name. - #SCHEMA# dstack._internal.cli.models.endpoints.EndpointModelRepo + #SCHEMA# dstack._internal.cli.models.configurations.PresetModelRepo overrides: show_root_heading: false +### `prompt` + +Custom agent instructions. Set to an inline string, or to a file: + +#SCHEMA# dstack._internal.cli.models.configurations.PresetPromptFile + overrides: + show_root_heading: false + ### `retry` #SCHEMA# dstack._internal.core.models.profiles.ProfileRetry diff --git a/mkdocs/docs/reference/env.md b/mkdocs/docs/reference/env.md index c547765cd..575fe12c6 100644 --- a/mkdocs/docs/reference/env.md +++ b/mkdocs/docs/reference/env.md @@ -216,7 +216,7 @@ $ find ~/.dstack/logs/cli/ - `DSTACK_PROJECT`{ #DSTACK_PROJECT } – Has the same effect as `--project`. Defaults to `None`. -- `DSTACK_AGENT_ANTHROPIC_API_KEY`{ #DSTACK_AGENT_ANTHROPIC_API_KEY } – The Anthropic API key used by the endpoint preset agent. If unset, the existing `claude` login is used. -- `DSTACK_AGENT_CLAUDE_PATH`{ #DSTACK_AGENT_CLAUDE_PATH } – The `claude` executable name or path used by the endpoint preset agent. Defaults to `claude` from `PATH`. -- `DSTACK_AGENT_ANTHROPIC_MODEL`{ #DSTACK_AGENT_ANTHROPIC_MODEL } – The Claude model used by the endpoint preset agent. Defaults to `claude-opus-4-8`. -- `DSTACK_AGENT_CLAUDE_EFFORT`{ #DSTACK_AGENT_CLAUDE_EFFORT } – The Claude effort level used by the endpoint preset agent. Can be `low`, `medium`, `high`, `xhigh`, or `max`. If unset, the `claude` CLI default is used. +- `DSTACK_AGENT_ANTHROPIC_API_KEY`{ #DSTACK_AGENT_ANTHROPIC_API_KEY } – The Anthropic API key used by the preset agent. If unset, the existing `claude` login is used. +- `DSTACK_AGENT_CLAUDE_PATH`{ #DSTACK_AGENT_CLAUDE_PATH } – The `claude` executable name or path used by the preset agent. Defaults to `claude` from `PATH`. +- `DSTACK_AGENT_ANTHROPIC_MODEL`{ #DSTACK_AGENT_ANTHROPIC_MODEL } – The Claude model used by the preset agent. Defaults to `claude-opus-4-8`. +- `DSTACK_AGENT_CLAUDE_EFFORT`{ #DSTACK_AGENT_CLAUDE_EFFORT } – The Claude effort level used by the preset agent. Can be `low`, `medium`, `high`, `xhigh`, or `max`. If unset, the `claude` CLI default is used. diff --git a/pyproject.toml b/pyproject.toml index 9dc99af54..14f40e082 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,8 +73,8 @@ artifacts = [ ] [tool.hatch.build.targets.wheel.force-include] -"skills/dstack/SKILL.md" = "dstack/_internal/cli/services/endpoints/resources/skills/dstack/SKILL.md" -"skills/dstack-prototyping/SKILL.md" = "dstack/_internal/cli/services/endpoints/resources/skills/dstack-prototyping/SKILL.md" +"skills/dstack/SKILL.md" = "dstack/_internal/cli/services/presets/resources/skills/dstack/SKILL.md" +"skills/dstack-prototyping/SKILL.md" = "dstack/_internal/cli/services/presets/resources/skills/dstack-prototyping/SKILL.md" [tool.hatch.metadata.hooks.fancy-pypi-readme] content-type = "text/markdown" diff --git a/skills/dstack-prototyping/SKILL.md b/skills/dstack-prototyping/SKILL.md index b70bb07a5..c41cf2465 100644 --- a/skills/dstack-prototyping/SKILL.md +++ b/skills/dstack-prototyping/SKILL.md @@ -22,10 +22,10 @@ through the dstack service URL. ## Choose Where To Run -Choose only VM-based backends, SSH fleets, or Kubernetes fleets because they support idle instances and/or instance volumes. That lets later runs reuse the provisioned/idle instance or instance volumes used by runs for caching model weights (and possibly other writes). You must follow this rule even if there are fleets/backends/offers that are cheaper. The only exception from this rule is when the required GPU class (regardless of the price) is not available through VM-based backend, SSH fleet, or Kubernetes fleet. +Pick the offer whose hardware best fits the goal at hand. Only when several offers fit comparably, choose a VM-based backend, an SSH fleet, or a Kubernetes fleet: they support idle instances and/or instance volumes, so later runs reuse the provisioned/idle instance or instance volumes for caching model weights (and possibly other writes), while container-based backends start clean on every run. -Read `https://dstack.ai/docs/concepts/backends.md` to know exactly which -backends are VM-based. +Fetch `https://dstack.ai/docs/concepts/backends.md` and classify backends +from the fetched document, not from memory. ## Check Serving Sources @@ -37,13 +37,13 @@ For vLLM and SGLang, use these as credible sources: - vLLM recipes and model index: `https://recipes.vllm.ai/` and `https://recipes.vllm.ai/models.json` -- vLLM recipe docs: `https://docs.vllm.ai/projects/recipes/en/stable/` -- SGLang docs and cookbook: `https://docs.sglang.ai/` and - `https://lmsysorg.mintlify.app/cookbook/intro` - -Use deeper serving-engine writeups, such as -`https://www.lmsys.org/blog/2026-07-02-agent-assisted-sglang-development`, when -these references do not explain the model, hardware, or serving failure. +- SGLang docs: `https://docs.sglang.io/` (fetch `/llms.txt` for the page + index) +- SGLang model recipes: `https://docs.sglang.io/cookbook/autoregressive/intro` +- Release notes: `https://github.com/vllm-project/vllm/releases` and + `https://github.com/sgl-project/sglang/releases` +- Performance-loop methodology (profiling, benchmark contracts): + `https://www.lmsys.org/blog/2026-07-02-agent-assisted-sglang-development` ## Use A Task Before Service diff --git a/src/dstack/_internal/cli/commands/endpoint.py b/src/dstack/_internal/cli/commands/endpoint.py deleted file mode 100644 index be609af5f..000000000 --- a/src/dstack/_internal/cli/commands/endpoint.py +++ /dev/null @@ -1,266 +0,0 @@ -import argparse -import os -from pathlib import Path - -from argcomplete import FilesCompleter # type: ignore[attr-defined] - -from dstack._internal.cli.commands import BaseCommand -from dstack._internal.cli.models.endpoint_presets import EndpointPresetListOutput -from dstack._internal.cli.models.endpoints import EndpointConfiguration -from dstack._internal.cli.services.completion import ProjectNameCompleter -from dstack._internal.cli.services.endpoints.apply import apply_endpoint_preset -from dstack._internal.cli.services.endpoints.create import create_endpoint_preset -from dstack._internal.cli.services.endpoints.output import print_endpoint_presets -from dstack._internal.cli.services.endpoints.store import ( - EndpointPresetStore, - load_endpoint_configuration, -) -from dstack._internal.cli.services.profile import ( - apply_profile_args, - load_profile_from_args, - register_profile_args, -) -from dstack._internal.cli.utils.common import confirm_ask, console -from dstack._internal.core.errors import CLIError -from dstack._internal.core.models.profiles import ProfileParams -from dstack._internal.core.services import is_valid_dstack_resource_name -from dstack.api import Client - - -class EndpointCommand(BaseCommand): - NAME = "endpoint" - DESCRIPTION = "Manage model inference endpoints" - - def _register(self) -> None: - self._parser.add_argument( - "--project", - help="The name of the project. Defaults to [code]$DSTACK_PROJECT[/]", - metavar="NAME", - default=os.getenv("DSTACK_PROJECT"), - ).completer = ProjectNameCompleter() # type: ignore[attr-defined] - self._parser.set_defaults(subfunc=lambda _: self._parser.print_help()) - subparsers = self._parser.add_subparsers(dest="action") - preset_parser = subparsers.add_parser( - "preset", - help="Manage endpoint presets", - formatter_class=self._parser.formatter_class, - ) - _add_list_args(preset_parser) - preset_parser.set_defaults(subfunc=self._list) - preset_subparsers = preset_parser.add_subparsers(dest="preset_action") - - list_parser = preset_subparsers.add_parser( - "list", - help="List endpoint presets", - formatter_class=self._parser.formatter_class, - ) - _add_list_args(list_parser) - list_parser.set_defaults(subfunc=self._list) - - get_parser = preset_subparsers.add_parser( - "get", - help="Get an endpoint preset", - formatter_class=self._parser.formatter_class, - ) - get_parser.add_argument("preset", metavar="ID", help="The preset ID") - get_parser.add_argument( - "--json", - action="store_true", - required=True, - help="Output in JSON format", - ) - get_parser.set_defaults(subfunc=self._get) - - create_parser = preset_subparsers.add_parser( - "create", - help="Create an endpoint preset", - formatter_class=self._parser.formatter_class, - ) - _add_configuration_args(create_parser) - register_profile_args(create_parser) - create_parser.add_argument( - "--keep-service", - action="store_true", - help="Leave the verified service running", - ) - create_parser.add_argument( - "--debug", - action="store_true", - help="Save the agent prompt and raw trace", - ) - create_parser.set_defaults(subfunc=self._create) - - apply_parser = preset_subparsers.add_parser( - "apply", - help="Apply an endpoint preset", - formatter_class=self._parser.formatter_class, - ) - _add_configuration_args(apply_parser) - register_profile_args(apply_parser) - apply_parser.add_argument("--preset", metavar="ID", help="The preset ID to use") - apply_parser.add_argument( - "-y", "--yes", action="store_true", help="Do not ask for confirmation" - ) - apply_parser.add_argument( - "--force", action="store_true", help="Force apply when no changes are detected" - ) - apply_parser.add_argument( - "-d", "--detach", action="store_true", help="Exit after submitting the service" - ) - apply_parser.add_argument( - "-v", "--verbose", action="store_true", help="Show all plan properties" - ) - apply_parser.set_defaults(subfunc=self._apply) - - delete_parser = preset_subparsers.add_parser( - "delete", - help="Delete endpoint presets", - formatter_class=self._parser.formatter_class, - ) - delete_target = delete_parser.add_mutually_exclusive_group(required=True) - delete_target.add_argument( - "preset", - nargs="?", - metavar="ID", - help="The preset ID", - ) - delete_target.add_argument( - "--model", - metavar="MODEL", - help="Delete all presets for a base model", - ) - delete_parser.add_argument( - "-y", "--yes", action="store_true", help="Do not ask for confirmation" - ) - delete_parser.set_defaults(subfunc=self._delete) - - def _command(self, args: argparse.Namespace) -> None: - super()._command(args) - try: - args.subfunc(args) - except KeyboardInterrupt: - console.print("\nOperation interrupted by user. Exiting...") - exit(0) - - def _list(self, args: argparse.Namespace) -> None: - presets = EndpointPresetStore().list() - if getattr(args, "json", False): - print(EndpointPresetListOutput(presets=presets).json()) - return - print_endpoint_presets(presets, verbose=getattr(args, "verbose", False)) - - def _create(self, args: argparse.Namespace) -> None: - _, configuration = load_endpoint_configuration(args.configuration_file) - configuration = _get_effective_configuration(configuration, args) - result = create_endpoint_preset( - api=Client.from_config(project_name=args.project), - configuration=configuration, - store=EndpointPresetStore(), - keep_service=args.keep_service, - debug=args.debug, - ) - console.print( - f"Endpoint preset [code]{result.preset.id}[/] for " - f"[code]{result.preset.base}[/] saved to [code]{result.path}[/]" - ) - if args.keep_service: - console.print(f"Final service [code]{result.final_run_name}[/] kept running") - - def _get(self, args: argparse.Namespace) -> None: - preset = EndpointPresetStore().get(args.preset) - if preset is None: - raise CLIError(f"Endpoint preset {args.preset!r} does not exist") - print(preset.json()) - - def _apply(self, args: argparse.Namespace) -> None: - configuration_path, configuration = load_endpoint_configuration(args.configuration_file) - configuration = _get_effective_configuration(configuration, args) - apply_endpoint_preset( - api=Client.from_config(project_name=args.project), - configuration=configuration, - configuration_path=configuration_path, - preset_id=args.preset or configuration.preset, - profile_name=args.profile, - command_args=args, - store=EndpointPresetStore(), - ) - - def _delete(self, args: argparse.Namespace) -> None: - store = EndpointPresetStore() - preset = None - if args.preset is not None: - preset = store.get(args.preset) - if preset is None: - raise CLIError(f"Endpoint preset {args.preset!r} does not exist") - message = f"Delete endpoint preset [code]{preset.id}[/] for [code]{preset.base}[/]?" - else: - presets = [preset for preset in store.list() if preset.base == args.model] - if not presets: - raise CLIError(f"No endpoint presets found for base model {args.model!r}") - message = ( - f"Delete {len(presets)} endpoint preset" - f"{'s' if len(presets) != 1 else ''} for [code]{args.model}[/]?" - ) - if not args.yes and not confirm_ask(message): - console.print("\nExiting...") - return - if args.preset is not None: - assert preset is not None - store.delete(args.preset) - console.print( - f"Endpoint preset [code]{preset.id}[/] for [code]{preset.base}[/] deleted" - ) - else: - count = store.delete_for_base(args.model) - console.print( - f"Deleted {count} endpoint preset{'s' if count != 1 else ''} " - f"for [code]{args.model}[/]" - ) - - -def _add_configuration_args(parser: argparse.ArgumentParser) -> None: - parser.add_argument( - "-f", - "--file", - required=True, - metavar="FILE", - dest="configuration_file", - help="The endpoint configuration file", - ).completer = FilesCompleter(allowednames=["*.yml", "*.yaml"]) # type: ignore[attr-defined] - parser.add_argument( - "-n", - "--name", - metavar="NAME", - help="The endpoint name. Required when the configuration omits name", - ) - - -def _add_list_args(parser: argparse.ArgumentParser) -> None: - parser.add_argument("-v", "--verbose", action="store_true") - parser.add_argument( - "--json", - action="store_true", - help="Output in JSON format", - ) - - -def _apply_name(configuration: EndpointConfiguration, name: str | None) -> None: - if name is not None: - configuration.name = name - if configuration.name is None: - raise CLIError("Endpoint name is required. Set `name` in the configuration or use --name") - if not is_valid_dstack_resource_name(configuration.name): - raise CLIError("Endpoint name must match '^[a-z][a-z0-9-]{1,40}$'") - - -def _get_effective_configuration( - configuration: EndpointConfiguration, - args: argparse.Namespace, -) -> EndpointConfiguration: - _apply_name(configuration, args.name) - profile = load_profile_from_args(args=args, repo_dir=Path.cwd()) - for field in ProfileParams.__fields__: - if getattr(configuration, field) is None: - setattr(configuration, field, getattr(profile, field)) - apply_profile_args(args, configuration) - return EndpointConfiguration.parse_obj(configuration.dict()) diff --git a/src/dstack/_internal/cli/commands/preset.py b/src/dstack/_internal/cli/commands/preset.py new file mode 100644 index 000000000..184df2592 --- /dev/null +++ b/src/dstack/_internal/cli/commands/preset.py @@ -0,0 +1,500 @@ +import argparse +import os +import time +from contextlib import suppress +from pathlib import Path + +from argcomplete import FilesCompleter # type: ignore[attr-defined] +from rich.live import Live + +from dstack._internal.cli.commands import BaseCommand +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.models.presets import ( + Preset, + PresetListOutput, +) +from dstack._internal.cli.services.completion import ProjectNameCompleter +from dstack._internal.cli.services.presets.apply import apply_preset +from dstack._internal.cli.services.presets.create import ( + CreationStopped, + create_preset, + find_preset_name_holders, + plan_preset, + reassign_preset_name, + reconcile_detached_sessions, + show_preset_session_logs, + stop_preset_session, +) +from dstack._internal.cli.services.presets.output import get_presets_table, print_presets +from dstack._internal.cli.services.presets.session import ( + list_agent_sessions, + load_resumable_agent_session, + resolve_session_ref, +) +from dstack._internal.cli.services.presets.store import ( + PresetStore, + load_preset_configuration, + resolve_preset_prompt, +) +from dstack._internal.cli.services.profile import ( + apply_profile_args, + load_profile_from_args, + register_profile_args, +) +from dstack._internal.cli.utils.common import ( + LIVE_TABLE_PROVISION_INTERVAL_SECS, + LIVE_TABLE_REFRESH_RATE_PER_SEC, + confirm_ask, + console, + warn, +) +from dstack._internal.core.errors import CLIError, ConfigurationError, ServerClientError +from dstack._internal.core.models.profiles import ProfileParams +from dstack._internal.core.services import validate_dstack_resource_name +from dstack.api import Client + + +class PresetCommand(BaseCommand): + NAME = "preset" + DESCRIPTION = "Manage model serving presets" + + def _register(self) -> None: + self._parser.add_argument( + "--project", + help="The name of the project. Defaults to [code]$DSTACK_PROJECT[/]", + metavar="NAME", + default=os.getenv("DSTACK_PROJECT"), + ).completer = ProjectNameCompleter() # type: ignore[attr-defined] + _add_list_args(self._parser) + self._parser.set_defaults(subfunc=self._list) + preset_subparsers = self._parser.add_subparsers(dest="action") + + list_parser = preset_subparsers.add_parser( + "list", + help="List presets", + formatter_class=self._parser.formatter_class, + ) + _add_list_args(list_parser) + list_parser.set_defaults(subfunc=self._list) + + get_parser = preset_subparsers.add_parser( + "get", + help="Get a preset", + formatter_class=self._parser.formatter_class, + ) + get_parser.add_argument("preset", metavar="ID", help="The preset ID or name") + get_parser.add_argument( + "--json", + action="store_true", + required=True, + help="Output in JSON format", + ) + get_parser.set_defaults(subfunc=self._get) + + create_parser = preset_subparsers.add_parser( + "create", + help="Create a preset", + formatter_class=self._parser.formatter_class, + ) + _add_configuration_args(create_parser) + register_profile_args(create_parser) + create_parser.add_argument( + "--keep-service", + action="store_true", + help="Leave the verified service running", + ) + create_parser.add_argument( + "--max-trials", + type=int, + metavar="N", + help="The maximum number of benchmarked trials before the best one is promoted", + ) + create_parser.add_argument( + "--debug", + action="store_true", + help="Save the agent prompt and raw trace", + ) + create_parser.add_argument( + "--resume", + metavar="ID", + help="Resume an interrupted preset creation by its preset ID", + ) + create_parser.add_argument( + "-y", "--yes", action="store_true", help="Do not ask for confirmation" + ) + create_parser.set_defaults(subfunc=self._create) + + logs_parser = preset_subparsers.add_parser( + "logs", + help="Show a preset's creation log", + formatter_class=self._parser.formatter_class, + ) + logs_parser.add_argument("preset", metavar="ID", help="The preset ID or name") + logs_parser.add_argument( + "-f", + "--follow", + action="store_true", + help="Follow to completion and save the preset", + ) + logs_parser.add_argument( + "--keep-service", + action="store_true", + help="Leave the verified service running (with -f)", + ) + logs_parser.set_defaults(subfunc=self._logs) + + stop_parser = preset_subparsers.add_parser( + "stop", + help="Stop a running preset creation", + formatter_class=self._parser.formatter_class, + ) + stop_parser.add_argument("preset", metavar="ID", help="The preset ID or name") + stop_parser.add_argument( + "-y", "--yes", action="store_true", help="Do not ask for confirmation" + ) + stop_parser.set_defaults(subfunc=self._stop) + + apply_parser = preset_subparsers.add_parser( + "apply", + help="Apply a preset", + formatter_class=self._parser.formatter_class, + ) + _add_configuration_args(apply_parser) + register_profile_args(apply_parser) + apply_parser.add_argument( + "--id", + dest="preset_id", + metavar="ID", + required=True, + help="The preset ID to deploy", + ) + apply_parser.add_argument( + "-y", "--yes", action="store_true", help="Do not ask for confirmation" + ) + apply_parser.add_argument( + "--force", action="store_true", help="Force apply when no changes are detected" + ) + apply_parser.add_argument( + "-d", "--detach", action="store_true", help="Exit after submitting the service" + ) + apply_parser.add_argument( + "-v", "--verbose", action="store_true", help="Show all plan properties" + ) + apply_parser.set_defaults(subfunc=self._apply) + + delete_parser = preset_subparsers.add_parser( + "delete", + help="Delete presets", + formatter_class=self._parser.formatter_class, + ) + delete_target = delete_parser.add_mutually_exclusive_group(required=True) + delete_target.add_argument( + "preset", + nargs="?", + metavar="ID", + help="The preset ID or name", + ) + delete_target.add_argument( + "--base", + metavar="MODEL", + help="Delete all presets for a base model", + ) + delete_target.add_argument( + "--repo", + metavar="REPO", + help="Delete all presets serving a model repo", + ) + delete_parser.add_argument( + "-y", "--yes", action="store_true", help="Do not ask for confirmation" + ) + delete_parser.set_defaults(subfunc=self._delete) + + def _command(self, args: argparse.Namespace) -> None: + super()._command(args) + try: + args.subfunc(args) + except KeyboardInterrupt: + console.print("\nOperation interrupted by user. Exiting...") + exit(0) + + def _reconcile(self) -> None: + """Finalize any detached/orphaned session whose agent already completed, + so a saved preset never depends on a foreground CLI surviving. Fully + best-effort — it must never make a read command fail.""" + with suppress(Exception): + reconcile_detached_sessions(PresetStore()) + + def _list(self, args: argparse.Namespace) -> None: + base = args.base + repo = args.repo + if args.json: + self._reconcile() + presets = _filter_presets(PresetStore().list(), base=base, repo=repo) + print(PresetListOutput(presets=presets).json()) + return + verbose = args.verbose + if not getattr(args, "watch", False): + presets, sessions = self._list_presets_and_sessions(base=base, repo=repo) + print_presets(presets, sessions=sessions, verbose=verbose) + return + with Live(console=console, refresh_per_second=LIVE_TABLE_REFRESH_RATE_PER_SEC) as live: + while True: + presets, sessions = self._list_presets_and_sessions(base=base, repo=repo) + live.update(get_presets_table(presets, sessions=sessions, verbose=verbose)) + time.sleep(LIVE_TABLE_PROVISION_INTERVAL_SECS) + + def _list_presets_and_sessions( + self, *, base: str | None, repo: str | None + ) -> tuple[list[Preset], list[dict]]: + self._reconcile() + presets = PresetStore().list() + sessions = list_agent_sessions() + if base or repo: + repo_to_base = {preset.model: preset.base for preset in presets} + presets = _filter_presets(presets, base=base, repo=repo) + sessions = [ + session + for session in sessions + if _session_matches_model(session, base=base, repo=repo, repo_to_base=repo_to_base) + ] + return presets, sessions + + def _create(self, args: argparse.Namespace) -> None: + configuration_path, configuration = load_preset_configuration(args.configuration_file) + configuration = _get_effective_configuration(configuration, args, require_name=False) + user_prompt = resolve_preset_prompt(configuration, configuration_path) + store = PresetStore() + resume_session = None + if getattr(args, "resume", None): + resume_session = load_resumable_agent_session(args.resume) + if getattr(args, "max_trials", None) is not None: + console.print( + "[warning]--max-trials is ignored when resuming: " + "the constraints are fixed at creation[/]" + ) + api = Client.from_config(project_name=args.project) + allowed_fleets = None + if resume_session is None: + if configuration.max_trials is None: + raise ConfigurationError( + "max_trials is required. Set it in the configuration or pass --max-trials" + ) + allowed_fleets = plan_preset(api=api, configuration=configuration) + if not _confirm_preset_creation(store, configuration.name, assume_yes=args.yes): + console.print("\nExiting...") + return + try: + result = create_preset( + api=api, + configuration=configuration, + store=store, + keep_service=args.keep_service, + debug=args.debug, + resume_session=resume_session, + user_prompt=user_prompt, + allowed_fleets=allowed_fleets, + ) + except KeyboardInterrupt: + return # the interrupt handler already reported detach / stop + except CreationStopped: + return # stopped from another CLI, which reported the interruption + # The log already told the story; like a finished run, success is silent. + if args.keep_service: + console.print(f"Final service [code]{result.final_run_name}[/] kept running") + + def _logs(self, args: argparse.Namespace) -> None: + try: + result = show_preset_session_logs( + project=args.project, + store=PresetStore(), + preset_id=resolve_session_ref(args.preset), + follow=args.follow, + keep_service=args.keep_service, + ) + except KeyboardInterrupt: + return # a log viewer: Ctrl+C just stops watching, quietly + except CreationStopped: + return # stopped from another CLI, which reported the interruption + if result is not None and args.keep_service: + console.print(f"Final service [code]{result.final_run_name}[/] kept running") + + def _stop(self, args: argparse.Namespace) -> None: + preset_id = resolve_session_ref(args.preset) + if not args.yes and not confirm_ask(f"Stop creating preset [code]{preset_id}[/]?"): + console.print("\nExiting...") + return + stop_preset_session(Client.from_config(project_name=args.project), preset_id) + + def _get(self, args: argparse.Namespace) -> None: + self._reconcile() + preset = PresetStore().find_by_id_or_name(args.preset) + if preset is None: + raise CLIError(f"Preset {args.preset!r} does not exist") + print(preset.json()) + + def _apply(self, args: argparse.Namespace) -> None: + self._reconcile() + configuration_path, configuration = load_preset_configuration(args.configuration_file) + configuration = _get_effective_configuration(configuration, args) + apply_preset( + api=Client.from_config(project_name=args.project), + configuration=configuration, + configuration_path=configuration_path, + preset_id=args.preset_id, + profile_name=args.profile, + command_args=args, + store=PresetStore(), + ) + + def _delete(self, args: argparse.Namespace) -> None: + store = PresetStore() + if args.preset is not None: + try: + preset = store.find_by_id_or_name(args.preset) + except CLIError as e: + # A corrupt preset file must still be removable by ID. + warn(str(e), stderr=True) + preset_ids = [args.preset] + description = f"preset [code]{args.preset}[/]" + else: + if preset is None: + raise CLIError(f"Preset {args.preset!r} does not exist") + preset_ids = [preset.id] + description = f"preset [code]{preset.id}[/] for [code]{preset.base}[/]" + else: + target = args.base or args.repo + presets = _filter_presets(store.list(), base=args.base, repo=args.repo) + if not presets: + kind = "base model" if args.base else "model repo" + raise CLIError(f"No presets found for {kind} {target!r}") + preset_ids = [preset.id for preset in presets] + count = f"{len(presets)} preset{'s' if len(presets) != 1 else ''}" + description = f"{count} for [code]{target}[/]" + if not args.yes and not confirm_ask(f"Delete {description}?"): + console.print("\nExiting...") + return + for preset_id in preset_ids: + store.delete(preset_id) + console.print(f"Deleted {description}") + + +def _add_configuration_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "-f", + "--file", + required=True, + metavar="FILE", + dest="configuration_file", + help="The preset configuration file", + ).completer = FilesCompleter(allowednames=["*.yml", "*.yaml"]) # type: ignore[attr-defined] + parser.add_argument( + "-n", + "--name", + metavar="NAME", + help="The service name. Required when the configuration omits name", + ) + + +def _add_list_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument("-v", "--verbose", action="store_true") + parser.add_argument( + "-w", + "--watch", + action="store_true", + help="Watch presets in realtime", + ) + parser.add_argument( + "--json", + action="store_true", + help="Output in JSON format", + ) + model_filter = parser.add_mutually_exclusive_group() + model_filter.add_argument( + "--base", + metavar="MODEL", + help="Show only presets for a base model", + ) + model_filter.add_argument( + "--repo", + metavar="REPO", + help="Show only presets serving a model repo", + ) + + +def _filter_presets( + presets: list[Preset], + *, + base: str | None, + repo: str | None, +) -> list[Preset]: + return [ + preset + for preset in presets + if (base is None or preset.base == base) and (repo is None or preset.model == repo) + ] + + +def _session_matches_model( + session: dict, + *, + base: str | None, + repo: str | None, + repo_to_base: dict[str, str], +) -> bool: + model = str(session.get("model") or "") + if not model: + return False + if repo is not None: + return model == repo + return model == base or repo_to_base.get(model) == base + + +def _apply_name(configuration: PresetConfiguration, name: str | None, *, required: bool) -> None: + if name is not None: + configuration.name = name + if configuration.name is None: + if required: + raise CLIError( + "The service name is required. Set `name` in the configuration or use --name" + ) + return + try: + validate_dstack_resource_name(configuration.name) + except ServerClientError as e: + raise CLIError(str(e)) from e + + +def _confirm_preset_creation(store: PresetStore, name: str | None, *, assume_yes: bool) -> bool: + """One apply-style confirmation; reassigns the name from any holder on yes.""" + holders = find_preset_name_holders(store, name) if name is not None else None + if holders is not None and holders.preset_ids: + used_by = ", ".join(f"preset [code]{preset_id}[/]" for preset_id in holders.preset_ids) + message = ( + f"The name [code]{name}[/] is already used by {used_by}. Reassign it to a new preset?" + ) + elif name is not None: + message = f"Create the preset [code]{name}[/]?" + else: + message = "Create the preset?" + if not assume_yes and not confirm_ask(message): + return False + if holders is not None: + reassign_preset_name(store, holders) + return True + + +def _get_effective_configuration( + configuration: PresetConfiguration, + args: argparse.Namespace, + *, + require_name: bool = True, +) -> PresetConfiguration: + _apply_name(configuration, args.name, required=require_name) + if getattr(args, "max_trials", None) is not None: + configuration.max_trials = args.max_trials + profile = load_profile_from_args(args=args, repo_dir=Path.cwd()) + for field in ProfileParams.__fields__: + if getattr(configuration, field) is None: + setattr(configuration, field, getattr(profile, field)) + apply_profile_args(args, configuration) + return PresetConfiguration.parse_obj(configuration.dict()) diff --git a/src/dstack/_internal/cli/main.py b/src/dstack/_internal/cli/main.py index 335f7693f..98e3b70f5 100644 --- a/src/dstack/_internal/cli/main.py +++ b/src/dstack/_internal/cli/main.py @@ -8,7 +8,6 @@ from dstack._internal.cli.commands.attach import AttachCommand from dstack._internal.cli.commands.completion import CompletionCommand from dstack._internal.cli.commands.delete import DeleteCommand -from dstack._internal.cli.commands.endpoint import EndpointCommand from dstack._internal.cli.commands.event import EventCommand from dstack._internal.cli.commands.export import ExportCommand from dstack._internal.cli.commands.fleet import FleetCommand @@ -19,6 +18,7 @@ from dstack._internal.cli.commands.logs import LogsCommand from dstack._internal.cli.commands.metrics import MetricsCommand from dstack._internal.cli.commands.offer import OfferCommand +from dstack._internal.cli.commands.preset import PresetCommand from dstack._internal.cli.commands.project import ProjectCommand from dstack._internal.cli.commands.ps import PsCommand from dstack._internal.cli.commands.run import RunCommand @@ -68,7 +68,6 @@ def main(): ApplyCommand.register(subparsers) AttachCommand.register(subparsers) DeleteCommand.register(subparsers) - EndpointCommand.register(subparsers) EventCommand.register(subparsers) ExportCommand.register(subparsers) FleetCommand.register(subparsers) @@ -79,6 +78,7 @@ def main(): LoginCommand.register(subparsers) LogsCommand.register(subparsers) MetricsCommand.register(subparsers) + PresetCommand.register(subparsers) ProjectCommand.register(subparsers) PsCommand.register(subparsers) RunCommand.register(subparsers) diff --git a/src/dstack/_internal/cli/models/configurations.py b/src/dstack/_internal/cli/models/configurations.py new file mode 100644 index 000000000..b0e08942b --- /dev/null +++ b/src/dstack/_internal/cli/models/configurations.py @@ -0,0 +1,224 @@ +from typing import Annotated, Any, Literal, Optional, Union + +from pydantic import Field, PositiveInt, root_validator, validator + +from dstack._internal.core.models.common import ( + CoreModel, + EntityReference, + generate_dual_core_model, +) +from dstack._internal.core.models.envs import Env +from dstack._internal.core.models.profiles import ProfileParams, ProfileParamsConfig +from dstack._internal.utils.json_schema import add_extra_schema_types + +DEFAULT_CONCURRENCY = 8 + + +class PresetModelRepo(CoreModel): + repo: Annotated[str, Field(description="The exact model repo or path to deploy")] + name: Annotated[ + Optional[str], Field(description="The client-facing model name. Defaults to `repo`") + ] = None + + @property + def api_model_name(self) -> str: + return self.name or self.repo + + @property + def exact_repo(self) -> str: + return self.repo + + @property + def allows_variant_selection(self) -> bool: + return False + + @validator("repo") + def validate_repo(cls, value: str) -> str: + return _validate_model(value, field="repo") + + @validator("name") + def validate_name(cls, value: Optional[str]) -> Optional[str]: + if value is None: + return None + return _validate_model(value, field="name") + + +class PresetModelBase(CoreModel): + base: Annotated[ + str, + Field(description="The base model for which the agent may select a compatible variant"), + ] + + @property + def api_model_name(self) -> str: + return self.base + + @property + def exact_repo(self) -> None: + return None + + @property + def allows_variant_selection(self) -> bool: + return True + + @validator("base") + def validate_base(cls, value: str) -> str: + return _validate_model(value, field="base") + + +PresetModelSpec = Union[PresetModelRepo, PresetModelBase] + +MAX_PROMPT_LENGTH = 10_000 + + +class PresetPromptFile(CoreModel): + path: Annotated[ + str, + Field(description="The path to a prompt file, relative to the configuration file"), + ] + + @validator("path") + def validate_path(cls, value: str) -> str: + if not value.strip(): + raise ValueError("Prompt path must be a non-empty string") + return value + + +class PresetConfigurationConfig(ProfileParamsConfig): + @staticmethod + def schema_extra(schema: dict[str, Any]): + ProfileParamsConfig.schema_extra(schema) + add_extra_schema_types( + schema["properties"]["model"], + extra_types=[{"type": "string"}], + ) + + +class PresetConfiguration( + ProfileParams, + generate_dual_core_model(PresetConfigurationConfig), +): + type: Annotated[Literal["preset"], Field(description="The configuration type")] = "preset" + name: Annotated[ + Optional[str], + Field(description="The service name. Required unless passed with `--name`"), + ] = None + model: Annotated[ + PresetModelSpec, + Field( + description=( + "The model to serve. Use a string or `repo` for an exact repo/path, " + "or `base` to allow compatible model variants. " + "Prefer the top-level `base`/`repo` shorthand unless a custom " + "client-facing model name is needed" + ) + ), + ] + base: Annotated[ + Optional[str], + Field( + description=( + "The base model repo; compatible variants are allowed. Shorthand for `model.base`" + ) + ), + ] = None + repo: Annotated[ + Optional[str], + Field(description="The exact model repo/path to serve. Shorthand for `model.repo`"), + ] = None + prompt: Annotated[ + Optional[Union[str, PresetPromptFile]], + Field( + description=( + "Additional instructions for the preset creation agent, inline or as a file `path`" + ) + ), + ] = None + context_length: Annotated[ + Optional[PositiveInt], Field(description="The minimum required context length") + ] = None + max_trials: Annotated[ + Optional[PositiveInt], + Field( + description=( + "The maximum number of benchmarked trials during preset creation" + " before the best one is promoted" + ) + ), + ] = None + concurrency: Annotated[ + Optional[PositiveInt], + Field( + description=( + "The number of simultaneous requests used for benchmarks during" + f" preset creation. Defaults to `{DEFAULT_CONCURRENCY}`" + ) + ), + ] = None + gateway: Annotated[ + Optional[Union[bool, EntityReference, str]], + Field( + description=( + "The name of the gateway. Specify boolean `false` to run without a gateway." + " Specify boolean `true` to run with the default gateway." + " Omit to run with the default gateway if there is one, or without a gateway otherwise" + ) + ), + ] = None + env: Annotated[Env, Field(description="The mapping or the list of environment variables")] = ( + Env() + ) + + @property + def effective_concurrency(self) -> int: + return self.concurrency if self.concurrency is not None else DEFAULT_CONCURRENCY + + @root_validator(pre=True) + def apply_model_shorthand(cls, values: Any) -> Any: + if not isinstance(values, dict): + return values + base, repo = values.get("base"), values.get("repo") + if base and repo: + raise ValueError("`base` and `repo` are mutually exclusive") + if base or repo: + if values.get("model") is not None: + raise ValueError("`model` cannot be combined with the `base`/`repo` shorthand") + values = dict(values) + values.pop("base", None) + values.pop("repo", None) + values["model"] = {"base": base} if base else {"repo": repo} + return values + + @validator("model", pre=True) + def parse_model(cls, value: Any) -> Any: + if isinstance(value, str): + return {"repo": _validate_model(value, field="model")} + return value + + @validator("prompt") + def validate_prompt(cls, value: Any) -> Any: + if isinstance(value, str): + if not value.strip(): + raise ValueError("Prompt must be a non-empty string") + if len(value) > MAX_PROMPT_LENGTH: + raise ValueError(f"Prompt must be at most {MAX_PROMPT_LENGTH} characters") + return value + + +class PresetConstraints(CoreModel): + """The effective constraints for preset creation, saved as `constraints.json` + in the agent workspace. Field semantics are documented in the agent system prompt.""" + + run_name_prefix: str + model: PresetModelSpec + context_length: Optional[PositiveInt] = None + max_trials: PositiveInt + concurrency: PositiveInt + fleets: list[str] = Field(min_items=1) + env: list[str] = [] + + +def _validate_model(value: Any, *, field: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"Preset model {field} must be a non-empty string") + return value diff --git a/src/dstack/_internal/cli/models/endpoints.py b/src/dstack/_internal/cli/models/endpoints.py deleted file mode 100644 index 8b2cb409e..000000000 --- a/src/dstack/_internal/cli/models/endpoints.py +++ /dev/null @@ -1,134 +0,0 @@ -from typing import Annotated, Any, Literal, Optional, Union - -from pydantic import Field, PositiveInt, validator - -from dstack._internal.core.models.common import ( - CoreModel, - EntityReference, - generate_dual_core_model, -) -from dstack._internal.core.models.envs import Env -from dstack._internal.core.models.profiles import ProfileParams, ProfileParamsConfig -from dstack._internal.utils.json_schema import add_extra_schema_types - - -class EndpointModelRepo(CoreModel): - repo: Annotated[str, Field(description="The exact model repo or path to deploy")] - name: Annotated[ - Optional[str], Field(description="The client-facing model name. Defaults to `repo`") - ] = None - - @property - def api_model_name(self) -> str: - return self.name or self.repo - - @property - def exact_repo(self) -> str: - return self.repo - - @property - def allows_variant_selection(self) -> bool: - return False - - @validator("repo") - def validate_repo(cls, value: str) -> str: - return _validate_model(value, field="repo") - - @validator("name") - def validate_name(cls, value: Optional[str]) -> Optional[str]: - if value is None: - return None - return _validate_model(value, field="name") - - -class EndpointModelBase(CoreModel): - base: Annotated[ - str, - Field(description="The base model for which the agent may select a compatible variant"), - ] - - @property - def api_model_name(self) -> str: - return self.base - - @property - def exact_repo(self) -> None: - return None - - @property - def allows_variant_selection(self) -> bool: - return True - - @validator("base") - def validate_base(cls, value: str) -> str: - return _validate_model(value, field="base") - - -EndpointModelSpec = Union[EndpointModelRepo, EndpointModelBase] - - -class EndpointConfigurationConfig(ProfileParamsConfig): - @staticmethod - def schema_extra(schema: dict[str, Any]): - ProfileParamsConfig.schema_extra(schema) - add_extra_schema_types( - schema["properties"]["model"], - extra_types=[{"type": "string"}], - ) - - -class EndpointConfiguration( - ProfileParams, - generate_dual_core_model(EndpointConfigurationConfig), -): - type: Annotated[Literal["endpoint"], Field(description="The configuration type")] = "endpoint" - name: Annotated[ - Optional[str], - Field(description="The endpoint name. Required unless passed with `--name`"), - ] = None - model: Annotated[ - EndpointModelSpec, - Field( - description=( - "The model to serve. Use a string or `repo` for an exact repo/path, " - "or `base` to allow compatible model variants." - ) - ), - ] - context_length: Annotated[ - Optional[PositiveInt], Field(description="The minimum required context length") - ] = None - preset: Annotated[ - Optional[str], Field(description="The preset ID to use when applying the endpoint") - ] = None - gateway: Annotated[ - Optional[Union[bool, EntityReference, str]], - Field( - description=( - "The name of the gateway. Specify boolean `false` to run without a gateway." - " Specify boolean `true` to run with the default gateway." - " Omit to run with the default gateway if there is one, or without a gateway otherwise" - ) - ), - ] = None - env: Annotated[Env, Field(description="The mapping or the list of environment variables")] = ( - Env() - ) - - @validator("model", pre=True) - def parse_model(cls, value: Any) -> Any: - if isinstance(value, str): - return {"repo": _validate_model(value, field="model")} - return value - - @validator("preset") - def validate_preset(cls, value: Optional[str]) -> Optional[str]: - if value is not None and not value.strip(): - raise ValueError("Endpoint preset must be a non-empty string") - return value - - -def _validate_model(value: Any, *, field: str) -> str: - if not isinstance(value, str) or not value.strip(): - raise ValueError(f"Endpoint model {field} must be a non-empty string") - return value diff --git a/src/dstack/_internal/cli/models/endpoint_agent.py b/src/dstack/_internal/cli/models/preset_agent.py similarity index 86% rename from src/dstack/_internal/cli/models/endpoint_agent.py rename to src/dstack/_internal/cli/models/preset_agent.py index c31a1c854..e75694124 100644 --- a/src/dstack/_internal/cli/models/endpoint_agent.py +++ b/src/dstack/_internal/cli/models/preset_agent.py @@ -1,9 +1,9 @@ import uuid -from typing import Optional +from typing import Any, Dict, Optional from pydantic import PositiveInt, root_validator -from dstack._internal.cli.models.endpoint_presets import EndpointBenchmark +from dstack._internal.cli.models.presets import PresetBenchmark from dstack._internal.core.models.common import CoreModel _LATENCY_JSON_SCHEMA = { @@ -97,7 +97,7 @@ class AgentFinalReport(CoreModel): base: Optional[str] = None model: Optional[str] = None context_length: Optional[PositiveInt] = None - benchmark: Optional[EndpointBenchmark] = None + benchmark: Optional[PresetBenchmark] = None failure_summary: Optional[str] = None @root_validator @@ -118,3 +118,23 @@ def validate_report(cls, values: dict) -> dict: elif not values.get("failure_summary"): raise ValueError("failed agent report must include failure_summary") return values + + +class PresetAgentInfo(CoreModel): + """Base information about the agent runtime that ran a preset creation + session, saved in the debug session directory.""" + + executable: str + version: Optional[str] = None + + +class ClaudeModelParams(CoreModel): + name: str + effort: str + + +class ClaudeAgentInfo(PresetAgentInfo): + """Claude agent runtime information, saved as `agent.json`.""" + + model: ClaudeModelParams + auth: Dict[str, Any] diff --git a/src/dstack/_internal/cli/models/endpoint_presets.py b/src/dstack/_internal/cli/models/presets.py similarity index 82% rename from src/dstack/_internal/cli/models/endpoint_presets.py rename to src/dstack/_internal/cli/models/presets.py index aea33f499..16866ecc3 100644 --- a/src/dstack/_internal/cli/models/endpoint_presets.py +++ b/src/dstack/_internal/cli/models/presets.py @@ -17,7 +17,7 @@ from dstack._internal.core.models.resources import CPUSpec, ResourcesSpec -class EndpointBenchmarkWorkload(CoreModel): +class PresetBenchmarkWorkload(CoreModel): api: Literal["chat_completions", "completions"] num_requests: PositiveInt input_tokens: PositiveInt @@ -25,38 +25,38 @@ class EndpointBenchmarkWorkload(CoreModel): concurrency: PositiveInt -class EndpointBenchmarkLatency(CoreModel): +class PresetBenchmarkLatency(CoreModel): mean: Annotated[float, Field(ge=0)] p50: Annotated[float, Field(ge=0)] p99: Annotated[float, Field(ge=0)] -class EndpointBenchmarkMetrics(CoreModel): +class PresetBenchmarkMetrics(CoreModel): successful_requests: Annotated[int, Field(ge=0)] failed_requests: Annotated[int, Field(ge=0)] duration_seconds: PositiveFloat total_input_tokens: Annotated[int, Field(ge=0)] total_output_tokens: Annotated[int, Field(ge=0)] - ttft_ms: EndpointBenchmarkLatency - tpot_ms: EndpointBenchmarkLatency + ttft_ms: PresetBenchmarkLatency + tpot_ms: PresetBenchmarkLatency -class EndpointBenchmarkTarget(CoreModel): +class PresetBenchmarkTarget(CoreModel): type: Literal["gateway", "server-proxy"] -class EndpointBenchmarkClient(CoreModel): +class PresetBenchmarkClient(CoreModel): type: Literal["local"] -class EndpointBenchmark(CoreModel): +class PresetBenchmark(CoreModel): tool: str tool_version: str command: str - workload: EndpointBenchmarkWorkload - metrics: EndpointBenchmarkMetrics - target: Optional[EndpointBenchmarkTarget] = None - client: Optional[EndpointBenchmarkClient] = None + workload: PresetBenchmarkWorkload + metrics: PresetBenchmarkMetrics + target: Optional[PresetBenchmarkTarget] = None + client: Optional[PresetBenchmarkClient] = None @validator("tool", "tool_version", "command") def validate_non_empty(cls, value: str) -> str: @@ -70,6 +70,10 @@ def validate_command_has_no_bearer_token(cls, value: str) -> str: token = match.group(1) if token.startswith("$") or "redacted" in token.lower() or set(token) == {"*"}: continue + # Prose such as "auth via bearer header from env" is not a + # credential: only credential-shaped values are rejected. + if len(token) < 16 or not any(char.isdigit() for char in token): + continue raise ValueError("command must not contain a bearer token value") return value @@ -77,6 +81,7 @@ def validate_command_has_no_bearer_token(cls, value: str) -> str: def validate_metrics(cls, values: dict) -> dict: metrics = values.get("metrics") workload = values.get("workload") + assert metrics is not None and workload is not None if metrics.failed_requests != 0: raise ValueError("benchmark must not include failed requests") if metrics.successful_requests != workload.num_requests: @@ -84,28 +89,30 @@ def validate_metrics(cls, values: dict) -> dict: return values -class EndpointPresetValidationReplica(CoreModel): +class PresetValidationReplica(CoreModel): resources: list[ResourcesSpec] """Exact resources for each running replica in this service replica group.""" -class EndpointPresetValidation(CoreModel): - replicas: list[EndpointPresetValidationReplica] +class PresetValidation(CoreModel): + replicas: list[PresetValidationReplica] """Ordered to match `ServiceConfiguration.replica_groups`.""" - benchmark: EndpointBenchmark + benchmark: PresetBenchmark -class EndpointPreset(CoreModel): +class Preset(CoreModel): base: str """Base model used for local preset lookup.""" id: str + name: Optional[str] = None + """Mutable human name; at most one preset or in-flight session holds it.""" model: str """Exact repo/path loaded by the service command.""" context_length: PositiveInt """Token context length this preset was verified to support.""" created_at: datetime service: ServiceConfiguration - validations: list[EndpointPresetValidation] + validations: list[PresetValidation] @validator("base", "id", "model") def validate_non_empty(cls, value: str) -> str: @@ -144,8 +151,8 @@ def validate_preset(cls, values: dict) -> dict: return values -class EndpointPresetListOutput(CoreModel): - presets: list[EndpointPreset] +class PresetListOutput(CoreModel): + presets: list[Preset] def _validate_exact_resources(resources: ResourcesSpec) -> None: diff --git a/src/dstack/_internal/cli/services/endpoints/agent.py b/src/dstack/_internal/cli/services/endpoints/agent.py deleted file mode 100644 index 546907b96..000000000 --- a/src/dstack/_internal/cli/services/endpoints/agent.py +++ /dev/null @@ -1,743 +0,0 @@ -import asyncio -import json -import os -import shutil -import signal -import subprocess -import sys -import tempfile -from contextlib import contextmanager, suppress -from dataclasses import dataclass, field -from datetime import datetime, timezone -from pathlib import Path -from typing import Any, Iterator, Optional, Sequence - -import psutil -import yaml -from rich.text import Text - -from dstack._internal.cli.models.endpoint_agent import AGENT_FINAL_REPORT_JSON_SCHEMA -from dstack._internal.cli.models.endpoints import EndpointConfiguration -from dstack._internal.cli.utils.common import console -from dstack._internal.compat import IS_WINDOWS -from dstack._internal.core.errors import CLIError -from dstack._internal.core.services.configs import ConfigManager -from dstack._internal.utils.common import get_dstack_dir -from dstack.api import Client - -_SKILL_NAMES = ("dstack", "dstack-prototyping") -_PROGRESS_FILENAME = "progress.jsonl" -_SUBMISSIONS_FILENAME = "submissions.jsonl" -_FINAL_REPORT_FILENAME = "final_report.json" -_PROGRESS_ENV = "DSTACK_ENDPOINT_PROGRESS_LOG" -_REDACTION = "[redacted]" -_CLAUDE_TOOLS = "Bash,Read,Write,Edit,WebFetch,WebSearch,StructuredOutput" -_CLAUDE_EFFORT_LEVELS = ("low", "medium", "high", "xhigh", "max") -_CLAUDE_STREAM_LIMIT = 16 * 1024 * 1024 -_MAX_RUN_NAME_LENGTH = 41 -_MAX_UNIX_SOCKET_PATH_BYTES = 103 -_INHERITED_ENV_NAMES = ( - "PATH", - "HOME", - "USER", - "SSL_CERT_FILE", - "REQUESTS_CA_BUNDLE", - "HTTP_PROXY", - "HTTPS_PROXY", - "NO_PROXY", - "http_proxy", - "https_proxy", - "no_proxy", -) -_WINDOWS_INHERITED_ENV_NAMES = ( - "APPDATA", - "COMSPEC", - "HOMEDRIVE", - "HOMEPATH", - "LOCALAPPDATA", - "PATHEXT", - "PROGRAMDATA", - "SYSTEMDRIVE", - "SYSTEMROOT", - "USERPROFILE", - "USERNAME", - "WINDIR", -) -_SENSITIVE_INHERITED_ENV_NAMES = ( - "HTTP_PROXY", - "HTTPS_PROXY", - "NO_PROXY", - "http_proxy", - "https_proxy", - "no_proxy", -) - - -@dataclass(frozen=True) -class EndpointAgentWorkspace: - path: Path - dstack_home: Path - - @property - def temp_path(self) -> Path: - return self.path / ".tmp" - - @property - def bin_path(self) -> Path: - return self.path / "bin" - - @property - def progress_path(self) -> Path: - return self.path / _PROGRESS_FILENAME - - @property - def submissions_path(self) -> Path: - return self.path / _SUBMISSIONS_FILENAME - - @property - def final_report_path(self) -> Path: - return self.path / _FINAL_REPORT_FILENAME - - -@dataclass -class EndpointAgentSession: - path: Path - timestamp: str - debug: bool - _log_enabled: bool = field(default=True, init=False, repr=False) - - @property - def log_path(self) -> Path: - return self.path / "agent.log" - - @property - def trace_path(self) -> Path: - return self.path / "trace.jsonl" - - def write_prompt(self, prompt: str) -> None: - _write_private_text(self.path / "prompt.md", prompt + "\n") - - def append_log(self, line: str) -> None: - if not self._log_enabled: - return - try: - with self.log_path.open("a", encoding="utf-8") as f: - f.write(line + "\n") - f.flush() - except OSError as e: - self._log_enabled = False - console.print(f"[warning]Could not write agent log {self.log_path}: {e}[/]") - - def finish(self, preset_id: Optional[str] = None) -> Path: - suffix = preset_id or "failed" - name = f"{self.timestamp}-{suffix}" - index = 0 - while True: - target = self.path.with_name(name if index == 0 else f"{name}-{index}") - if target.exists(): - index += 1 - continue - try: - self.path.rename(target) - except FileExistsError: - index += 1 - continue - self.path = target - return target - - -@dataclass(frozen=True) -class ClaudeAuth: - api_key: Optional[str] - executable: str - effort: Optional[str] - model: str - - -@dataclass -class EndpointAgentProcessOutput: - report_data: Optional[dict[str, Any]] = None - error: Optional[str] = None - - -@contextmanager -def endpoint_agent_workspace() -> Iterator[EndpointAgentWorkspace]: - with tempfile.TemporaryDirectory(prefix="dpe-", dir=_get_short_temp_dir()) as directory: - root = Path(directory) - _validate_control_socket_path(root) - workspace = EndpointAgentWorkspace(path=root / "w", dstack_home=root / "h") - _prepare_workspace(workspace) - yield workspace - - -def create_endpoint_agent_session( - configuration: EndpointConfiguration, - *, - debug: bool = False, -) -> EndpointAgentSession: - if configuration.name is None: - raise CLIError("Endpoint name is required to save agent output") - timestamp = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S-%fZ") - parent = get_dstack_dir() / "agent" / configuration.name - path: Optional[Path] = None - try: - parent.mkdir(mode=0o700, parents=True, exist_ok=True) - path = _create_agent_session_directory(parent, timestamp) - _write_private_text(path / "agent.log", "") - if debug: - data = json.loads(configuration.json(exclude_none=True)) - if configuration.env: - data["env"] = list(configuration.env) - else: - data.pop("env", None) - _write_private_text( - path / "endpoint.dstack.yml", - yaml.safe_dump(data, sort_keys=False), - ) - _write_private_text(path / "trace.jsonl", "") - except OSError as e: - if path is not None: - shutil.rmtree(path, ignore_errors=True) - raise CLIError(f"Could not create agent output under {parent}: {e}") from e - assert path is not None - return EndpointAgentSession(path=path, timestamp=timestamp, debug=debug) - - -def _create_agent_session_directory(parent: Path, timestamp: str) -> Path: - index = 0 - while True: - name = f"{timestamp}-running" if index == 0 else f"{timestamp}-{index}-running" - path = parent / name - try: - path.mkdir(mode=0o700) - except FileExistsError: - index += 1 - continue - return path - - -def get_claude_auth() -> ClaudeAuth: - api_key = os.getenv("DSTACK_AGENT_ANTHROPIC_API_KEY") or None - configured_path = os.getenv("DSTACK_AGENT_CLAUDE_PATH") or "claude" - executable = shutil.which(configured_path) - if executable is None: - raise CLIError(f"Claude executable not found: {configured_path}") - effort = os.getenv("DSTACK_AGENT_CLAUDE_EFFORT") or None - if effort is not None and effort not in _CLAUDE_EFFORT_LEVELS: - raise CLIError( - f"DSTACK_AGENT_CLAUDE_EFFORT must be one of: {', '.join(_CLAUDE_EFFORT_LEVELS)}" - ) - return ClaudeAuth( - api_key=api_key, - executable=executable, - effort=effort, - model=os.getenv("DSTACK_AGENT_ANTHROPIC_MODEL", "claude-opus-4-8"), - ) - - -def build_endpoint_agent_env( - *, - api: Client, - endpoint_env: dict[str, str], - auth: ClaudeAuth, - workspace: EndpointAgentWorkspace, - token: str, -) -> dict[str, str]: - config_manager = ConfigManager(workspace.dstack_home / ".dstack") - config_manager.configure_project( - name=api.project, - url=api.client.base_url, - token=token, - default=True, - ) - config_manager.save() - env = {name: value for name in _INHERITED_ENV_NAMES if (value := os.getenv(name))} - env.update(endpoint_env) - if IS_WINDOWS: - env.update( - {name: value for name in _WINDOWS_INHERITED_ENV_NAMES if (value := os.getenv(name))} - ) - env["PATH"] = os.pathsep.join([str(workspace.bin_path), env.get("PATH", "")]) - env["DSTACK_SERVER_URL"] = api.client.base_url - env["DSTACK_PROJECT"] = api.project - env["DSTACK_TOKEN"] = token - env["DSTACK_ENDPOINT_SERVER_URL"] = api.client.base_url - env["DSTACK_ENDPOINT_BEARER_TOKEN"] = token - env[_PROGRESS_ENV] = str(workspace.progress_path) - for name in ["TMPDIR", "TEMP", "TMP"]: - env[name] = str(workspace.temp_path) - if auth.api_key is not None: - env["ANTHROPIC_API_KEY"] = auth.api_key - env["HOME"] = str(workspace.dstack_home) - if IS_WINDOWS: - env["USERPROFILE"] = str(workspace.dstack_home) - else: - env["HOME"] = str(Path.home()) - if IS_WINDOWS: - env["USERPROFILE"] = str(Path.home()) - return env - - -async def run_endpoint_agent( - *, - prompt: str, - env: dict[str, str], - workspace: EndpointAgentWorkspace, - auth: ClaudeAuth, - redacted_values: Sequence[str], - agent_session: EndpointAgentSession, -) -> EndpointAgentProcessOutput: - command = _prepare_subprocess_command(_build_claude_command(auth=auth)) - progress_tailer = _ProgressTailer( - path=workspace.progress_path, - redacted_values=redacted_values, - agent_session=agent_session, - ) - progress_task = asyncio.create_task(progress_tailer.run()) - proc: Optional[asyncio.subprocess.Process] = None - try: - proc = await asyncio.create_subprocess_exec( - *command, - cwd=workspace.path, - env=env, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - start_new_session=not IS_WINDOWS, - limit=_CLAUDE_STREAM_LIMIT, - ) - assert proc.stdin is not None - assert proc.stdout is not None - assert proc.stderr is not None - proc.stdin.write(prompt.encode()) - with suppress(BrokenPipeError, ConnectionResetError): - await proc.stdin.drain() - proc.stdin.close() - stdout_task = asyncio.create_task( - _read_process_stream( - stream=proc.stdout, - stream_name="stdout", - parse_result=True, - redacted_values=redacted_values, - agent_session=agent_session, - ) - ) - stderr_task = asyncio.create_task( - _read_process_stream( - stream=proc.stderr, - stream_name="stderr", - parse_result=False, - redacted_values=redacted_values, - agent_session=agent_session, - ) - ) - stdout_output, stderr_output, returncode = await asyncio.gather( - stdout_task, - stderr_task, - proc.wait(), - ) - except BaseException: - if proc is not None and proc.returncode is None: - await _terminate_process(proc) - raise - finally: - progress_task.cancel() - with suppress(asyncio.CancelledError): - await progress_task - progress_tailer.flush() - - output = stdout_output - if output.report_data is None: - output.report_data = stderr_output.report_data - if output.error is None: - output.error = stderr_output.error - if output.report_data is None and returncode != 0: - output.error = output.error or f"Claude exited with return code {returncode}" - return output - - -def print_endpoint_progress(message: str, *, agent_session: EndpointAgentSession) -> None: - timestamp = datetime.now(timezone.utc).astimezone().strftime("%Y-%m-%d %H:%M:%S") - message = message.rstrip("\r\n") - agent_session.append_log(f"[{timestamp}] {message}") - console.print( - Text(f"[{timestamp}]", style="log.time"), - Text(message, style="log.message"), - soft_wrap=True, - ) - - -def get_redacted_values(values: Sequence[str]) -> tuple[str, ...]: - return tuple(sorted({value for value in values if value}, key=len, reverse=True)) - - -def contains_redacted_value(value: Any, redacted_values: Sequence[str]) -> bool: - if isinstance(value, str): - return any( - value == redacted or (len(redacted) >= 8 and redacted in value) - for redacted in redacted_values - ) - if isinstance(value, dict): - return any(contains_redacted_value(item, redacted_values) for item in value.values()) - if isinstance(value, list): - return any(contains_redacted_value(item, redacted_values) for item in value) - return False - - -def get_sensitive_inherited_env_values() -> list[str]: - return [value for name in _SENSITIVE_INHERITED_ENV_NAMES if (value := os.getenv(name))] - - -def redact(value: str, redacted_values: Sequence[str]) -> str: - for redacted_value in redacted_values: - if value == redacted_value: - return _REDACTION - # Replacing short values such as "1" or "false" corrupts unrelated diagnostics. - if len(redacted_value) >= 8: - value = value.replace(redacted_value, _REDACTION) - return value - - -def _validate_control_socket_path(build_root: Path) -> None: - if IS_WINDOWS: - return - path = build_root / "h" / ".dstack" / "ssh" / f"{'x' * _MAX_RUN_NAME_LENGTH}.control.sock" - if len(os.fsencode(path)) > _MAX_UNIX_SOCKET_PATH_BYTES: - raise CLIError(f"Temporary path is too long for an SSH control socket: {build_root}") - - -def _prepare_workspace(workspace: EndpointAgentWorkspace) -> None: - workspace.path.mkdir(mode=0o700, parents=True, exist_ok=False) - workspace.dstack_home.mkdir(mode=0o700) - workspace.temp_path.mkdir(mode=0o700) - for path in [ - workspace.progress_path, - workspace.submissions_path, - ]: - path.touch() - workspace.bin_path.mkdir() - _install_python_command(workspace.bin_path, "progress", _get_progress_script()) - (workspace.dstack_home / ".ssh").mkdir(mode=0o700) - _install_dstack_wrapper(workspace.bin_path, workspace.dstack_home) - _install_home_wrapper(workspace.bin_path, "ssh", workspace.dstack_home) - _install_skills(workspace.path) - - -def _install_dstack_wrapper(bin_dir: Path, home: Path) -> None: - script = f"""#!{sys.executable} -import os -import sys - -from dstack._internal.cli.main import main - -os.environ["HOME"] = {json.dumps(str(home))} -if os.name == "nt": - os.environ["USERPROFILE"] = {json.dumps(str(home))} -sys.argv[0] = "dstack" -raise SystemExit(main()) -""" - _install_python_command(bin_dir, "dstack", script) - - -def _install_home_wrapper(bin_dir: Path, command: str, home: Path) -> None: - executable = shutil.which(command) - if executable is None: - script = f"""#!{sys.executable} -import sys - -print("Endpoint preset creation could not find the {command} executable.", file=sys.stderr) -raise SystemExit(127) -""" - else: - script = f"""#!{sys.executable} -import os -import subprocess -import sys - -os.environ["HOME"] = {json.dumps(str(home))} -if os.name == "nt": - os.environ["USERPROFILE"] = {json.dumps(str(home))} -raise SystemExit(subprocess.call([{json.dumps(executable)}, *sys.argv[1:]])) -""" - _install_python_command(bin_dir, command, script) - - -def _install_python_command(bin_dir: Path, name: str, script: str) -> None: - if not IS_WINDOWS: - path = bin_dir / name - path.write_text(script, encoding="utf-8") - path.chmod(0o755) - return - - # Avoid shadowing packages with the same name when Python adds this directory to sys.path. - script_path = bin_dir / f"_{name}.py" - script_path.write_text(script, encoding="utf-8") - command = subprocess.list2cmdline([sys.executable, str(script_path)]) - (bin_dir / f"{name}.cmd").write_text( - f"@echo off\n{command} %*\n", - encoding="utf-8", - ) - - -def _get_short_temp_dir() -> Optional[str]: - # Keep SSH control sockets below the Unix-domain socket path limit on macOS. - if IS_WINDOWS: - return None - return "/tmp" if Path("/tmp").is_dir() else None - - -def _get_progress_script() -> str: - return f"""#!{sys.executable} -import json -import os -from pathlib import Path -import sys - -message = " ".join(sys.argv[1:]).strip() -if not message and not sys.stdin.isatty(): - message = sys.stdin.read().strip() -if not message: - print("Usage: progress ", file=sys.stderr) - raise SystemExit(2) -path = Path(os.environ.get("{_PROGRESS_ENV}", "{_PROGRESS_FILENAME}")) -with path.open("a", encoding="utf-8") as f: - f.write(json.dumps({{"message": message}}, ensure_ascii=False) + "\\n") -""" - - -def _install_skills(workspace: Path) -> None: - source_dir = _get_skills_dir() - target_dir = workspace / ".claude" / "skills" - target_dir.mkdir(parents=True) - for skill_name in _SKILL_NAMES: - source = source_dir / skill_name - if not (source / "SKILL.md").is_file(): - raise CLIError(f"Missing endpoint agent skill: {skill_name}") - shutil.copytree(source, target_dir / skill_name) - - -def _get_skills_dir() -> Path: - source_path = Path(__file__).resolve() - candidates = ( - source_path.parent / "resources" / "skills", - source_path.parents[6] / "skills", - ) - for candidate in candidates: - if (candidate / "dstack" / "SKILL.md").is_file(): - return candidate - raise CLIError("Could not find packaged dstack skills") - - -def _build_claude_command(*, auth: ClaudeAuth) -> list[str]: - command = [ - auth.executable, - "-p", - "--output-format", - "stream-json", - "--verbose", - "--tools", - _CLAUDE_TOOLS, - "--allowedTools", - _CLAUDE_TOOLS, - "--disallowedTools", - "Task,NotebookEdit", - "--permission-mode", - "bypassPermissions", - "--model", - auth.model, - "--json-schema", - json.dumps(AGENT_FINAL_REPORT_JSON_SCHEMA), - ] - if auth.api_key is None: - command[2:2] = ["--setting-sources", "project,local"] - else: - command[2:2] = ["--bare"] - if auth.effort is not None: - command[2:2] = ["--effort", auth.effort] - return command - - -def _prepare_subprocess_command(command: list[str]) -> list[str]: - if not IS_WINDOWS or Path(command[0]).suffix.lower() not in {".bat", ".cmd"}: - return command - comspec = os.getenv("COMSPEC") or shutil.which("cmd.exe") - if comspec is None: - raise CLIError("Cannot run the Claude batch launcher because cmd.exe was not found") - return [comspec, "/d", "/s", "/c", subprocess.list2cmdline(command)] - - -def _write_private_text(path: Path, content: str) -> None: - path.write_text(content, encoding="utf-8") - if not IS_WINDOWS: - path.chmod(0o600) - - -def _write_debug_trace( - session: EndpointAgentSession, - *, - stream_name: str, - text: str, - redacted_values: Sequence[str], -) -> None: - timestamp = ( - datetime.now(timezone.utc).isoformat(timespec="microseconds").replace("+00:00", "Z") - ) - try: - event = _redact_trace_value(json.loads(text), redacted_values) - record = {"timestamp": timestamp, "stream": stream_name, "event": event} - except json.JSONDecodeError: - record = { - "timestamp": timestamp, - "stream": stream_name, - "text": redact(text.rstrip("\r\n"), redacted_values), - } - line = json.dumps(record, ensure_ascii=False, separators=(",", ":")) - with session.trace_path.open("a", encoding="utf-8") as f: - f.write(line + "\n") - f.flush() - - -def _redact_trace_value(value: Any, redacted_values: Sequence[str]) -> Any: - if isinstance(value, str): - return redact(value, redacted_values) - if isinstance(value, list): - return [_redact_trace_value(item, redacted_values) for item in value] - if isinstance(value, dict): - return { - redact(key, redacted_values): _redact_trace_value(item, redacted_values) - for key, item in value.items() - } - return value - - -async def _read_process_stream( - *, - stream: asyncio.StreamReader, - stream_name: str, - parse_result: bool, - redacted_values: Sequence[str], - agent_session: EndpointAgentSession, -) -> EndpointAgentProcessOutput: - output = EndpointAgentProcessOutput() - while True: - line = await stream.readline() - if not line: - return output - text = line.decode(errors="replace") - if agent_session.debug: - _write_debug_trace( - agent_session, - stream_name=stream_name, - text=text, - redacted_values=redacted_values, - ) - if not parse_result: - continue - try: - message = json.loads(text) - except json.JSONDecodeError: - continue - if not isinstance(message, dict) or message.get("type") != "result": - continue - if message.get("is_error"): - error = message.get("result") or "Claude failed" - output.error = redact(str(error), redacted_values) - structured_output = message.get("structured_output") - if isinstance(structured_output, dict): - output.report_data = structured_output - continue - result = message.get("result") - if isinstance(result, str): - try: - parsed = json.loads(result) - except json.JSONDecodeError: - continue - if isinstance(parsed, dict): - output.report_data = parsed - - -async def _terminate_process(proc: asyncio.subprocess.Process) -> None: - if IS_WINDOWS: - await asyncio.to_thread(_terminate_windows_process_tree, proc.pid) - await proc.wait() - return - # The Windows branch returns above; Pyright still checks these POSIX-only APIs on Windows. - if hasattr(os, "killpg"): - with suppress(ProcessLookupError): - os.killpg(proc.pid, signal.SIGTERM) # pyright: ignore[reportAttributeAccessIssue] - else: - proc.terminate() - try: - await asyncio.wait_for(proc.wait(), timeout=3) - except asyncio.TimeoutError: - if hasattr(os, "killpg"): - with suppress(ProcessLookupError): - os.killpg( # pyright: ignore[reportAttributeAccessIssue] - proc.pid, - signal.SIGKILL, # pyright: ignore[reportAttributeAccessIssue] - ) - else: - proc.kill() - await proc.wait() - - -def _terminate_windows_process_tree(pid: int) -> None: - try: - root = psutil.Process(pid) - except psutil.NoSuchProcess: - return - processes = [*root.children(recursive=True), root] - for process in processes: - with suppress(psutil.NoSuchProcess): - process.terminate() - _, alive = psutil.wait_procs(processes, timeout=3) - for process in alive: - with suppress(psutil.NoSuchProcess): - process.kill() - psutil.wait_procs(alive, timeout=3) - - -class _ProgressTailer: - def __init__( - self, - *, - path: Path, - redacted_values: Sequence[str], - agent_session: EndpointAgentSession, - ) -> None: - self._path = path - self._redacted_values = redacted_values - self._agent_session = agent_session - self._offset = 0 - - async def run(self) -> None: - while True: - self.flush() - await asyncio.sleep(1) - - def flush(self) -> None: - if not self._path.exists(): - return - with self._path.open("r", encoding="utf-8", errors="replace") as f: - f.seek(self._offset) - lines = f.readlines() - self._offset = f.tell() - for line in lines: - message = _parse_progress(line) - if message is not None: - print_endpoint_progress( - redact(message, self._redacted_values), - agent_session=self._agent_session, - ) - - -def _parse_progress(line: str) -> Optional[str]: - try: - value = json.loads(line) - except json.JSONDecodeError: - return line.strip() or None - if isinstance(value, str): - return value.strip() or None - if isinstance(value, dict) and isinstance(value.get("message"), str): - return value["message"].strip() or None - return None diff --git a/src/dstack/_internal/cli/services/endpoints/apply.py b/src/dstack/_internal/cli/services/endpoints/apply.py deleted file mode 100644 index 350a47dd9..000000000 --- a/src/dstack/_internal/cli/services/endpoints/apply.py +++ /dev/null @@ -1,157 +0,0 @@ -import argparse -from dataclasses import dataclass -from typing import Optional - -from rich.markup import escape - -from dstack._internal.cli.models.endpoint_presets import EndpointPreset -from dstack._internal.cli.models.endpoints import EndpointConfiguration -from dstack._internal.cli.services.configurators.run import ServiceConfigurator -from dstack._internal.cli.services.endpoints.output import ( - format_endpoint_benchmark, - format_endpoint_context_length, -) -from dstack._internal.cli.services.endpoints.store import EndpointPresetStore -from dstack._internal.core.errors import CLIError -from dstack._internal.core.models.configurations import ServiceConfiguration -from dstack._internal.core.models.profiles import ProfileParams -from dstack._internal.core.models.repos.base import Repo -from dstack._internal.core.models.runs import RunPlan -from dstack.api import Client - - -@dataclass(frozen=True) -class _PresetPlan: - preset: EndpointPreset - run_plan: RunPlan - repo: Repo - - -def apply_endpoint_preset( - *, - api: Client, - configuration: EndpointConfiguration, - configuration_path: str, - preset_id: Optional[str], - profile_name: Optional[str], - command_args: argparse.Namespace, - store: EndpointPresetStore, -) -> None: - presets = _get_matching_presets( - store.list(), - configuration=configuration, - preset_id=preset_id, - ) - if not presets: - qualifier = f" preset {preset_id!r}" if preset_id else "" - raise CLIError( - f"No matching endpoint preset{qualifier} for {configuration.model.api_model_name}" - ) - - configurator = ServiceConfigurator(api_client=api) - service_args = configurator.get_parser().parse_args([]) - service_args.profile = profile_name - selected = _select_plan( - configuration=configuration, - configuration_path=configuration_path, - presets=presets, - configurator=configurator, - service_args=service_args, - ) - configurator.apply_plan( - run_plan=selected.run_plan, - repo=selected.repo, - command_args=command_args, - configurator_args=service_args, - plan_properties={ - "Model": _format_requested_model(configuration), - "Preset": _format_selected_preset(selected.preset), - }, - ) - - -def _get_matching_presets( - presets: list[EndpointPreset], - *, - configuration: EndpointConfiguration, - preset_id: Optional[str], -) -> list[EndpointPreset]: - model_name = configuration.model.api_model_name - matches = [] - for preset in presets: - service_model = preset.service.model - if preset_id is not None and preset.id != preset_id: - continue - if service_model is None or service_model.name.lower() != model_name.lower(): - continue - if configuration.context_length is not None: - if preset.context_length < configuration.context_length: - continue - if configuration.model.allows_variant_selection: - if preset.base.lower() != model_name.lower(): - continue - elif preset.model != configuration.model.exact_repo: - continue - matches.append(preset) - return matches - - -def _select_plan( - *, - configuration: EndpointConfiguration, - configuration_path: str, - presets: list[EndpointPreset], - configurator: ServiceConfigurator, - service_args: argparse.Namespace, -) -> _PresetPlan: - first_plan: Optional[_PresetPlan] = None - for preset in presets: - service = _build_service(configuration, preset) - run_plan, repo = configurator.get_plan( - conf=service, - configuration_path=configuration_path, - configurator_args=service_args, - ) - plan = _PresetPlan(preset=preset, run_plan=run_plan, repo=repo) - if first_plan is None: - first_plan = plan - if _has_available_offers(run_plan): - return plan - assert first_plan is not None - return first_plan - - -def _build_service( - configuration: EndpointConfiguration, - preset: EndpointPreset, -) -> ServiceConfiguration: - service = preset.service.copy(deep=True) - service.name = configuration.name - service.gateway = configuration.gateway - service.env.update(configuration.env) - for field in ProfileParams.__fields__: - value = getattr(configuration, field) - if value is not None: - setattr(service, field, value) - return service - - -def _has_available_offers(plan: RunPlan) -> bool: - return bool(plan.job_plans) and all( - any(offer.availability.is_available() for offer in job_plan.offers) - for job_plan in plan.job_plans - ) - - -def _format_requested_model(configuration: EndpointConfiguration) -> str: - model = escape(configuration.model.api_model_name) - if configuration.model.allows_variant_selection: - return f"{model} ([secondary]base[/])" - return model - - -def _format_selected_preset(preset: EndpointPreset) -> str: - details = ( - f"context={format_endpoint_context_length(preset)}, {format_endpoint_benchmark(preset)}" - ) - return f"{escape(preset.id)} ([secondary]{details}[/])" diff --git a/src/dstack/_internal/cli/services/endpoints/create.py b/src/dstack/_internal/cli/services/endpoints/create.py deleted file mode 100644 index 9c1ef179e..000000000 --- a/src/dstack/_internal/cli/services/endpoints/create.py +++ /dev/null @@ -1,336 +0,0 @@ -import asyncio -import json -import os -import secrets -import uuid -from contextlib import suppress -from dataclasses import dataclass -from pathlib import Path -from typing import Optional, Sequence - -from dstack._internal.cli.models.endpoint_agent import AgentFinalReport -from dstack._internal.cli.models.endpoint_presets import EndpointPreset -from dstack._internal.cli.models.endpoints import EndpointConfiguration -from dstack._internal.cli.services.endpoints.agent import ( - EndpointAgentSession, - EndpointAgentWorkspace, - build_endpoint_agent_env, - contains_redacted_value, - create_endpoint_agent_session, - endpoint_agent_workspace, - get_claude_auth, - get_redacted_values, - get_sensitive_inherited_env_values, - print_endpoint_progress, - run_endpoint_agent, -) -from dstack._internal.cli.services.endpoints.presets import endpoint_preset_to_data -from dstack._internal.cli.services.endpoints.prompt import ( - format_endpoint_constraints, - get_endpoint_agent_system_prompt, -) -from dstack._internal.cli.services.endpoints.store import EndpointPresetStore -from dstack._internal.cli.services.endpoints.verify import ( - build_verified_endpoint_preset, - load_endpoint_agent_report, -) -from dstack._internal.cli.utils.common import console, warn -from dstack._internal.core.errors import CLIError, ConfigurationError -from dstack._internal.core.models.envs import EnvSentinel -from dstack._internal.core.models.fleets import FleetStatus -from dstack.api import Client - -_RUN_STOP_TIMEOUT_SECONDS = 10 * 60 - - -@dataclass(frozen=True) -class EndpointPresetCreateResult: - preset: EndpointPreset - path: Path - final_run_id: uuid.UUID - final_run_name: str - - -def create_endpoint_preset( - *, - api: Client, - configuration: EndpointConfiguration, - store: EndpointPresetStore, - keep_service: bool = False, - build_name: Optional[str] = None, - debug: bool = False, -) -> EndpointPresetCreateResult: - agent_session = create_endpoint_agent_session(configuration, debug=debug) - try: - resolved_configuration = _resolve_endpoint_env(configuration) - result = asyncio.run( - _create_endpoint_preset( - api=api, - configuration=resolved_configuration, - source_configuration=configuration, - store=store, - keep_service=keep_service, - build_name=build_name, - agent_session=agent_session, - ) - ) - except BaseException: - _finish_agent_session(agent_session) - raise - _finish_agent_session(agent_session, result.preset.id) - return result - - -async def _create_endpoint_preset( - *, - api: Client, - configuration: EndpointConfiguration, - store: EndpointPresetStore, - source_configuration: Optional[EndpointConfiguration] = None, - keep_service: bool = False, - build_name: Optional[str] = None, - agent_session: EndpointAgentSession, -) -> EndpointPresetCreateResult: - source_configuration = source_configuration or configuration - build_name = build_name or _get_build_name(configuration.name) - allowed_fleets = _get_allowed_fleets(api, configuration) - if not allowed_fleets: - raise CLIError("The project has no active fleets available for preset creation") - auth = get_claude_auth() - - endpoint_env = configuration.env.as_dict() - token = getattr(api.client, "_token", None) - if not isinstance(token, str) or not token: - raise CLIError("The configured dstack client has no authentication token") - redacted_values = get_redacted_values( - [ - token, - auth.api_key or "", - *endpoint_env.values(), - *get_sensitive_inherited_env_values(), - ] - ) - report: Optional[AgentFinalReport] = None - preset: Optional[EndpointPreset] = None - preset_path: Optional[Path] = None - creation_succeeded = False - cleanup_error: Optional[str] = None - with endpoint_agent_workspace() as workspace: - env = build_endpoint_agent_env( - api=api, - endpoint_env=endpoint_env, - auth=auth, - workspace=workspace, - token=token, - ) - prompt = _build_prompt( - configuration=configuration, - build_name=build_name, - allowed_fleets=allowed_fleets, - ) - if agent_session.debug: - agent_session.write_prompt(prompt) - print_endpoint_progress( - f"Starting endpoint preset creation for {configuration.model.api_model_name}. " - f"Allowed fleets: {', '.join(allowed_fleets)}.", - agent_session=agent_session, - ) - try: - process_output = await run_endpoint_agent( - prompt=prompt, - env=env, - workspace=workspace, - auth=auth, - redacted_values=redacted_values, - agent_session=agent_session, - ) - report = load_endpoint_agent_report( - output=process_output, - workspace=workspace, - redacted_values=redacted_values, - ) - run = api.client.runs.get(api.project, report.run_name) - preset = build_verified_endpoint_preset( - run=run, - endpoint_configuration=source_configuration, - report=report, - ) - if contains_redacted_value(endpoint_preset_to_data(preset), redacted_values): - raise CLIError("Generated endpoint preset contains a secret value") - preset_path = store.save(preset) - print_endpoint_progress( - f"Saved endpoint preset {preset.id} for {preset.base} at {preset_path}.", - agent_session=agent_session, - ) - creation_succeeded = True - finally: - keep_final_service = keep_service and creation_succeeded - try: - await _cleanup_runs( - api=api, - build_name=build_name, - workspace=workspace, - final_run_name=report.run_name if report is not None else None, - keep_final_service=keep_final_service, - agent_session=agent_session, - ) - except Exception as e: - cleanup_error = str(e) - if keep_final_service: - with suppress(Exception): - await _cleanup_runs( - api=api, - build_name=build_name, - workspace=workspace, - final_run_name=report.run_name if report is not None else None, - agent_session=agent_session, - ) - - if cleanup_error is not None: - raise CLIError(f"Failed to clean up preset creation runs: {cleanup_error}") - assert preset is not None - assert preset_path is not None - assert report is not None - assert report.run_id is not None - assert report.run_name is not None - return EndpointPresetCreateResult( - preset=preset, - path=preset_path, - final_run_id=report.run_id, - final_run_name=report.run_name, - ) - - -def _resolve_endpoint_env(configuration: EndpointConfiguration) -> EndpointConfiguration: - configuration = configuration.copy(deep=True) - for key, value in configuration.env.items(): - if not isinstance(value, EnvSentinel): - continue - try: - configuration.env[key] = value.from_env(os.environ) - except ValueError as e: - raise ConfigurationError(str(e)) from e - return configuration - - -def _finish_agent_session( - session: EndpointAgentSession, - preset_id: Optional[str] = None, -) -> None: - try: - path = session.finish(preset_id) - except OSError as e: - path = session.path - warn(f"Could not finalize agent output. Files remain at {path}: {e}") - console.print(f"Agent log saved to [code]{path / 'agent.log'}[/]") - - -def _get_build_name(endpoint_name: Optional[str]) -> str: - if endpoint_name is None: - raise CLIError("Endpoint name is required. Set `name` in the configuration or use --name") - suffix = secrets.token_hex(3) - # Leave room for the numeric submission suffix while retaining a recognizable prefix. - prefix = endpoint_name[:28].rstrip("-") - return f"{prefix}-{suffix}" - - -def _get_allowed_fleets(api: Client, configuration: EndpointConfiguration) -> tuple[str, ...]: - if configuration.fleets is not None: - return tuple( - fleet.format() if hasattr(fleet, "format") else str(fleet) - for fleet in configuration.fleets - ) - fleets = api.client.fleets.list(api.project, include_imported=True) - return tuple( - fleet.name if fleet.project_name == api.project else f"{fleet.project_name}/{fleet.name}" - for fleet in fleets - if fleet.status == FleetStatus.ACTIVE - ) - - -def _build_prompt( - *, - configuration: EndpointConfiguration, - build_name: str, - allowed_fleets: Sequence[str], -) -> str: - context_lines = [f"- service_model_name: {configuration.model.api_model_name}"] - if configuration.model.allows_variant_selection: - context_lines.append(f"- base_model: {configuration.model.api_model_name}") - else: - context_lines.append(f"- model_repo: {configuration.model.exact_repo}") - if configuration.context_length is not None: - context_lines.append(f"- context_length: {configuration.context_length}") - return f"""{get_endpoint_agent_system_prompt()} - -Endpoint context: -- endpoint_name: {build_name} -{chr(10).join(context_lines)} - -{ - format_endpoint_constraints( - configuration, - configuration.env.as_dict(), - allowed_fleets=allowed_fleets, - ) - } -""" - - -async def _cleanup_runs( - *, - api: Client, - build_name: str, - workspace: EndpointAgentWorkspace, - final_run_name: Optional[str], - agent_session: EndpointAgentSession, - keep_final_service: bool = False, -) -> None: - run_names = _load_submitted_run_names(workspace.submissions_path) - if final_run_name is not None: - run_names.append(final_run_name) - run_names = list(dict.fromkeys(run_names)) - expected_prefix = f"{build_name}-" - run_names = [name for name in run_names if name.startswith(expected_prefix)] - if keep_final_service: - run_names = [name for name in run_names if name != final_run_name] - active_names = [] - for name in run_names: - run = api.runs.get(name) - if run is not None and not run.status.is_finished(): - active_names.append(name) - if not active_names: - return - print_endpoint_progress( - f"Stopping preset creation runs: {', '.join(active_names)}.", - agent_session=agent_session, - ) - api.client.runs.stop(api.project, active_names, abort=False) - deadline = asyncio.get_running_loop().time() + _RUN_STOP_TIMEOUT_SECONDS - pending = set(active_names) - while pending: - if asyncio.get_running_loop().time() >= deadline: - raise CLIError(f"Timed out waiting for runs to stop: {', '.join(sorted(pending))}") - for name in list(pending): - run = api.runs.get(name) - if run is None or run.status.is_finished(): - pending.remove(name) - if pending: - await asyncio.sleep(2) - print_endpoint_progress("All preset creation runs stopped.", agent_session=agent_session) - - -def _load_submitted_run_names(path: Path) -> list[str]: - if not path.exists(): - return [] - names = [] - for line in path.read_text(encoding="utf-8", errors="replace").splitlines(): - try: - value = json.loads(line) - except json.JSONDecodeError: - continue - if isinstance(value, dict) and isinstance(value.get("name"), str): - name = value["name"].strip() - if name: - names.append(name) - return names diff --git a/src/dstack/_internal/cli/services/endpoints/output.py b/src/dstack/_internal/cli/services/endpoints/output.py deleted file mode 100644 index 3d11c02e8..000000000 --- a/src/dstack/_internal/cli/services/endpoints/output.py +++ /dev/null @@ -1,150 +0,0 @@ -from collections import defaultdict - -from rich.table import Table - -from dstack._internal.cli.models.endpoint_presets import ( - EndpointPreset, - EndpointPresetValidation, -) -from dstack._internal.cli.utils.common import add_row_from_dict, console -from dstack._internal.utils.common import pretty_date, pretty_resources - - -def print_endpoint_presets(presets: list[EndpointPreset], verbose: bool = False) -> None: - table = Table(box=None) - table.add_column("MODEL", no_wrap=True) - table.add_column("RESOURCES" if verbose else "GPU") - table.add_column("CONTEXT", justify="right") - table.add_column("BENCHMARK", min_width=len("concurrency=1"), overflow="fold") - table.add_column("CREATED", no_wrap=True) - presets_by_base: dict[str, list[EndpointPreset]] = defaultdict(list) - for preset in presets: - presets_by_base[preset.base].append(preset) - - for base, base_presets in presets_by_base.items(): - add_row_from_dict(table, {"MODEL": f"[bold]{base}[/]"}) - for preset in base_presets: - _add_preset(table, preset, verbose=verbose) - console.print(table) - console.print() - - -def _add_preset(table: Table, preset: EndpointPreset, *, verbose: bool) -> None: - groups = preset.service.replica_groups - column = "RESOURCES" if verbose else "GPU" - add_row_from_dict( - table, - { - "MODEL": f"[secondary] preset={preset.id}[/]", - column: _format_resources(groups[0].resources, verbose=verbose), - "CONTEXT": format_endpoint_context_length(preset), - "BENCHMARK": format_endpoint_benchmark(preset, verbose=verbose), - "CREATED": pretty_date(preset.created_at), - }, - ) - if preset.model != preset.base: - add_row_from_dict( - table, - {"MODEL": f" repo={preset.model}"}, - style="secondary", - ) - if len(groups) > 1: - for group in groups: - add_row_from_dict( - table, - { - "MODEL": f" group={group.name}", - column: _format_resources(group.resources, verbose=verbose), - }, - style="secondary", - ) - - -def format_endpoint_context_length(preset: EndpointPreset) -> str: - return _format_token_count(preset.context_length) - - -def format_endpoint_benchmark(preset: EndpointPreset, *, verbose: bool = False) -> str: - validation = preset.validations[0] - benchmark = validation.benchmark - workload = benchmark.workload - metrics = benchmark.metrics - requests_per_second = metrics.successful_requests / metrics.duration_seconds - output_tokens_per_second = metrics.total_output_tokens / metrics.duration_seconds - parts = [ - f"concurrency={workload.concurrency}", - f"{_format_number(output_tokens_per_second)} tok/s", - f"TTFT {_format_latency(metrics.ttft_ms.p50)}", - ] - if verbose: - ttft = _format_latency_summary( - metrics.ttft_ms.mean, metrics.ttft_ms.p50, metrics.ttft_ms.p99 - ) - tpot = _format_latency_summary( - metrics.tpot_ms.mean, metrics.tpot_ms.p50, metrics.tpot_ms.p99 - ) - parts.extend( - [ - f"hardware={_format_validation_gpus(validation)}", - f"api={workload.api}", - f"n={workload.num_requests}", - f"{_format_token_count(workload.input_tokens)}" - f"->{_format_token_count(workload.output_tokens)}", - f"{_format_number(requests_per_second)} req/s", - f"duration={_format_number(metrics.duration_seconds)}s", - f"TTFT mean/p50/p99={ttft}", - f"TPOT mean/p50/p99={tpot}", - f"{benchmark.tool} {benchmark.tool_version}", - ] - ) - return " ".join(parts) - - -def _format_validation_gpus(validation: EndpointPresetValidation) -> str: - gpus = [ - _format_resources(resources, verbose=False) - for replica_group in validation.replicas - for resources in replica_group.resources - ] - return "+".join(gpus) or "-" - - -def _format_token_count(value: int) -> str: - for divisor, suffix in ((1024 * 1024, "M"), (1024, "K")): - if value >= divisor and value % divisor == 0: - return f"{value // divisor}{suffix}" - return str(value) - - -def _format_number(value: float) -> str: - return f"{value:.3g}" - - -def _format_latency(value_ms: float) -> str: - if value_ms >= 1000: - return f"{_format_number(value_ms / 1000)}s" - return f"{_format_number(value_ms)}ms" - - -def _format_latency_summary(*values_ms: float) -> str: - divisor, unit = (1000, "s") if max(values_ms) >= 1000 else (1, "ms") - return "/".join(_format_number(value / divisor) for value in values_ms) + unit - - -def _format_resources(resources, *, verbose: bool) -> str: - if resources is None: - return "-" - if verbose: - return resources.pretty_format() - gpu = resources.gpu - if gpu is None or gpu.count.max == 0: - return "-" - formatted = pretty_resources( - gpu_vendor=gpu.vendor, - gpu_name=",".join(gpu.name) if gpu.name else None, - gpu_count=gpu.count, - gpu_memory=gpu.memory, - total_gpu_memory=gpu.total_memory, - compute_capability=gpu.compute_capability, - ) - return formatted.removeprefix("gpu=") or "-" diff --git a/src/dstack/_internal/cli/services/endpoints/prompt.py b/src/dstack/_internal/cli/services/endpoints/prompt.py deleted file mode 100644 index 89e66964d..000000000 --- a/src/dstack/_internal/cli/services/endpoints/prompt.py +++ /dev/null @@ -1,58 +0,0 @@ -import enum -import json -from pathlib import Path -from typing import Any, Sequence - -from dstack._internal.cli.models.endpoints import EndpointConfiguration -from dstack._internal.core.models.profiles import ProfileParams - -_SYSTEM_PROMPT_PATH = Path(__file__).resolve().parent / "resources" / "system_prompt.md" - - -def get_endpoint_agent_system_prompt() -> str: - return _SYSTEM_PROMPT_PATH.read_text(encoding="utf-8").strip() - - -def format_endpoint_constraints( - configuration: EndpointConfiguration, - endpoint_env: dict[str, str], - *, - allowed_fleets: Sequence[str], -) -> str: - lines = [ - "Fixed endpoint constraints:", - "- Do not submit any task or service that conflicts with these values.", - ] - for field in ProfileParams.__fields__: - if field == "fleets": - continue - value = getattr(configuration, field) - if value is not None: - lines.append(f"- {field}: {_format_constraint_value(value)}") - if allowed_fleets: - lines.append(f"- fleets: {', '.join(allowed_fleets)}") - if configuration.gateway is not None: - lines.append(f"- gateway: {_format_constraint_value(configuration.gateway)}") - lines.append("- endpoint_env_keys: " + (", ".join(endpoint_env) if endpoint_env else "none")) - return "\n".join(lines) - - -def _format_constraint_value(value: Any) -> str: - data = _constraint_value_to_data(value) - if isinstance(data, (bool, dict, list)): - return json.dumps(data, sort_keys=True) - return str(data) - - -def _constraint_value_to_data(value: Any) -> Any: - if isinstance(value, enum.Enum): - return value.value - if hasattr(value, "format"): - return value.format() - if hasattr(value, "json"): - return json.loads(value.json(exclude_none=True)) - if isinstance(value, dict): - return {key: _constraint_value_to_data(item) for key, item in value.items()} - if isinstance(value, (list, tuple)): - return [_constraint_value_to_data(item) for item in value] - return value diff --git a/src/dstack/_internal/cli/services/endpoints/resources/system_prompt.md b/src/dstack/_internal/cli/services/endpoints/resources/system_prompt.md deleted file mode 100644 index a6ff533eb..000000000 --- a/src/dstack/_internal/cli/services/endpoints/resources/system_prompt.md +++ /dev/null @@ -1,332 +0,0 @@ -# Objective - -You are the endpoint preset creation agent for dstack. Produce one final dstack -service that can be saved as a reusable endpoint preset. Report success only -after that service answers a real request using `service_model_name` from -`Endpoint context:` through the dstack service URL. - -# Requested Model - -The `Endpoint context:` block contains either `model_repo` or `base_model`. - -- If it contains `model_repo`, deploy that repo/path exactly. -- If it contains `base_model`, choose a repo/path compatible with `base_model` - that best fits performance and hardware within the endpoint constraints, - allowed fleets, backends, and offers. A variant can be the base repo itself, a - different precision or quantization, or another trusted repo compatible with - `base_model`. - -If `Endpoint context:` contains `context_length`, the selected repo/path and -final service must support at least that context length. - -Use the real `dstack` CLI and shell commands in this workspace. Load and follow -`/dstack` for dstack CLI/YAML syntax. Load and follow `/dstack-prototyping` for -how to test a model-serving configuration with tasks before verifying it as a -service. The skill files are installed under `.claude/skills` if you need to -inspect them directly. - -# Progress - -Write endpoint progress with: - -```bash -progress "message" -``` - -The helper appends to `progress.jsonl`; the caller shows only the message text. - -Write progress before your first investigation action, whenever you submit a -run, when new evidence changes the plan, when model verification succeeds or -fails, and before the final report. Messages should explain what you tried, what -happened, and what you will do next. Do not put raw YAML, command output, long -tables, traces, or secrets in progress. - -Write a short progress message for every meaningful choice or action. Each -message should say what you did or chose, why, what evidence you used, what -happened, and what you will do next. Include the run name when a run is -involved. Include the fleet or backend when a fleet or backend is involved. Do -not use generic phase labels as the message. - -# Workspace Files - -Create local files only in the current workspace. - -`submissions.jsonl` is append-only. For every dstack task or service you submit, -append one JSON line when you submit it and more JSON lines when you learn its -run id, status, URL, or final outcome. - -Use these fields: - -- `event`: `submit`, `update`, or `final` -- `name`: submitted dstack run name -- `type`: `task` or `service` -- `status`: current known status -- `config_path`: YAML file path, when applicable -- `run_id`: run id, when known -- `reason`: why this run exists or why its status changed -- `service_url`: service URL, when known - -Example: - -```json -{"event":"submit","name":"qwen-endpoint-1","type":"task","status":"submitted","config_path":"qwen-endpoint-1.dstack.yml","reason":"test the serving image and local model request on an allowed fleet"} -{"event":"update","name":"qwen-endpoint-1","type":"task","status":"running","run_id":"..."} -``` - -Do not rewrite previous `submissions.jsonl` lines; the latest line for a run is -the current record. Stop runs you no longer need unless they are still needed for -attach/SSH debugging, logs, or backend diagnosis. - -After stopping a task or service, follow `/dstack` structured status guidance -and confirm that the run reached a terminal status before continuing. - -# Run Names - -Do not use the endpoint name itself as a submitted run name. - -Use only `-` for dstack runs submitted for -this endpoint. The first submitted run is `-1`, the second is -`-2`, and so on. Do not add framework, hardware, role, or purpose -suffixes to run names; record those details in workspace files and progress. - -# Endpoint Constraints - -Use existing allowed fleets only. Do not create, delete, apply, or edit fleets, -including `nodes`, `target`, `idle_duration`, backends, resources, max nodes, or -ownership. - -If the endpoint config lists fleets, use only those fleets. Otherwise, use the -existing project/imported fleets supplied in the request. - -The request also lists fixed endpoint constraints such as max price, spot -policy, backends, regions, instance types, fleets, env keys, tags, or backend -options. Do not submit a task or service that conflicts with those values. - -Put each fixed constraint that has a dstack service YAML field into the final -service YAML. For example, if the endpoint is limited to a fleet, max price, or -spot policy, the final `service_yaml` should include the corresponding -`service_yaml.fleets`, `service_yaml.max_price`, or `service_yaml.spot_policy` -fields. Use `/dstack` for exact field names. - -If you cannot submit a useful task or service within the allowed fleets and -constraints, write a failed `final_report.json`. The -`final_report.json.failure_summary` value should say which fleet or constraint -blocked the run and what the user/admin would need to change. - -# Task Usage - -Use `/dstack-prototyping` to learn how to use tasks. Keep in mind that using a -task is a must. This means `sleep infinity` and directly attaching inside the -task via SSH to run commands as the only way to use tasks. If there is any -ambiguity, follow `/dstack-prototyping` and do nothing that contradicts it. - -# Backend/Fleet Selection, Idle Duration, Instance Volumes - -Use only the fleets allowed by the endpoint request. If the endpoint request -also has constraints such as `backends`, `regions`, `instance_types`, -`spot_policy`, or `max_price`, apply them when running `dstack offer` if the CLI -has matching flags, and always apply them when submitting runs. - -## Backend And Fleet Choice - -Use `/dstack-prototyping` to learn how to select backends and fleets. - -Follow `/dstack-prototyping` skill on using a VM-based backend, Kubernetes backend, or SSH fleet that supports idle instances/instance volumes if there is such an option for the required GPU class. - -If backend allows (see above), use instance volumes to mount cache and model weights between runs. - -## Validating Offers - -When selecting backends/fleets and evaluating hardware that can be used to run -the model, use `dstack offer --json` and pass `--fleet` explicitly. If the -endpoint request has `fleets`, pass those fleets. Otherwise, pass every existing -fleet. If `--fleet` is not passed, `dstack offer` can show offers that are not -applicable to the fleets allowed by the endpoint request. Use these offers when -selecting fleet, backend, and hardware. - -Choosing fleet/backend is gated by classifying each against `https://dstack.ai/docs/concepts/backends.md`. VM-based backends are listed under `## VM-based` (they support idle instances and instance volumes). Kubernetes backend is listed under `## Container-based`, but supports instance volumes and thus is preferred over other container-based backends. -SSH fleets can be treated as VM-based backends as they support both idle instances (its equivalent) and instance volumes. - -## Model, Image, Serving Configuration, And Compute Fit - -Choose the model variant when allowed, image, serving configuration, and compute -to optimize expected performance within endpoint constraints and current offers. - -If `Endpoint context:` contains `model_repo`, choose fleet/backend/hardware that -can run `model_repo` within endpoint constraints. - -If `Endpoint context:` contains `base_model`, choose the repo/path and compute -together. You may pick the base repo or a compatible variant that fits the -allowed fleets, backends, and offers. - -If a task or service shows that the selected repo/path is a bad fit and -`Endpoint context:` contains `base_model`, pick another compatible variant if -available and test it in a task before submitting another service. - -## Submitting run - -When submitting a task or service, pass exact `fleets`, `backends`, and an -intentional `resources` range based on the choice made from offers, so the run -does not land outside the intended fleet/backend/hardware. - -## Decision Progress - -Progress should explain meaningful decisions and actions. - -When choosing a repo/path for `base_model`, fleet, backend, or hardware, write a -progress message that includes: - -- `service_model_name` and the selected repo/path, when they differ; -- the selected fleet, backend, and resource range; -- the offer/docs evidence used for the choice; -- the viable alternatives not selected and the exact reason; -- the fleet, backend, and resources that will be used in the submitted YAML. -- how the selected and rejected fleets/backends were classified; - -Backend-choice progress example (provide the same level of explanation for -repo/path, fleet, and hardware choice): - -"I chose backend ... on fleet ... because ...; I did not choose ... because ...; -I will submit YAML with fleets=..., backends=..., resources=...." - -# Final Service - -The final `service_yaml` is used to build the endpoint preset. It must -contain the full service config: `type: service`, the final run name, the service -model name, image/commands/port, resources, env references, and the fixed -endpoint constraints that apply to service YAML. - -Set final `service_yaml.name` to the final service run name. -If final `service_yaml.model` is a string, set it to `service_model_name` from -`Endpoint context:`. If final `service_yaml.model` is an object, set -`service_yaml.model.name` to `service_model_name`. - -Before submitting the final service, choose service resources from the least -restrictive requirements supported by the evidence, not from the exact machine -that happened to run. For example, if the model worked on an A40 but the -evidence only says it needs an NVIDIA GPU with at least 16GB memory, use that -broader requirement. Use an exact GPU, region, backend, or instance type only -when the endpoint constraints require it or the tested configuration depends on -that exact choice. - -Use run status to know whether the final service is still starting, running, or -failed. Use logs to understand failures. When dstack exposes the final service -run's `service.url`, build its absolute URL using `DSTACK_ENDPOINT_SERVER_URL` -and send a real model request using `DSTACK_ENDPOINT_BEARER_TOKEN` as the bearer -token. - -Verify the context length that the final service actually supports and report -it as `final_report.json.context_length`. - -# Benchmark - -Send benchmark requests to the final service run's absolute dstack service URL -using the bearer authentication used for service verification. Do not send them -to a server used during task prototyping or an SSH-forwarded local port. - -Run one benchmark with streaming responses. Choose a benchmark tool and workload -that can produce every required field below. - -Report the benchmark as `final_report.json.benchmark` using exactly the -following structure and field names (values are illustrative): - -```json -{ - "tool": "vllm bench serve", - "tool_version": "0.11.0", - "command": "vllm bench serve ...", - "workload": {"api": "chat_completions", "num_requests": 16, "input_tokens": 1024, "output_tokens": 128, "concurrency": 1}, - "metrics": { - "successful_requests": 16, "failed_requests": 0, "duration_seconds": 4.0, - "total_input_tokens": 16384, "total_output_tokens": 2048, - "ttft_ms": {"mean": 110.9, "p50": 108.2, "p99": 121.6}, - "tpot_ms": {"mean": 7.5, "p50": 7.4, "p99": 8.1} - } -} -``` - -Set `tool` to the command name and subcommands, without options or values (for -example, `vllm bench serve`). Set `tool_version` to the exact version and -`command` to the secret-free invocation. `api` must be `chat_completions` or -`completions`. `num_requests` is the number of measured requests; -`input_tokens` and `output_tokens` are the selected per-request lengths; and -`concurrency` is the maximum number of simultaneous requests. - -Calculate all metrics from the `num_requests` benchmark requests only. Exclude -setup, health-check, and warmup requests. `successful_requests` and -`failed_requests` are request counts; all requests must succeed. -`duration_seconds` is the elapsed wall-clock time from starting the first -measured request until the last measured request completes. `total_input_tokens` -and `total_output_tokens` are actual measured token totals. `ttft_ms` is the -time-to-first-token distribution across requests. `tpot_ms` is the -time-per-output-token distribution: for each request, divide the time from the -first output token to the last by one less than the actual output token count. -For both distributions, `mean`, `p50`, and `p99` are the arithmetic mean, 50th -percentile, and 99th percentile. Do not invent missing values. - -If benchmarking fails, write a failed `final_report.json`. - -Do not write a successful `final_report.json` until the final service answers a -request using `service_model_name` from `Endpoint context:` through the dstack -service URL. - -# Secrets - -The dstack CLI is already configured. Do not inspect `~/.dstack/config.yml` or -print, copy, or summarize tokens, secrets, or environment variable values. - -Do not expose the value of `DSTACK_TOKEN`, `DSTACK_ENDPOINT_BEARER_TOKEN`, or -the value of any environment variable listed under `endpoint_env_keys` in -`Fixed endpoint constraints:`. - -Do not put secret values in `final_report.json` or print them. Use env -references in `final_report.json.service_yaml`; use environment variable names or redacted values in -`final_report.json.benchmark.command`. - -# Resume - -On startup or resume, inspect `final_report.json` and `submissions.jsonl` -before submitting anything new. If `final_report.json` already reports success -for the current endpoint configuration, do not submit another run. Otherwise use -the next run number and record why the old run is not enough. - -# Final Report - -`final_report.json` may contain only `success`, `run_id`, `run_name`, -`service_yaml`, `base`, `model`, `context_length`, `benchmark`, and -`failure_summary`. - -On success, include exactly: - -- `success`: `true` -- `run_id`: the final verified service run ID -- `run_name`: the final verified service run name -- `service_yaml`: the full reusable service YAML described in `# Final Service` -- `base`: the base model repo, determined by the rules below -- `model`: the exact repo/path loaded by the final service command -- `context_length`: the context length verified for the final service -- `benchmark`: the final service benchmark described in `# Benchmark` - -Set `final_report.json.base` as follows: - -- If `Endpoint context:` contains `base_model`, set `final_report.json.base` to - `base_model`. -- If `Endpoint context:` contains `model_repo`, inspect the repo metadata, model - card, config, or another reliable source to identify the base model repo. -- If `model_repo` is itself the base model repo, set `final_report.json.base` to - `model_repo`. -- Do not infer `final_report.json.base` only from the repo name. - -On failure, include exactly: - -- `success`: `false` -- `failure_summary`: the reason a preset could not be created and any change - required from the user or administrator - -Write the report to `final_report.json`, then submit the identical JSON object -through `StructuredOutput`. - -Verify that `final_report.json` is correct and matches the required schema. - -Stop after one correct service is verified and benchmarked. P/D disaggregation -is not covered by the current endpoint agent or `/dstack-prototyping` skill. diff --git a/src/dstack/_internal/cli/services/endpoints/store.py b/src/dstack/_internal/cli/services/endpoints/store.py deleted file mode 100644 index bcb7c24db..000000000 --- a/src/dstack/_internal/cli/services/endpoints/store.py +++ /dev/null @@ -1,136 +0,0 @@ -import os -import sys -import tempfile -from pathlib import Path -from typing import List, TextIO - -import yaml -from pydantic import ValidationError - -from dstack._internal.cli.models.endpoint_presets import EndpointPreset -from dstack._internal.cli.models.endpoints import EndpointConfiguration -from dstack._internal.cli.services.endpoints.presets import endpoint_preset_to_data -from dstack._internal.core.errors import CLIError, ConfigurationError -from dstack._internal.utils.common import get_dstack_dir - - -class EndpointPresetStore: - def __init__(self, root: Path | None = None) -> None: - self.root = root or get_dstack_dir() / "presets" - - def list(self) -> list[EndpointPreset]: - if not self.root.exists(): - return [] - presets = [self._load(path) for path in self.root.glob("models--*/*.yaml")] - return sorted(presets, key=lambda preset: (preset.base.lower(), preset.id)) - - def get(self, preset_id: str) -> EndpointPreset | None: - paths = self._find_preset_paths(preset_id) - if not paths: - return None - if len(paths) > 1: - raise CLIError(f"Endpoint preset ID {preset_id!r} is not unique") - path = paths[0] - preset = self._load(path) - if preset.id != preset_id: - raise CLIError(f"Endpoint preset file {path} does not match its path") - return preset - - def save(self, preset: EndpointPreset) -> Path: - path = self._path(preset.base, preset.id) - path.parent.mkdir(parents=True, exist_ok=True) - content = yaml.safe_dump(endpoint_preset_to_data(preset), sort_keys=False) - fd, temporary_path = tempfile.mkstemp( - dir=path.parent, - prefix=f".{preset.id}.", - suffix=".tmp", - ) - try: - with os.fdopen(fd, "w", encoding="utf-8") as f: - f.write(content) - f.flush() - os.fsync(f.fileno()) - os.replace(temporary_path, path) - finally: - try: - Path(temporary_path).unlink() - except FileNotFoundError: - pass - return path - - def delete(self, preset_id: str) -> bool: - preset = self.get(preset_id) - if preset is None: - return False - path = self._path(preset.base, preset.id) - path.unlink() - try: - path.parent.rmdir() - except OSError: - pass - return True - - def delete_for_base(self, base: str) -> int: - directory = self._directory(base) - paths = list(directory.glob("*.yaml")) - presets = [self._load(path) for path in paths] - if any(preset.base != base for preset in presets): - raise CLIError(f"Endpoint preset directory {directory} contains another base model") - for path in paths: - path.unlink() - try: - directory.rmdir() - except OSError: - pass - return len(presets) - - def _load(self, path: Path) -> EndpointPreset: - try: - with path.open(encoding="utf-8") as f: - return EndpointPreset.parse_obj(yaml.safe_load(f)) - except (OSError, ValidationError, yaml.YAMLError) as e: - raise CLIError(f"Invalid endpoint preset file {path}: {e}") from e - - def _path(self, base: str, preset_id: str) -> Path: - if not preset_id or any(char in preset_id for char in "/\\"): - raise CLIError("Endpoint preset ID must not contain path separators") - return self._directory(base) / f"{preset_id}.yaml" - - def _find_preset_paths(self, preset_id: str) -> List[Path]: - if not preset_id or any(char in preset_id for char in "/\\"): - raise CLIError("Endpoint preset ID must not contain path separators") - return [ - path - for directory in self.root.glob("models--*") - if (path := directory / f"{preset_id}.yaml").is_file() - ] - - def _directory(self, base: str) -> Path: - directory = "models--" + base.replace("/", "--").replace("\\", "--") - return self.root / directory - - -def load_endpoint_configuration(path: str) -> tuple[str, EndpointConfiguration]: - if path == "-": - return "-", _parse_endpoint_configuration(sys.stdin) - configuration_path = Path(path) - if not configuration_path.is_file(): - raise ConfigurationError(f"Configuration file {path} does not exist") - try: - with configuration_path.open(encoding="utf-8") as f: - configuration = _parse_endpoint_configuration(f) - except OSError as e: - raise ConfigurationError(f"Failed to load configuration from {path}") from e - return str(configuration_path.resolve()), configuration - - -def _parse_endpoint_configuration(stream: TextIO) -> EndpointConfiguration: - try: - data = yaml.safe_load(stream) - if not isinstance(data, dict): - raise ConfigurationError("Endpoint configuration must be a YAML object") - return EndpointConfiguration.parse_obj(data) - except ValidationError as e: - raise ConfigurationError(e) from e - except yaml.YAMLError as e: - raise ConfigurationError(f"Invalid endpoint configuration: {e}") from e diff --git a/src/dstack/_internal/cli/services/endpoints/__init__.py b/src/dstack/_internal/cli/services/presets/__init__.py similarity index 100% rename from src/dstack/_internal/cli/services/endpoints/__init__.py rename to src/dstack/_internal/cli/services/presets/__init__.py diff --git a/src/dstack/_internal/cli/services/presets/agent.py b/src/dstack/_internal/cli/services/presets/agent.py new file mode 100644 index 000000000..201c2ed35 --- /dev/null +++ b/src/dstack/_internal/cli/services/presets/agent.py @@ -0,0 +1,630 @@ +import asyncio +import json +import os +import shutil +import signal +import subprocess +import time +from contextlib import asynccontextmanager, suppress +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, AsyncIterator, Callable, Optional, Sequence + +import psutil + +from dstack._internal.cli.models.preset_agent import AGENT_FINAL_REPORT_JSON_SCHEMA +from dstack._internal.cli.services.presets.redaction import redact, redact_structure +from dstack._internal.cli.services.presets.session import ( + PresetAgentSession, + _pid_alive, + _pid_running, + _process_started_at, + print_preset_progress, +) +from dstack._internal.cli.services.presets.tail import ( + _FileLineReader, + _OffsetStore, + _ProgressTailer, + _RecordMirror, + open_session_offsets, +) +from dstack._internal.cli.services.presets.workspace import ( + _PROGRESS_ENV, + PresetAgentWorkspace, +) +from dstack._internal.compat import IS_WINDOWS +from dstack._internal.core.errors import CLIError +from dstack._internal.core.services.configs import ConfigManager +from dstack.api import Client + +_CLAUDE_TOOLS = "Bash,Read,Write,Edit,WebFetch,WebSearch,StructuredOutput" +_CLAUDE_EFFORT_LEVELS = ("low", "medium", "high", "xhigh", "max") +_RESUME_DELAYS_SECONDS: tuple[int, ...] = (30, 60, 120) +_TERMINATE_GRACE_SECONDS = 3 +_RESUME_PROMPT = ( + "The previous agent process was interrupted. Continue where you left off. " + "Re-check the states of your runs before relying on them: time may have " + "passed, and tasks or instances may have stopped in the meantime." +) +_INHERITED_ENV_NAMES = ( + "PATH", + "HOME", + "USER", + "SSL_CERT_FILE", + "REQUESTS_CA_BUNDLE", + "HTTP_PROXY", + "HTTPS_PROXY", + "NO_PROXY", + "http_proxy", + "https_proxy", + "no_proxy", +) +_WINDOWS_INHERITED_ENV_NAMES = ( + "APPDATA", + "COMSPEC", + "HOMEDRIVE", + "HOMEPATH", + "LOCALAPPDATA", + "PATHEXT", + "PROGRAMDATA", + "SYSTEMDRIVE", + "SYSTEMROOT", + "USERPROFILE", + "USERNAME", + "WINDIR", +) + + +@dataclass(frozen=True) +class ClaudeAuth: + api_key: Optional[str] + executable: str + effort: Optional[str] + model: str + + +@dataclass +class PresetAgentProcessOutput: + report_data: Optional[dict[str, Any]] = None + error: Optional[str] = None + session_id: Optional[str] = None + made_progress: bool = False + + +def _get_claude_version(auth: "ClaudeAuth") -> Optional[str]: + try: + result = subprocess.run( + [auth.executable, "--version"], + capture_output=True, + text=True, + timeout=15, + ) + return result.stdout.strip() or None + except (OSError, subprocess.SubprocessError): + return None + + +def _get_claude_auth_status(auth: "ClaudeAuth") -> dict[str, Any]: + if auth.api_key: + return {"authMethod": "api-key"} + try: + result = subprocess.run( + [auth.executable, "auth", "status", "--json"], + capture_output=True, + text=True, + timeout=15, + ) + status = json.loads(result.stdout) + if isinstance(status, dict): + return status + except (OSError, subprocess.SubprocessError, json.JSONDecodeError): + pass + return {"authMethod": "unknown"} + + +def get_claude_auth() -> ClaudeAuth: + api_key = os.getenv("DSTACK_AGENT_ANTHROPIC_API_KEY") or None + configured_path = os.getenv("DSTACK_AGENT_CLAUDE_PATH") or "claude" + executable = shutil.which(configured_path) + if executable is None: + raise CLIError(f"Claude executable not found: {configured_path}") + effort = os.getenv("DSTACK_AGENT_CLAUDE_EFFORT") or None + if effort is not None and effort not in _CLAUDE_EFFORT_LEVELS: + raise CLIError( + f"DSTACK_AGENT_CLAUDE_EFFORT must be one of: {', '.join(_CLAUDE_EFFORT_LEVELS)}" + ) + return ClaudeAuth( + api_key=api_key, + executable=executable, + effort=effort, + model=os.getenv("DSTACK_AGENT_ANTHROPIC_MODEL", "claude-opus-4-8"), + ) + + +def build_preset_agent_env( + *, + api: Client, + preset_env: dict[str, str], + auth: ClaudeAuth, + workspace: PresetAgentWorkspace, + token: str, +) -> dict[str, str]: + config_manager = ConfigManager(workspace.dstack_home / ".dstack") + config_manager.configure_project( + name=api.project, + url=api.client.base_url, + token=token, + default=True, + ) + config_manager.save() + env = {name: value for name in _INHERITED_ENV_NAMES if (value := os.getenv(name))} + env.update(preset_env) + if IS_WINDOWS: + env.update( + {name: value for name in _WINDOWS_INHERITED_ENV_NAMES if (value := os.getenv(name))} + ) + env["PATH"] = os.pathsep.join([str(workspace.bin_path), env.get("PATH", "")]) + env["DSTACK_SERVER_URL"] = api.client.base_url + env["DSTACK_PROJECT"] = api.project + env["DSTACK_TOKEN"] = token + env[_PROGRESS_ENV] = str(workspace.progress_path) + for name in ["TMPDIR", "TEMP", "TMP"]: + env[name] = str(workspace.temp_path) + if auth.api_key is not None: + env["ANTHROPIC_API_KEY"] = auth.api_key + env["HOME"] = str(workspace.dstack_home) + if IS_WINDOWS: + env["USERPROFILE"] = str(workspace.dstack_home) + else: + env["HOME"] = str(Path.home()) + if IS_WINDOWS: + env["USERPROFILE"] = str(Path.home()) + return env + + +async def run_preset_agent( + *, + prompt: str, + env: dict[str, str], + workspace: PresetAgentWorkspace, + auth: ClaudeAuth, + redacted_values: Sequence[str], + agent_session: PresetAgentSession, + initial_resume_session_id: Optional[str] = None, +) -> PresetAgentProcessOutput: + offset_store = open_session_offsets(agent_session) + async with _session_tailers( + workspace=workspace, + agent_session=agent_session, + redacted_values=redacted_values, + offset_store=offset_store, + ): + resume_session_id: Optional[str] = initial_resume_session_id + attempt_prompt = prompt if resume_session_id is None else _RESUME_PROMPT + retry_delays = list(_RESUME_DELAYS_SECONDS) + while True: + command = _prepare_subprocess_command( + _build_claude_command(auth=auth, resume_session_id=resume_session_id) + ) + output, returncode = await _run_claude_process( + command=command, + prompt=attempt_prompt, + env=env, + workspace=workspace, + redacted_values=redacted_values, + agent_session=agent_session, + offset_store=offset_store, + ) + if output.report_data is None and returncode != 0: + output.error = output.error or f"Claude exited with return code {returncode}" + # Retry any process death without a submitted report; a terminal + # failure report from the agent returns immediately. + if output.report_data is not None or output.error is None: + return output + # A failed attempt that produced agent work is a new outage, not a + # continuation of the previous one: restore the full retry budget. + # Attempts that fail without any work drain it, so the loop always + # terminates when the network stays down. + if output.made_progress: + retry_delays = list(_RESUME_DELAYS_SECONDS) + # An externally recorded stop is a decision, not an outage: never + # resurrect an agent another CLI just terminated. + if agent_session.read_manifest().get("status") == "interrupted": + return output + if not retry_delays: + return output + delay = retry_delays.pop(0) + session_id = output.session_id or resume_session_id + if session_id is not None: + resume_session_id = session_id + attempt_prompt = _RESUME_PROMPT + action = "resuming" + else: + action = "retrying" + print_preset_progress( + f"Agent process exited without a report; {action} in {delay}s.", + agent_session=agent_session, + ) + await asyncio.sleep(delay) + + +async def _run_claude_process( + *, + command: list[str], + prompt: str, + env: dict[str, str], + workspace: PresetAgentWorkspace, + redacted_values: Sequence[str], + agent_session: PresetAgentSession, + offset_store: _OffsetStore, +) -> tuple[PresetAgentProcessOutput, int]: + proc: Optional[asyncio.subprocess.Process] = None + try: + # The agent's streams go to workspace files rather than pipes, so the + # agent survives CLI death (detach) and a later attach can continue + # parsing from the persisted offsets. + with ( + workspace.agent_stdout_path.open("ab") as stdout_file, + workspace.agent_stderr_path.open("ab") as stderr_file, + ): + proc = await asyncio.create_subprocess_exec( + *command, + cwd=workspace.path, + env=env, + stdin=asyncio.subprocess.PIPE, + stdout=stdout_file, + stderr=stderr_file, + start_new_session=not IS_WINDOWS, + # Inherit only the redirected std handles, not the CLI's other fds. + # Without this the untrusted agent inherits our open descriptors, and + # on Windows the broad inheritance flakes CreateProcess (WinError 87). + close_fds=True, + ) + agent_session.update_manifest( + agent_pid=proc.pid, agent_started_at=_process_started_at(proc.pid) + ) + assert proc.stdin is not None + proc.stdin.write(prompt.encode()) + with suppress(BrokenPipeError, ConnectionResetError): + await proc.stdin.drain() + proc.stdin.close() + + def agent_alive() -> bool: + return proc.returncode is None + + collect_task = asyncio.create_task( + _collect_agent_output( + workspace=workspace, + agent_session=agent_session, + redacted_values=redacted_values, + is_alive=agent_alive, + offset_store=offset_store, + ) + ) + try: + returncode = await proc.wait() + output = await collect_task + finally: + # Never leave the collector orphaned when an await above raises + # (cancellation, interrupt, or a wait failure). + if not collect_task.done(): + collect_task.cancel() + with suppress(asyncio.CancelledError): + await collect_task + except (KeyboardInterrupt, asyncio.CancelledError): + # The stop-or-detach decision belongs to the interrupt handler; the + # agent must stay alive here in case the user detaches. + raise + except BaseException: + if proc is not None and proc.returncode is None: + await _terminate_process(proc) + raise + + return output, returncode + + +def _build_claude_command( + *, auth: ClaudeAuth, resume_session_id: Optional[str] = None +) -> list[str]: + command = [ + auth.executable, + "-p", + "--output-format", + "stream-json", + "--verbose", + "--tools", + _CLAUDE_TOOLS, + "--allowedTools", + _CLAUDE_TOOLS, + "--disallowedTools", + "Task,NotebookEdit", + "--permission-mode", + "bypassPermissions", + "--model", + auth.model, + "--json-schema", + json.dumps(AGENT_FINAL_REPORT_JSON_SCHEMA), + ] + if auth.api_key is None: + command[2:2] = ["--setting-sources", "project,local"] + else: + command[2:2] = ["--bare"] + if auth.effort is not None: + command[2:2] = ["--effort", auth.effort] + if resume_session_id is not None: + command += ["--resume", resume_session_id] + return command + + +def _prepare_subprocess_command(command: list[str]) -> list[str]: + if not IS_WINDOWS or Path(command[0]).suffix.lower() not in {".bat", ".cmd"}: + return command + comspec = os.getenv("COMSPEC") or shutil.which("cmd.exe") + if comspec is None: + raise CLIError("Cannot run the Claude batch launcher because cmd.exe was not found") + return [comspec, "/d", "/s", "/c", subprocess.list2cmdline(command)] + + +def _write_debug_trace( + session: PresetAgentSession, + *, + stream_name: str, + text: str, + redacted_values: Sequence[str], +) -> None: + timestamp = ( + datetime.now(timezone.utc).isoformat(timespec="microseconds").replace("+00:00", "Z") + ) + try: + event = redact_structure(json.loads(text), redacted_values) + record = {"timestamp": timestamp, "stream": stream_name, "event": event} + except json.JSONDecodeError: + record = { + "timestamp": timestamp, + "stream": stream_name, + "text": redact(text.rstrip("\r\n"), redacted_values), + } + line = json.dumps(record, ensure_ascii=False, separators=(",", ":")) + with session.trace_path.open("a", encoding="utf-8") as f: + f.write(line + "\n") + f.flush() + + +@asynccontextmanager +async def _session_tailers( + *, + workspace: PresetAgentWorkspace, + agent_session: PresetAgentSession, + redacted_values: Sequence[str], + offset_store: _OffsetStore, +) -> AsyncIterator[None]: + """Mirrors the session's progress and record files while the body runs.""" + progress_tailer = _ProgressTailer( + path=workspace.progress_path, + redacted_values=redacted_values, + agent_session=agent_session, + offset_store=offset_store, + ) + record_mirrors = [ + _RecordMirror( + source=workspace.runs_path, + target=agent_session.runs_path, + redacted_values=redacted_values, + offset_store=offset_store, + offset_key="runs", + echo=agent_session.echo, + ), + _RecordMirror( + source=workspace.trials_path, + target=agent_session.trials_path, + redacted_values=redacted_values, + offset_store=offset_store, + offset_key="trials", + echo=agent_session.echo, + ), + ] + tailer_tasks = [ + asyncio.create_task(tailer.run()) for tailer in [progress_tailer, *record_mirrors] + ] + try: + yield + finally: + for task in tailer_tasks: + task.cancel() + for task in tailer_tasks: + with suppress(asyncio.CancelledError): + await task + progress_tailer.flush() + for mirror in record_mirrors: + mirror.flush() + + +async def _collect_agent_output( + *, + workspace: PresetAgentWorkspace, + agent_session: PresetAgentSession, + redacted_values: Sequence[str], + is_alive: Callable[[], bool], + offset_store: _OffsetStore, +) -> PresetAgentProcessOutput: + """Parses the agent's stream files until it exits; safe alongside a live + process or over the remains of a finished one.""" + stdout_output, _ = await asyncio.gather( + _read_process_stream( + stream=_FileLineReader( + workspace.agent_stdout_path, + offset_store=offset_store, + offset_key="agent_stdout", + is_alive=is_alive, + ), + stream_name="stdout", + parse_result=True, + redacted_values=redacted_values, + agent_session=agent_session, + ), + _read_process_stream( + stream=_FileLineReader( + workspace.agent_stderr_path, + offset_store=offset_store, + offset_key="agent_stderr", + is_alive=is_alive, + ), + stream_name="stderr", + parse_result=False, + redacted_values=redacted_values, + agent_session=agent_session, + ), + ) + # stderr is tailed with parse_result=False — it feeds the debug trace and + # advances the persisted offset, but can never contribute report data. + return stdout_output + + +async def attach_preset_agent( + *, + workspace: PresetAgentWorkspace, + redacted_values: Sequence[str], + agent_session: PresetAgentSession, +) -> PresetAgentProcessOutput: + """Follows a detached session's agent to completion, like + `run_preset_agent` without owning the process.""" + offset_store = open_session_offsets(agent_session) + async with _session_tailers( + workspace=workspace, + agent_session=agent_session, + redacted_values=redacted_values, + offset_store=offset_store, + ): + manifest = agent_session.read_manifest() + + def agent_alive() -> bool: + return _pid_alive(manifest.get("agent_pid"), manifest.get("agent_started_at")) + + return await _collect_agent_output( + workspace=workspace, + agent_session=agent_session, + redacted_values=redacted_values, + is_alive=agent_alive, + offset_store=offset_store, + ) + + +async def _read_process_stream( + *, + stream: "_FileLineReader", + stream_name: str, + parse_result: bool, + redacted_values: Sequence[str], + agent_session: PresetAgentSession, +) -> PresetAgentProcessOutput: + output = PresetAgentProcessOutput() + while True: + line = await stream.readline() + if not line: + return output + text = line.decode(errors="replace") + if agent_session.debug: + _write_debug_trace( + agent_session, + stream_name=stream_name, + text=text, + redacted_values=redacted_values, + ) + if not parse_result: + continue + try: + message = json.loads(text) + except json.JSONDecodeError: + continue + if not isinstance(message, dict): + continue + if output.session_id is None: + session_id = message.get("session_id") + if isinstance(session_id, str) and session_id: + output.session_id = session_id + agent_session.record_claude_session_id(session_id) + if message.get("type") == "assistant": + output.made_progress = True + if message.get("type") != "result": + continue + if message.get("is_error"): + error = message.get("result") or "Claude failed" + output.error = redact(str(error), redacted_values) + structured_output = message.get("structured_output") + if isinstance(structured_output, dict): + output.report_data = structured_output + continue + result = message.get("result") + if isinstance(result, str): + try: + parsed = json.loads(result) + except json.JSONDecodeError: + continue + if isinstance(parsed, dict): + output.report_data = parsed + + +async def _terminate_process(proc: asyncio.subprocess.Process) -> None: + """SIGTERM, a grace period, then SIGKILL — the same ladder as + `terminate_agent_process`, driven through the owned process handle.""" + if IS_WINDOWS: + await asyncio.to_thread(_terminate_windows_process_tree, proc.pid) + await proc.wait() + return + # The Windows branch returns above; Pyright still checks these POSIX-only APIs on Windows. + if hasattr(os, "killpg"): + with suppress(ProcessLookupError): + os.killpg(proc.pid, signal.SIGTERM) # pyright: ignore[reportAttributeAccessIssue] + else: + proc.terminate() + try: + await asyncio.wait_for(proc.wait(), timeout=_TERMINATE_GRACE_SECONDS) + except asyncio.TimeoutError: + if hasattr(os, "killpg"): + with suppress(ProcessLookupError): + os.killpg( # pyright: ignore[reportAttributeAccessIssue] + proc.pid, + signal.SIGKILL, # pyright: ignore[reportAttributeAccessIssue] + ) + else: + proc.kill() + await proc.wait() + + +def terminate_agent_process(manifest: dict[str, Any]) -> None: + """Terminates the session's agent process tree, if alive. The same + SIGTERM-grace-SIGKILL ladder as `_terminate_process`, driven by pid because + the caller (`preset stop`) never owned the process.""" + agent_pid = manifest.get("agent_pid") + if not isinstance(agent_pid, int) or not _pid_alive( + agent_pid, manifest.get("agent_started_at") + ): + return + if IS_WINDOWS: + _terminate_windows_process_tree(agent_pid) + return + with suppress(OSError): + os.killpg(agent_pid, signal.SIGTERM) # pyright: ignore[reportAttributeAccessIssue] + for _ in range(_TERMINATE_GRACE_SECONDS * 10): + if not _pid_running(agent_pid): + return + time.sleep(0.1) + with suppress(OSError): + os.killpg(agent_pid, signal.SIGKILL) # pyright: ignore[reportAttributeAccessIssue] + + +def _terminate_windows_process_tree(pid: int) -> None: + try: + root = psutil.Process(pid) + except psutil.NoSuchProcess: + return + processes = [*root.children(recursive=True), root] + for process in processes: + with suppress(psutil.NoSuchProcess): + process.terminate() + _, alive = psutil.wait_procs(processes, timeout=3) + for process in alive: + with suppress(psutil.NoSuchProcess): + process.kill() + psutil.wait_procs(alive, timeout=3) diff --git a/src/dstack/_internal/cli/services/presets/apply.py b/src/dstack/_internal/cli/services/presets/apply.py new file mode 100644 index 000000000..40495a357 --- /dev/null +++ b/src/dstack/_internal/cli/services/presets/apply.py @@ -0,0 +1,98 @@ +import argparse +from typing import Optional + +from rich.markup import escape + +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.models.presets import Preset +from dstack._internal.cli.services.configurators.run import ServiceConfigurator +from dstack._internal.cli.services.presets.output import ( + format_preset_benchmark, +) +from dstack._internal.cli.services.presets.store import PresetStore +from dstack._internal.core.errors import CLIError +from dstack._internal.core.models.configurations import ServiceConfiguration +from dstack._internal.core.models.profiles import ProfileParams +from dstack.api import Client + + +def apply_preset( + *, + api: Client, + configuration: PresetConfiguration, + configuration_path: str, + preset_id: str, + profile_name: Optional[str], + command_args: argparse.Namespace, + store: PresetStore, +) -> None: + preset = store.get(preset_id) + if preset is None: + raise CLIError(f"Preset {preset_id} does not exist") + _validate_preset_matches(preset, configuration=configuration) + + configurator = ServiceConfigurator(api_client=api) + service_args = configurator.get_parser().parse_args([]) + service_args.profile = profile_name + service = _build_service(configuration, preset) + run_plan, repo = configurator.get_plan( + conf=service, + configuration_path=configuration_path, + configurator_args=service_args, + ) + configurator.apply_plan( + run_plan=run_plan, + repo=repo, + command_args=command_args, + configurator_args=service_args, + plan_properties={ + "Model": _format_requested_model(configuration), + "Preset": _format_selected_preset(preset), + }, + ) + + +def _validate_preset_matches(preset: Preset, *, configuration: PresetConfiguration) -> None: + """The referenced preset must serve what the configuration asks for.""" + model_name = configuration.model.api_model_name + service_model = preset.service.model + if service_model is None or service_model.name.lower() != model_name.lower(): + raise CLIError(f"Preset {preset.id} does not serve {model_name}") + if configuration.context_length is not None: + if preset.context_length < configuration.context_length: + raise CLIError( + f"Preset {preset.id} does not support context length" + f" {configuration.context_length}" + ) + if configuration.model.allows_variant_selection: + if preset.base.lower() != model_name.lower(): + raise CLIError(f"Preset {preset.id} does not serve base model {model_name}") + elif preset.model != configuration.model.exact_repo: + raise CLIError(f"Preset {preset.id} does not serve repo {configuration.model.exact_repo}") + + +def _build_service( + configuration: PresetConfiguration, + preset: Preset, +) -> ServiceConfiguration: + service = preset.service.copy(deep=True) + service.name = configuration.name + service.gateway = configuration.gateway + service.env.update(configuration.env) + for field in ProfileParams.__fields__: + value = getattr(configuration, field) + if value is not None: + setattr(service, field, value) + return service + + +def _format_requested_model(configuration: PresetConfiguration) -> str: + model = escape(configuration.model.api_model_name) + if configuration.model.allows_variant_selection: + return f"{model} ([secondary]base[/])" + return model + + +def _format_selected_preset(preset: Preset) -> str: + details = format_preset_benchmark(preset, verbose=True) + return f"{escape(preset.id)} ([secondary]{details}[/])" diff --git a/src/dstack/_internal/cli/services/presets/create.py b/src/dstack/_internal/cli/services/presets/create.py new file mode 100644 index 000000000..cf14cfebf --- /dev/null +++ b/src/dstack/_internal/cli/services/presets/create.py @@ -0,0 +1,942 @@ +import asyncio +import dataclasses +import json +import os +import re +import time +import uuid +from contextlib import nullcontext, suppress +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Optional, Sequence + +import yaml +from rich.table import Table +from rich.text import Text + +from dstack._internal.cli.models.configurations import ( + PresetConfiguration, + PresetConstraints, +) +from dstack._internal.cli.models.preset_agent import AgentFinalReport +from dstack._internal.cli.models.presets import Preset +from dstack._internal.cli.services.presets.agent import ( + ClaudeAuth, + PresetAgentProcessOutput, + attach_preset_agent, + build_preset_agent_env, + get_claude_auth, + run_preset_agent, + terminate_agent_process, +) +from dstack._internal.cli.services.presets.presets import preset_to_data +from dstack._internal.cli.services.presets.prompt import get_preset_agent_system_prompt +from dstack._internal.cli.services.presets.redaction import ( + contains_redacted_value, + get_redacted_values, + get_sensitive_inherited_env_values, + redact, +) +from dstack._internal.cli.services.presets.session import ( + PresetAgentSession, + SessionBusyError, + claimed_session_name, + create_preset_agent_session, + find_session_name_claims, + iter_agent_sessions, + load_agent_session, + load_attachable_agent_session, + mark_session_owner, + print_preset_progress, + print_session_log, + release_session_claim, + session_process_alive, + session_report_exists, + try_claim_session, +) +from dstack._internal.cli.services.presets.store import PresetStore +from dstack._internal.cli.services.presets.verify import ( + build_verified_preset, + load_preset_agent_report, +) +from dstack._internal.cli.services.presets.workspace import ( + PresetAgentWorkspace, + attach_agent_workspace, + create_agent_workspace, + remove_agent_workspace, + scrub_workspace_token, +) +from dstack._internal.cli.utils.common import NO_OFFERS_WARNING, confirm_ask, console, warn +from dstack._internal.cli.utils.offers import print_offers_table +from dstack._internal.core.errors import CLIError, ConfigurationError +from dstack._internal.core.models.configurations import TaskConfiguration +from dstack._internal.core.models.envs import Env, EnvSentinel +from dstack._internal.core.models.fleets import FleetStatus +from dstack._internal.core.models.runs import RunSpec +from dstack.api import Client + +_RUN_STOP_TIMEOUT_SECONDS = 10 * 60 +_NO_FLEETS_ERROR = "The project has no fleets. Create one before creating a preset" + + +@dataclass(frozen=True) +class PresetCreateResult: + preset: Preset + path: Path + final_run_id: uuid.UUID + final_run_name: str + + +class CreationStopped(Exception): + """Another CLI stopped this creation; it already recorded the interruption.""" + + +class AgentExitedWithoutReport(Exception): + """A detached agent died without submitting a report; the session is + resumable rather than failed.""" + + def __init__(self, error: Optional[str]) -> None: + super().__init__(error or "The agent exited without a report") + self.error = error + + +def follow_preset( + *, + api: Client, + store: PresetStore, + preset_id: str, + keep_service: bool = False, + wait_for_run_stop: bool = True, + echo: bool = True, +) -> PresetCreateResult: + """Re-owns a detached session: follows its agent to completion, then + verifies and saves the preset (the finalize role, which must run CLI-side + for secret-scrubbing and server-verified preset building). + + Always takes the exclusive finalize lock so a concurrent `logs -f` and + reconcile can't both finalize the same session. `wait_for_run_stop=False` + and `echo=False` make it non-blocking and silent for reconcile.""" + agent_session = load_attachable_agent_session(preset_id) + agent_session.echo = echo + lock = try_claim_session(agent_session) + if lock is None: + raise SessionBusyError(f"Preset {preset_id} is being finalized by another process") + try: + configuration = _load_session_configuration(agent_session) + try: + result = asyncio.run( + _create_preset( + api=api, + configuration=_resolve_preset_env(configuration, strict=False), + source_configuration=configuration, + store=store, + keep_service=keep_service, + agent_session=agent_session, + attach=True, + wait_for_run_stop=wait_for_run_stop, + ) + ) + except KeyboardInterrupt: + # `logs -f` is a viewer: Ctrl+C detaches and leaves the agent + # running (reconcile finalizes it later), never stops it. + _detach_agent_session(agent_session) + raise + except AgentExitedWithoutReport as e: + _stop_active_session_runs(api, agent_session) + _suspend_agent_session(agent_session) + raise CLIError(str(e)) from e + except CLIError: + # Definitive: a failure/invalid report, an unverifiable service, or a + # leaked secret — the preset genuinely cannot be built, so fail it. + _close_agent_session(agent_session, "failed") + raise + # A transient error (network / OS) propagates untouched: the completed + # session and its report stay intact for a later follow or reconcile. + _close_agent_session(agent_session, "success") + return result + finally: + release_session_claim(lock) + + +def _load_session_configuration(agent_session: PresetAgentSession) -> PresetConfiguration: + configuration_path = agent_session.path / "preset.dstack.yml" + if not configuration_path.is_file(): + raise CLIError( + f"Preset {agent_session.preset_id} has no saved configuration and cannot be" + f" followed; resume it with --resume {agent_session.preset_id} instead" + ) + # The session copy is canonical output, not user input: parse it without + # the user-facing deprecation warnings. + try: + return PresetConfiguration.parse_obj( + yaml.safe_load(configuration_path.read_text(encoding="utf-8")) + ) + except (OSError, ValueError) as e: + raise CLIError(f"Could not read the preset configuration: {e}") from e + + +def show_preset_session_logs( + *, + project: Optional[str], + store: PresetStore, + preset_id: str, + follow: bool, + keep_service: bool, +) -> Optional[PresetCreateResult]: + """`logs`: dump a session's log (any status). With `follow`, a still-live + session is re-owned, followed to completion, and its preset saved; a + finished session just prints its log. Returns the saved preset, if any.""" + session = load_agent_session(preset_id) + status = session.read_manifest().get("status") + if not follow or status in ("success", "failed", "interrupted"): + print_session_log(session) + return None + # Following a live session: print the log so far, then stream new progress + # (a future --since could bound this). The client is built only here, so a + # read-only dump never needs a server or authentication. + print_session_log(session) + try: + return follow_preset( + api=Client.from_config(project_name=project), + store=store, + preset_id=preset_id, + keep_service=keep_service, + ) + except SessionBusyError: + # Another CLI already owns the finalize; follow read-only instead of + # refusing, so any number of viewers can watch the same preset at once. + _follow_session_log_readonly(session) + return None + + +def _follow_session_log_readonly(session: PresetAgentSession) -> None: + """Read-only follow: another CLI owns the finalize, so just stream the log it + writes until the preset reaches a terminal state.""" + try: + offset = session.log_path.stat().st_size + except OSError: + offset = 0 + while True: + try: + with session.log_path.open("r", encoding="utf-8", errors="replace") as f: + f.seek(offset) + chunk = f.read() + offset = f.tell() + except OSError: + chunk = "" + if chunk: + console.print(Text(chunk.rstrip("\n")), soft_wrap=True) + manifest = session.read_manifest() + if manifest.get("status") in ("success", "failed", "interrupted"): + return + if not chunk and not session_process_alive(manifest): + # The owner died without recording a terminal status; stop tailing + # rather than poll forever. + console.print( + f"The process creating preset [code]{session.preset_id}[/] exited;" + f" follow it again with [code]dstack preset logs -f {session.preset_id}[/]" + ) + return + time.sleep(1) + + +def reconcile_detached_sessions(store: PresetStore) -> None: + """Finalizes sessions whose agent completed while no CLI was attached + (graceful detach, or an ungraceful CLI death). This is what makes the saved + preset independent of a foreground process: any read command runs it, and + the work materializes from the on-disk report. + + Best-effort and parallel-safe — finalize takes an exclusive claim, and every + error is swallowed so the calling read command never fails. + """ + for session in iter_agent_sessions(): + if _is_reconcilable(session.read_manifest()): + _reconcile_session(session, store) + + +def _is_reconcilable(manifest: dict[str, Any]) -> bool: + # An orphaned session (no live owner) whose agent left a completion report. + # A session interrupted mid-work has no report and stays resumable; one + # stopped *after* the agent finished is finalized by `stop` itself, not here. + # Sessions created before finalize context was persisted lack `project` and + # are skipped — they finalize interactively via `logs -f`. + return ( + manifest.get("status") == "running" + and bool(manifest.get("project")) + and session_report_exists(manifest) + and not session_process_alive(manifest) + ) + + +def _reconcile_session(session: PresetAgentSession, store: PresetStore) -> None: + manifest = session.read_manifest() + try: + api = Client.from_config(project_name=str(manifest.get("project") or "")) + except Exception: # noqa: BLE001 — offline/misconfigured must not break the read command + return + # follow_preset takes the finalize claim (so a concurrent `logs -f` + # or reconcile can't double-finalize), records the terminal status itself, + # and leaves the session intact on a transient error. Every outcome is silent + # here — the result shows in the list that follows. + with suppress(Exception): + follow_preset( + api=api, + store=store, + preset_id=session.preset_id, + keep_service=bool(manifest.get("keep_service")), + wait_for_run_stop=False, + echo=False, + ) + + +def stop_preset_session(api: Client, preset_id: str) -> None: + session = load_agent_session(preset_id) + manifest = session.read_manifest() + status = manifest.get("status") + if status == "success": + console.print(f"Preset [code]{preset_id}[/] is already created") + return + if status == "failed": + console.print(f"Preset [code]{preset_id}[/] creation failed") + return + if status == "interrupted": + console.print(f"Preset [code]{preset_id}[/] creation was interrupted.") + return + if session_report_exists(manifest) and not session_process_alive(manifest): + # The agent already finished; finalize like reconcile would instead of + # leaving a not-yet-saved intermediate state behind. + follow_preset( + api=api, + store=PresetStore(), + preset_id=preset_id, + keep_service=bool(manifest.get("keep_service")), + echo=False, + ) + console.print(f"Preset [code]{preset_id}[/] is already created") + return + # Stop wins, like `dstack stop`: record the intent first so a live owner's + # retry loop exits instead of resurrecting the agent, then terminate. + _finish_agent_session(session, "interrupted") + terminate_agent_process(manifest) + _stop_active_session_runs(api, session) + _suspend_agent_session(session) + + +def _stop_active_session_runs(api: Client, session: PresetAgentSession) -> None: + """Stops the session's non-terminal runs (with a spinner), like `dstack + stop`. Keeping a trial instance warm for resume is the detach path, not this.""" + names = _load_submitted_run_names(session.runs_path) + active = [] + for name in names: + try: + run = api.client.runs.get(api.project, name) + except Exception: # noqa: BLE001 + continue + if not run.status.is_finished(): + active.append(name) + if not active: + return + with console.status("Stopping runs..."): + api.client.runs.stop(api.project, active, abort=False) + + +def _resolve_preset_env( + configuration: PresetConfiguration, *, strict: bool = True +) -> PresetConfiguration: + """Resolves `EnvSentinel` entries from the process environment. Non-strict + drops unresolvable entries instead of raising — for attach, where env values + only feed redaction and the agent already runs.""" + configuration = configuration.copy(deep=True) + resolved: dict[str, str] = {} + for key, value in configuration.env.items(): + if isinstance(value, EnvSentinel): + try: + resolved[key] = value.from_env(os.environ) + except ValueError as e: + if strict: + raise ConfigurationError(str(e)) from e + else: + resolved[key] = value + configuration.env = Env.parse_obj(resolved) + return configuration + + +def create_preset( + *, + api: Client, + configuration: PresetConfiguration, + store: PresetStore, + keep_service: bool = False, + build_name: Optional[str] = None, + debug: bool = False, + resume_session: Optional[PresetAgentSession] = None, + user_prompt: Optional[str] = None, + allowed_fleets: Optional[tuple[str, ...]] = None, +) -> PresetCreateResult: + agent_session = resume_session or create_preset_agent_session(configuration, debug=debug) + try: + resolved_configuration = _resolve_preset_env(configuration) + result = asyncio.run( + _create_preset( + api=api, + configuration=resolved_configuration, + source_configuration=configuration, + store=store, + keep_service=keep_service, + build_name=build_name, + agent_session=agent_session, + resume=resume_session is not None, + user_prompt=user_prompt, + allowed_fleets=allowed_fleets, + ) + ) + except KeyboardInterrupt: + _stop_or_detach_agent_session(agent_session, api) + raise + except CreationStopped: + # The stopping CLI already suspended the session. + raise + except BaseException: + _close_agent_session(agent_session, "failed") + raise + _close_agent_session(agent_session, "success") + return result + + +@dataclass(frozen=True) +class _CreationSetup: + """Per-mode inputs to the shared creation path, built by one of + `_fresh_setup`, `_resume_setup`, or `_attach_setup`.""" + + auth: Optional[ClaudeAuth] + workspace: PresetAgentWorkspace + build_name: str + allowed_fleets: tuple[str, ...] + user_prompt: Optional[str] + initial_resume_session_id: Optional[str] + write_constraints: bool # True only for fresh creations + + +def _fresh_setup( + api: Client, + configuration: PresetConfiguration, + agent_session: PresetAgentSession, + build_name: Optional[str], + allowed_fleets: Optional[tuple[str, ...]], + user_prompt: Optional[str], +) -> _CreationSetup: + if allowed_fleets is None: + allowed_fleets = _get_allowed_fleets(api, configuration) + if not allowed_fleets: + raise CLIError(_NO_FLEETS_ERROR) + auth = get_claude_auth() + workspace = create_agent_workspace(agent_session) + build_name = build_name or _get_build_name( + configuration.name, configuration.model.api_model_name, agent_session.preset_id + ) + return _CreationSetup( + auth=auth, + workspace=workspace, + build_name=build_name, + allowed_fleets=allowed_fleets, + user_prompt=user_prompt, + initial_resume_session_id=None, + write_constraints=True, + ) + + +def _resume_setup( + agent_session: PresetAgentSession, + build_name: Optional[str], + user_prompt: Optional[str], +) -> _CreationSetup: + # The prompt is fixed at session creation, like the constraints. + pinned_prompt = agent_session.read_user_prompt() + if user_prompt is not None and user_prompt != pinned_prompt: + warn( + "The configuration prompt is ignored when resuming: the preset keeps its original prompt" + ) + user_prompt = pinned_prompt + auth = get_claude_auth() + workspace = attach_agent_workspace(agent_session) + manifest = agent_session.read_manifest() + claude_model = manifest.get("claude_model") + if isinstance(claude_model, str) and claude_model: + auth = dataclasses.replace(auth, model=claude_model) + initial_resume_session_id: Optional[str] = None + claude_session_id = manifest.get("claude_session_id") + if isinstance(claude_session_id, str) and claude_session_id: + initial_resume_session_id = claude_session_id + return _CreationSetup( + auth=auth, + workspace=workspace, + build_name=build_name or _load_build_name(workspace), + allowed_fleets=(), + user_prompt=user_prompt, + initial_resume_session_id=initial_resume_session_id, + write_constraints=False, + ) + + +def _attach_setup( + agent_session: PresetAgentSession, + build_name: Optional[str], +) -> _CreationSetup: + workspace = attach_agent_workspace(agent_session) + return _CreationSetup( + auth=None, + workspace=workspace, + build_name=build_name or _load_build_name(workspace), + allowed_fleets=(), + user_prompt=None, + initial_resume_session_id=None, + write_constraints=False, + ) + + +async def _create_preset( + *, + api: Client, + configuration: PresetConfiguration, + store: PresetStore, + source_configuration: Optional[PresetConfiguration] = None, + keep_service: bool = False, + build_name: Optional[str] = None, + agent_session: PresetAgentSession, + resume: bool = False, + attach: bool = False, + wait_for_run_stop: bool = True, + user_prompt: Optional[str] = None, + allowed_fleets: Optional[tuple[str, ...]] = None, +) -> PresetCreateResult: + source_configuration = source_configuration or configuration + if attach: + setup = _attach_setup(agent_session, build_name) + elif resume: + setup = _resume_setup(agent_session, build_name, user_prompt) + else: + setup = _fresh_setup( + api, configuration, agent_session, build_name, allowed_fleets, user_prompt + ) + # Record ownership + the finalize context (project, keep-service) so a later + # detached reconcile can complete this session from disk alone. + mark_session_owner( + agent_session, + project=api.project, + keep_service=keep_service, + claude_model=setup.auth.model if setup.auth is not None else None, + ) + + preset_env = configuration.env.as_dict() + token = getattr(api.client, "_token", None) + if not isinstance(token, str) or not token: + raise CLIError("The configured dstack client has no authentication token") + redacted_values = get_redacted_values( + [ + token, + (setup.auth.api_key if setup.auth is not None else None) or "", + *preset_env.values(), + *get_sensitive_inherited_env_values(), + ] + ) + env: dict[str, str] = {} + report: Optional[AgentFinalReport] = None + preset: Optional[Preset] = None + preset_path: Optional[Path] = None + creation_succeeded = False + interrupted = False + cleanup_error: Optional[str] = None + if setup.auth is not None: + env = build_preset_agent_env( + api=api, + preset_env=preset_env, + auth=setup.auth, + workspace=setup.workspace, + token=token, + ) + prompt = get_preset_agent_system_prompt(user_prompt=setup.user_prompt) + if setup.write_constraints: + if setup.user_prompt: + agent_session.write_user_prompt(setup.user_prompt) + constraints_text = _build_constraints( + configuration=configuration, + build_name=setup.build_name, + allowed_fleets=setup.allowed_fleets, + ) + setup.workspace.constraints_path.write_text(constraints_text, encoding="utf-8") + if agent_session.debug: + agent_session.write_prompt(prompt) + agent_session.write_constraints(constraints_text) + if setup.auth is not None: + agent_session.write_agent_info(setup.auth) + try: + if attach: + process_output = await attach_preset_agent( + workspace=setup.workspace, + redacted_values=redacted_values, + agent_session=agent_session, + ) + _raise_if_externally_stopped(agent_session, process_output) + if ( + process_output.report_data is None + and not setup.workspace.final_report_path.exists() + ): + raise AgentExitedWithoutReport(process_output.error) + else: + assert setup.auth is not None + process_output = await run_preset_agent( + prompt=prompt, + env=env, + workspace=setup.workspace, + auth=setup.auth, + redacted_values=redacted_values, + agent_session=agent_session, + initial_resume_session_id=setup.initial_resume_session_id, + ) + _raise_if_externally_stopped(agent_session, process_output) + report = load_preset_agent_report( + output=process_output, + workspace=setup.workspace, + redacted_values=redacted_values, + ) + run = api.client.runs.get(api.project, report.run_name) + preset = build_verified_preset( + run=run, + preset_configuration=source_configuration, + report=report, + preset_id=agent_session.preset_id or None, + name=claimed_session_name(agent_session.read_manifest()), + ) + if contains_redacted_value(preset_to_data(preset), redacted_values): + raise CLIError("Generated preset contains a secret value") + preset_path = store.save(preset) + creation_succeeded = True + except (KeyboardInterrupt, asyncio.CancelledError): + interrupted = True + raise + finally: + if agent_session.debug: + _save_final_report_copy( + workspace=setup.workspace, + agent_session=agent_session, + redacted_values=redacted_values, + ) + if not interrupted: + keep_final_service = keep_service and creation_succeeded + try: + await _cleanup_runs( + api=api, + build_name=setup.build_name, + workspace=setup.workspace, + final_run_name=report.run_name if report is not None else None, + keep_final_service=keep_final_service, + agent_session=agent_session, + wait_for_stop=wait_for_run_stop, + ) + except Exception as e: + cleanup_error = str(e) + + if cleanup_error is not None: + # The preset is already saved by this point; a failed cleanup only means + # trial runs may still be running. Warn rather than fail the (successful) + # session — otherwise a transient blip would discard completed work. + if agent_session.echo: + warn(f"Failed to stop preset creation runs: {cleanup_error}") + assert preset is not None + assert preset_path is not None + assert report is not None + assert report.run_id is not None + assert report.run_name is not None + return PresetCreateResult( + preset=preset, + path=preset_path, + final_run_id=report.run_id, + final_run_name=report.run_name, + ) + + +def _raise_if_externally_stopped( + session: PresetAgentSession, output: "PresetAgentProcessOutput" +) -> None: + if output.report_data is None and session.read_manifest().get("status") == "interrupted": + raise CreationStopped + + +def _finish_agent_session( + session: PresetAgentSession, + status: str, +) -> None: + try: + session.finish(status) + except OSError as e: + if session.echo: + warn(f"Could not finalize agent output. Files remain at {session.path}: {e}") + + +def _close_agent_session(session: PresetAgentSession, status: str) -> None: + """Records the terminal status and removes the workspace alias.""" + _finish_agent_session(session, status) + remove_agent_workspace(session) + + +def _detach_agent_session(session: PresetAgentSession) -> None: + """Releases ownership but leaves the agent running — it stays visible and + reconcilable in `dstack preset`. Silent: `logs -f` calls this on Ctrl+C, and + a viewer that just stops watching shouldn't announce anything.""" + session.update_manifest(pid=None) + + +def _stop_or_detach_agent_session( + session: PresetAgentSession, api: Optional[Client] = None +) -> None: + """`create` interrupt: stop the session, or detach and leave the agent + working — it stays visible as a running session in `dstack preset`.""" + manifest = session.read_manifest() + agent_alive = session_process_alive({**manifest, "pid": None}) + stop = True + if agent_alive: + try: + stop = confirm_ask(f"Stop creating preset [code]{session.preset_id}[/]?") + except (KeyboardInterrupt, EOFError): + stop = True + if not stop: + _detach_agent_session(session) + console.print( + f"\nDetached. Follow with [code]dstack preset logs -f {session.preset_id}[/]" + f" or stop with [code]dstack preset stop {session.preset_id}[/]." + ) + return + terminate_agent_process(manifest) + if api is not None: + _stop_active_session_runs(api, session) + _suspend_agent_session(session) + + +def _suspend_agent_session(session: PresetAgentSession) -> None: + try: + session.finish("interrupted") + except OSError as e: + warn(f"Could not record the interrupted preset state: {e}") + # The kept workspace must not retain a live credential while suspended. + scrub_workspace_token(session) + console.print(f"\nPreset [code]{session.preset_id}[/] creation interrupted.") + console.print( + f"Resume it with [code]dstack preset create -f --resume {session.preset_id}[/]." + ) + + +def _get_build_name(name: Optional[str], model_name: str, suffix: str) -> str: + base = name or _model_slug(model_name) + # Leave room for the preset id and numeric submission suffix while retaining + # a recognizable prefix. + prefix = base[:26].rstrip("-") + return f"{prefix}-{suffix}" + + +def _model_slug(model_name: str) -> str: + """A run-name-safe slug for name-less presets, from the model's basename.""" + basename = model_name.rsplit("/", 1)[-1] + slug = re.sub(r"[^a-z0-9]+", "-", basename.lower()).strip("-") + if not slug or not slug[0].isalpha(): + slug = f"model-{slug}".strip("-") + return slug + + +def _load_build_name(workspace: PresetAgentWorkspace) -> str: + try: + data = json.loads(workspace.constraints_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as e: + raise CLIError(f"The preset creation constraints cannot be read: {e}") from e + prefix = data.get("run_name_prefix") if isinstance(data, dict) else None + if not isinstance(prefix, str) or not prefix: + raise CLIError("The preset creation constraints do not contain a run name prefix") + return prefix + + +@dataclass(frozen=True) +class PresetNameHolders: + """Current claims on a preset name: at most one saved preset, plus any + creation sessions claiming it (excluding the holder preset's own session).""" + + name: str + preset: Optional[Preset] + sessions: list[PresetAgentSession] + + @property + def preset_ids(self) -> list[str]: + ids = [self.preset.id] if self.preset is not None else [] + ids.extend(session.preset_id for session in self.sessions) + return ids + + +def find_preset_name_holders(store: PresetStore, name: str) -> PresetNameHolders: + preset = store.find_by_name(name) + sessions = [ + session + for session in find_session_name_claims(name) + if preset is None or session.preset_id != preset.id + ] + return PresetNameHolders(name=name, preset=preset, sessions=sessions) + + +def reassign_preset_name(store: PresetStore, holders: PresetNameHolders) -> None: + """Releases the name from every holder so a new preset can claim it.""" + if holders.preset is not None: + store.release_name(holders.name) + for session in holders.sessions: + session.update_manifest(name=None) + + +def plan_preset(*, api: Client, configuration: PresetConfiguration) -> tuple[str, ...]: + """Resolves the allowed fleets and shows what the agent will have to work + with — Project, User, the effective fleets, and their offers. Agent-free.""" + allowed_fleets = _get_allowed_fleets(api, configuration) + if not allowed_fleets: + raise CLIError(_NO_FLEETS_ERROR) + _print_fleet_offers(api, allowed_fleets) + return allowed_fleets + + +def _print_fleet_offers(api: Client, allowed_fleets: tuple[str, ...]) -> None: + try: + # Image and user are set so the server neither defaults gpu.vendor to + # nvidia nor pulls image config from a registry (as in `dstack offer`). + offer_configuration = TaskConfiguration(commands=[":"], image="scratch", user="root") + offer_configuration.fleets = list(allowed_fleets) + run_spec = RunSpec(configuration=offer_configuration, profile=None) + with console.status("Getting offers..."): + run_plan = api.client.runs.get_plan(api.project, run_spec, max_offers=10) + props = Table(box=None, show_header=False) + props.add_column(no_wrap=True) + props.add_column() + props.add_row("[bold]Project[/bold]", run_plan.project_name) + props.add_row("[bold]User[/bold]", run_plan.user) + props.add_row("[bold]Fleets[/bold]", ", ".join(allowed_fleets)) + console.print(props) + console.print() + job_plan = run_plan.job_plans[0] + if job_plan.offers: + print_offers_table( + offers=job_plan.offers, + total_offers=job_plan.total_offers, + max_price=job_plan.max_price or 0.0, + mute_tail_rows=False, + ) + else: + console.print(NO_OFFERS_WARNING) + except Exception as e: # noqa: BLE001 + warn(f"Could not list offers for the allowed fleets: {e}") + + +def _get_allowed_fleets(api: Client, configuration: PresetConfiguration) -> tuple[str, ...]: + if configuration.fleets is not None: + return tuple(fleet.format() for fleet in configuration.fleets) + fleets = api.client.fleets.list(api.project, include_imported=True) + return tuple( + fleet.name if fleet.project_name == api.project else f"{fleet.project_name}/{fleet.name}" + for fleet in fleets + if fleet.status == FleetStatus.ACTIVE + ) + + +def _build_constraints( + *, + configuration: PresetConfiguration, + build_name: str, + allowed_fleets: Sequence[str], +) -> str: + constraints = PresetConstraints.parse_obj( + { + "run_name_prefix": build_name, + "model": json.loads(configuration.model.json(exclude_none=True)), + "context_length": configuration.context_length, + "max_trials": configuration.max_trials, + "concurrency": configuration.effective_concurrency, + "fleets": list(allowed_fleets), + "env": list(configuration.env), + } + ) + # All fields are always present; unset optional constraints render as null. + return json.dumps(json.loads(constraints.json()), indent=2) + "\n" + + +def _save_final_report_copy( + *, + workspace: PresetAgentWorkspace, + agent_session: PresetAgentSession, + redacted_values: Sequence[str], +) -> None: + if not workspace.final_report_path.exists(): + return + try: + report_text = workspace.final_report_path.read_text(encoding="utf-8", errors="replace") + agent_session.write_final_report(redact(report_text, redacted_values)) + except OSError as e: + warn(f"Could not save a final report copy: {e}") + + +async def _cleanup_runs( + *, + api: Client, + build_name: str, + workspace: PresetAgentWorkspace, + final_run_name: Optional[str], + agent_session: PresetAgentSession, + keep_final_service: bool = False, + wait_for_stop: bool = True, +) -> None: + run_names = _load_submitted_run_names(workspace.runs_path) + if final_run_name is not None: + run_names.append(final_run_name) + run_names = list(dict.fromkeys(run_names)) + expected_prefix = f"{build_name}-" + run_names = [name for name in run_names if name.startswith(expected_prefix)] + if keep_final_service: + run_names = [name for name in run_names if name != final_run_name] + active_names = [] + for name in run_names: + run = api.runs.get(name) + if run is not None and not run.status.is_finished(): + active_names.append(name) + if not active_names: + return + api.client.runs.stop(api.project, active_names, abort=False) + if not wait_for_stop: + # Background reconcile issues the stop but must not block the read + # command it runs inside; the runs terminate on the server regardless. + return + deadline = asyncio.get_running_loop().time() + _RUN_STOP_TIMEOUT_SECONDS + pending = set(active_names) + # The same spinner the stop command shows: without it the CLI looks hung + # for however long the runs take to terminate. + spinner = console.status("Stopping runs...") if agent_session.echo else nullcontext() + with spinner: + while pending: + if asyncio.get_running_loop().time() >= deadline: + raise CLIError(f"Timed out waiting for runs to stop: {', '.join(sorted(pending))}") + for name in list(pending): + run = api.runs.get(name) + if run is None or run.status.is_finished(): + pending.remove(name) + if pending: + await asyncio.sleep(2) + if agent_session.debug: + print_preset_progress("All preset creation runs stopped.", agent_session=agent_session) + + +def _load_submitted_run_names(path: Path) -> list[str]: + try: + lines = path.read_text(encoding="utf-8", errors="replace").splitlines() + except OSError: + return [] + names = [] + for line in lines: + try: + value = json.loads(line) + except json.JSONDecodeError: + continue + if isinstance(value, dict) and isinstance(value.get("name"), str): + name = value["name"].strip() + if name: + names.append(name) + return names diff --git a/src/dstack/_internal/cli/services/presets/output.py b/src/dstack/_internal/cli/services/presets/output.py new file mode 100644 index 000000000..4a9da5252 --- /dev/null +++ b/src/dstack/_internal/cli/services/presets/output.py @@ -0,0 +1,229 @@ +from collections import defaultdict +from datetime import datetime +from typing import Any, Optional + +from rich.table import Table + +from dstack._internal.cli.models.presets import ( + Preset, +) +from dstack._internal.cli.utils.common import add_row_from_dict, console +from dstack._internal.utils.common import pretty_date, pretty_resources + +_STATUS_DISPLAY = { + "ready": ("verified", "grey"), + "running": ("clauding", "bold sea_green3"), + "verifying": ("verifying", "bold deep_sky_blue1"), + "interrupted": ("interrupted", "bold gold1"), + "failed": ("failed", "indian_red1"), +} + + +def _format_status(status: str) -> str: + text, style = _STATUS_DISPLAY.get(status, (status, None)) + return f"[{style}]{text}[/]" if style else text + + +def _trials_exhausted(session: dict[str, Any]) -> bool: + trials = session.get("trials") + max_trials = session.get("max_trials") + return ( + isinstance(trials, dict) + and isinstance(max_trials, int) + and isinstance(trials.get("count"), int) + and trials["count"] >= max_trials + ) + + +def _format_trial_progress(session: Optional[dict[str, Any]]) -> str: + """The ` (N/M)` suffix; stays outside the status markup to render in the + default color.""" + if not isinstance(session, dict): + return "" + trials = session.get("trials") + max_trials = session.get("max_trials") + if not isinstance(trials, dict) or not (trials.get("count") or isinstance(max_trials, int)): + return "" + progress = str(trials.get("count") or 0) + if isinstance(max_trials, int): + progress += f"/{max_trials}" + return f" [secondary]({progress})[/]" + + +def print_presets( + presets: list[Preset], + sessions: Optional[list[dict[str, Any]]] = None, + verbose: bool = False, +) -> None: + console.print(get_presets_table(presets, sessions=sessions, verbose=verbose)) + console.print() + + +def get_presets_table( + presets: list[Preset], + sessions: Optional[list[dict[str, Any]]] = None, + verbose: bool = False, +) -> Table: + table = Table(box=None) + table.add_column("BASE", no_wrap=True) + table.add_column("ID", no_wrap=True) + table.add_column("RESOURCES" if verbose else "GPU", style="secondary") + table.add_column("BENCHMARK", min_width=len("con=1"), overflow="fold") + table.add_column("STATUS", no_wrap=True) + table.add_column("SUBMITTED", no_wrap=True, style="secondary") + table.add_column("NAME", no_wrap=True, style="secondary") + presets_by_base: dict[str, list[Preset]] = defaultdict(list) + repo_to_base: dict[str, str] = {} + for preset in presets: + presets_by_base[preset.base].append(preset) + repo_to_base[preset.model] = preset.base + sessions_by_model: dict[str, list[dict[str, Any]]] = defaultdict(list) + creations_by_id: dict[str, dict[str, Any]] = {} + for session in sessions or []: + if str(session.get("status")) == "success": + # A completed creation session decorates its preset row. + creations_by_id[str(session.get("id"))] = session + continue + model = str(session.get("model") or "unknown") + sessions_by_model[repo_to_base.get(model, model)].append(session) + + for base in sorted({*presets_by_base, *sessions_by_model}, key=str.lower): + add_row_from_dict(table, {"BASE": base}) + # Newest first within a group, as in `dstack ps`. + for preset in sorted( + presets_by_base.get(base, []), key=lambda p: p.created_at, reverse=True + ): + _add_preset(table, preset, verbose=verbose, creation=creations_by_id.get(preset.id)) + for session in sorted( + sessions_by_model.get(base, []), + key=lambda s: str(s.get("created_at") or ""), + reverse=True, + ): + _add_session(table, session) + return table + + +def _add_session(table: Table, session: dict[str, Any]) -> None: + created = "" + created_at = session.get("created_at") + if isinstance(created_at, str): + try: + created = pretty_date(datetime.fromisoformat(created_at)) + except ValueError: + created = created_at + benchmark = "" + gpu = "" + status_key = str(session.get("status", "")) + if status_key == "running" and _trials_exhausted(session): + # The trial budget is spent, so the agent is deploying and verifying + # the final service. + status_key = "verifying" + status = _format_status(status_key) + _format_trial_progress(session) + trials = session.get("trials") + if isinstance(trials, dict): + best = trials.get("best") + if isinstance(best, dict): + parts = ["best trial:"] + if best.get("concurrency"): + parts.append(f"con={best['concurrency']}") + parts.append(f"{_format_number(best['tok_s'])} tok/s") + benchmark = " ".join(parts) + gpu = best.get("gpu") or "" + add_row_from_dict( + table, + { + "ID": str(session.get("id", "")), + "NAME": str(session.get("name") or ""), + "GPU": gpu, + "RESOURCES": gpu, + "BENCHMARK": benchmark, + "STATUS": status, + "SUBMITTED": created, + }, + ) + + +def _add_preset( + table: Table, + preset: Preset, + *, + verbose: bool, + creation: Optional[dict[str, Any]] = None, +) -> None: + groups = preset.service.replica_groups + column = "RESOURCES" if verbose else "GPU" + row = { + "ID": preset.id, + "NAME": preset.name or "", + column: _format_resources(groups[0].resources, verbose=verbose), + "STATUS": _format_status("ready") + _format_trial_progress(creation), + "BENCHMARK": format_preset_benchmark(preset, verbose=verbose), + "SUBMITTED": pretty_date(preset.created_at), + } + if verbose and preset.model != preset.base: + row["BASE"] = f"[secondary] repo={preset.model}[/]" + add_row_from_dict(table, row) + if len(groups) > 1: + for group in groups: + add_row_from_dict( + table, + { + "BASE": f" group={group.name}", + column: _format_resources(group.resources, verbose=verbose), + }, + style="secondary", + ) + + +def format_preset_benchmark(preset: Preset, *, verbose: bool = False) -> str: + benchmark = preset.validations[0].benchmark + workload = benchmark.workload + metrics = benchmark.metrics + output_tokens_per_second = metrics.total_output_tokens / metrics.duration_seconds + parts = [ + f"con={workload.concurrency}", + f"{_format_number(output_tokens_per_second)} tok/s", + f"TTFT {_format_latency(metrics.ttft_ms.p50)}", + ] + if verbose: + parts.insert(0, f"ctx={_format_token_count(preset.context_length)}") + return " ".join(parts) + + +def _format_token_count(value: int) -> str: + for divisor, suffix in ((1024 * 1024, "M"), (1024, "K")): + if value >= divisor and value % divisor == 0: + return f"{value // divisor}{suffix}" + return str(value) + + +def _format_number(value: float) -> str: + # `g` switches to scientific notation for values >= 1000 (after rounding). + if abs(value) >= 999.5: + return f"{value:.0f}" + return f"{value:.3g}" + + +def _format_latency(value_ms: float) -> str: + if value_ms >= 1000: + return f"{_format_number(value_ms / 1000)}s" + return f"{_format_number(value_ms)}ms" + + +def _format_resources(resources, *, verbose: bool) -> str: + if resources is None: + return "-" + if verbose: + return resources.pretty_format() + gpu = resources.gpu + if gpu is None or gpu.count.max == 0: + return "-" + formatted = pretty_resources( + gpu_vendor=gpu.vendor, + gpu_name=",".join(gpu.name) if gpu.name else None, + gpu_count=gpu.count, + gpu_memory=gpu.memory, + total_gpu_memory=gpu.total_memory, + compute_capability=gpu.compute_capability, + ) + return formatted.removeprefix("gpu=") or "-" diff --git a/src/dstack/_internal/cli/services/endpoints/presets.py b/src/dstack/_internal/cli/services/presets/presets.py similarity index 89% rename from src/dstack/_internal/cli/services/endpoints/presets.py rename to src/dstack/_internal/cli/services/presets/presets.py index 727133fa3..3b4db39ee 100644 --- a/src/dstack/_internal/cli/services/endpoints/presets.py +++ b/src/dstack/_internal/cli/services/presets/presets.py @@ -1,14 +1,14 @@ import hashlib import json -from typing import Any +from typing import Any, Optional import gpuhunt -from dstack._internal.cli.models.endpoint_presets import ( - EndpointBenchmark, - EndpointPreset, - EndpointPresetValidation, - EndpointPresetValidationReplica, +from dstack._internal.cli.models.presets import ( + Preset, + PresetBenchmark, + PresetValidation, + PresetValidationReplica, ) from dstack._internal.core.models.configurations import ServiceConfiguration from dstack._internal.core.models.envs import EnvSentinel @@ -18,28 +18,31 @@ from dstack._internal.utils.common import format_mib_as_gb, get_current_datetime -def build_endpoint_preset( +def build_preset( *, service: ServiceConfiguration, - validation_replicas: list[EndpointPresetValidationReplica], + validation_replicas: list[PresetValidationReplica], base_model: str, model: str, context_length: int, - benchmark: EndpointBenchmark, -) -> EndpointPreset: + benchmark: PresetBenchmark, + preset_id: Optional[str] = None, + name: Optional[str] = None, +) -> Preset: service = service.copy(deep=True) service.name = None service.gateway = None for field in ProfileParams.__fields__: setattr(service, field, None) - validation = EndpointPresetValidation( + validation = PresetValidation( replicas=validation_replicas, benchmark=benchmark, ) set_service_gpu_vendors_from_validations(service, [validation]) - return EndpointPreset( + return Preset( + name=name, base=base_model, - id=make_endpoint_preset_id(service, context_length=context_length), + id=preset_id or make_preset_id(service, context_length=context_length), model=model, context_length=context_length, created_at=get_current_datetime(), @@ -48,7 +51,7 @@ def build_endpoint_preset( ) -def make_endpoint_preset_id( +def make_preset_id( service: ServiceConfiguration, context_length: int, ) -> str: @@ -63,10 +66,11 @@ def make_endpoint_preset_id( return hashlib.sha256(payload.encode()).hexdigest()[:8] -def endpoint_preset_to_data(preset: EndpointPreset) -> dict[str, Any]: +def preset_to_data(preset: Preset) -> dict[str, Any]: return { "base": preset.base, "id": preset.id, + **({"name": preset.name} if preset.name else {}), "model": preset.model, "context_length": preset.context_length, "created_at": preset.created_at.isoformat(), @@ -115,7 +119,7 @@ def resources_spec_from_instance_resources(resources: Resources) -> ResourcesSpe or gpu.vendor != first_gpu.vendor for gpu in resources.gpus ): - raise ValueError("endpoint preset cannot be built from mixed-GPU instances") + raise ValueError("preset cannot be built from mixed-GPU instances") data["gpu"] = { "name": first_gpu.name, "memory": format_mib_as_gb(first_gpu.memory_mib), @@ -130,7 +134,7 @@ def resources_spec_from_instance_resources(resources: Resources) -> ResourcesSpe def set_service_gpu_vendors_from_validations( service: ServiceConfiguration, - validations: list[EndpointPresetValidation], + validations: list[PresetValidation], ) -> None: for group_num, group in enumerate(service.replica_groups): resources = group.resources @@ -153,7 +157,7 @@ def _env_item_to_preset_data(key: str, value: str | EnvSentinel) -> str: def _get_validation_group_gpu_vendor( - validations: list[EndpointPresetValidation], + validations: list[PresetValidation], group_num: int, ) -> gpuhunt.AcceleratorVendor | None: vendors = { diff --git a/src/dstack/_internal/cli/services/presets/prompt.py b/src/dstack/_internal/cli/services/presets/prompt.py new file mode 100644 index 000000000..0d0052c32 --- /dev/null +++ b/src/dstack/_internal/cli/services/presets/prompt.py @@ -0,0 +1,36 @@ +import re +from pathlib import Path +from typing import Optional + +from dstack._internal.core.errors import CLIError + +_SYSTEM_PROMPT_PATH = Path(__file__).resolve().parent / "resources" / "system_prompt.md" + +# `` emits CONTENT (with `{NAME}` interpolated) when the +# variable NAME is set, and nothing otherwise. The conditional text lives in +# the document; this module only applies the rule. +_DIRECTIVE_PATTERN = re.compile(r"", re.DOTALL) + + +# TODO: reintroduce a `# Resume` section in system_prompt.md once session resume +# (seeded from `runs.jsonl` and `trials.jsonl`) is designed. +def get_preset_agent_system_prompt(user_prompt: Optional[str] = None) -> str: + text = _SYSTEM_PROMPT_PATH.read_text(encoding="utf-8").strip() + variables = {"prompt": user_prompt.strip() if user_prompt else None} + applied = 0 + + def substitute(match: re.Match) -> str: + nonlocal applied + name, content = match.group(1), match.group(2) + if name not in variables: + raise CLIError(f"Unknown variable {name!r} in the agent system prompt") + value = variables[name] + if not value: + return "" + applied += 1 + return content.replace("{" + name + "}", value) + + rendered = _DIRECTIVE_PATTERN.sub(substitute, text) + if variables["prompt"] and not applied: + raise CLIError("The agent system prompt has no place for the user prompt") + return re.sub(r"\n{3,}", "\n\n", rendered) diff --git a/src/dstack/_internal/cli/services/presets/redaction.py b/src/dstack/_internal/cli/services/presets/redaction.py new file mode 100644 index 000000000..f62738810 --- /dev/null +++ b/src/dstack/_internal/cli/services/presets/redaction.py @@ -0,0 +1,61 @@ +"""Secret redaction for everything copied out of the agent workspace.""" + +import os +from typing import Any, Sequence + +_REDACTION = "[redacted]" +# Replacing shorter values such as "1" or "false" corrupts unrelated diagnostics. +_MIN_REDACTED_SUBSTRING_LENGTH = 8 +_SENSITIVE_INHERITED_ENV_NAMES = ( + "HTTP_PROXY", + "HTTPS_PROXY", + "NO_PROXY", + "http_proxy", + "https_proxy", + "no_proxy", +) + + +def get_redacted_values(values: Sequence[str]) -> tuple[str, ...]: + return tuple(sorted({value for value in values if value}, key=len, reverse=True)) + + +def contains_redacted_value(value: Any, redacted_values: Sequence[str]) -> bool: + if isinstance(value, str): + return any( + value == redacted + or (len(redacted) >= _MIN_REDACTED_SUBSTRING_LENGTH and redacted in value) + for redacted in redacted_values + ) + if isinstance(value, dict): + return any(contains_redacted_value(item, redacted_values) for item in value.values()) + if isinstance(value, list): + return any(contains_redacted_value(item, redacted_values) for item in value) + return False + + +def get_sensitive_inherited_env_values() -> list[str]: + return [value for name in _SENSITIVE_INHERITED_ENV_NAMES if (value := os.getenv(name))] + + +def redact(value: str, redacted_values: Sequence[str]) -> str: + for redacted_value in redacted_values: + if value == redacted_value: + return _REDACTION + if len(redacted_value) >= _MIN_REDACTED_SUBSTRING_LENGTH: + value = value.replace(redacted_value, _REDACTION) + return value + + +def redact_structure(value: Any, redacted_values: Sequence[str]) -> Any: + """Recursively redacts every string (including dict keys) in a JSON-like value.""" + if isinstance(value, str): + return redact(value, redacted_values) + if isinstance(value, list): + return [redact_structure(item, redacted_values) for item in value] + if isinstance(value, dict): + return { + redact(key, redacted_values): redact_structure(item, redacted_values) + for key, item in value.items() + } + return value diff --git a/src/dstack/_internal/cli/services/presets/resources/system_prompt.md b/src/dstack/_internal/cli/services/presets/resources/system_prompt.md new file mode 100644 index 000000000..238709970 --- /dev/null +++ b/src/dstack/_internal/cli/services/presets/resources/system_prompt.md @@ -0,0 +1,382 @@ +# Objective + +Your goal is to find a model serving configuration with the best +performance through sequential experimental trials. Once the best-performing +candidate is found, it is deployed as a `dstack` service for the final +benchmark and saved as a reusable preset. + +# Constraints + +`constraints.json` in the workspace root contains the effective constraints +for this session. No submitted run or final service may conflict with them. + +Field semantics: + +- `run_name_prefix`: the required prefix for all submitted run names, + followed by the run counter; see `# Runs`. +- `model`: the model to serve. If it has `repo`, deploy that repo/path + exactly. If it has `base`, choose a compatible repo/path that best fits + performance and hardware within the constraints: the base repo itself, a + different precision or quantization, or another trusted compatible repo. + The client-facing model name of the final service is `model.name` when + set, otherwise `model.repo` or `model.base`. +- `context_length`: the minimum context length the selected repo/path and + the final service must support. `null` means no minimum is required. +- `max_trials`: the maximum number of trials in this session. +- `concurrency`: the number of simultaneous requests for every benchmark in + this session. It is fixed so that benchmark results are comparable. +- `fleets`: use these existing `dstack` fleets only. Do not create, delete, + apply, or edit fleets. +- `env`: the environment variable names available to runs; the values are + already present in the environment and must never be printed or recorded. + +During the trials and experimentation aimed at the best performance, you may +pick the hardware (the best available within the allowed `dstack` fleets), +the model variant (only if `model` has `base`), the serving framework, the +Docker image and dependencies, the serving framework parameters, and +anything else within these constraints — except generating custom kernels, +patching drivers, or patching serving framework source code. + + +## CLI And Skills + +All trials and the final verification are done using `dstack`. This includes +fetching information about `dstack` fleets, offers, and runs, as well as +submitting `dstack` runs (e.g. tasks for trials and services for the final +verification). + +To do this, use the real `dstack` CLI and shell commands in this workspace. +Load and follow `/dstack` for `dstack` CLI/YAML syntax. Load and follow +`/dstack-prototyping` for how to test a model-serving configuration with +`dstack` tasks before verifying it as a `dstack` service. If a skill cannot +be loaded with its +slash command, read it from `.claude/skills//SKILL.md` and +follow it the same way. + +# Workspace Files + +Files provided to you in the workspace root (read them; never edit them): + +- `constraints.json`: the effective constraints; see `# Constraints`. + +Files you are expected to maintain in the workspace root: + +- `runs.jsonl`: the append-only record of submitted `dstack` runs; see + `# Runs`. +- `progress.jsonl`: progress messages, written through the `progress` helper; + see `# Progress`. +- `trials.jsonl`: the append-only record of completed trials; see `# Trials`. +- `final_report.json`: the final report; see `# Final Report`. + +You may create any other working files (run YAML files, benchmark output, +notes) in the workspace, and only there: do not deliberately save files +elsewhere on this machine. Incidental writes made by the tools you run +(caches, temporary files, SSH configuration) are fine wherever those tools +keep them. Files inside running `dstack` tasks or services are not subject +to this rule. + +# Runs + +If you submit `dstack` runs via the `dstack` CLI, e.g. to create tasks or +services, name them as described below and record them in `runs.jsonl`. + +Use only `-` as the run name, with +`run_name_prefix` from `constraints.json` and `` counting all +runs submitted in this session: the first submitted run is +`-1`, the second is `-2`, and so on. The +next counter value is the number of lines in `runs.jsonl` plus one. Do not +add framework, hardware, role, or purpose suffixes to run names; record those +details in workspace files and progress. + +`runs.jsonl` is append-only. Immediately after submitting a run, fetch its +id with `dstack run get --json` and append one JSON line: + +```json +{"name":"qwen-preset-1","id":""} +``` + +Never edit or delete existing lines. Stop runs you no longer need unless they +are still needed for attach/SSH debugging, logs, or backend diagnosis. + +After stopping a `dstack` task or service, follow `/dstack` structured status +guidance and confirm that the run reached a terminal status before continuing. + +# Progress + +From the very beginning to the very end — through every trial, the service +verification, and each transition between them — report all major +intentions, decisions, and results, concisely but clearly, to +`progress.jsonl`. Write each message with the `progress` helper; it appends +to `progress.jsonl`, and the caller shows only the message text: + +```bash +progress "" +``` + +Make sure progress messages explain your choices. This explicitly includes, +but is not limited to, the choice of fleet, backend, hardware, model repo, +and serving framework: name what you chose, why, what evidence you used, and +what you rejected. + +Do not put raw YAML, command output, long tables, traces, or secrets in +progress messages. + +# Trials (Main Section) + +The end goal is to find a model serving configuration that matches the +constraints and delivers the best performance. The search is done via +so-called trials, where each trial is formed around a substantive idea on +how to get better performance than the previous trials. Do not consider +P/D disaggregation setups yet. + + + +Trial ideas must not rely only on what you already know. Research how to +get the best performance for the chosen model, serving framework, and +hardware in trustworthy sources. Start with these: + +- vLLM recipes: `https://recipes.vllm.ai/` (model index: + `https://recipes.vllm.ai/models.json`) +- SGLang docs: `https://docs.sglang.io/` (fetch `/llms.txt` for the page + index) +- SGLang model recipes: `https://docs.sglang.io/cookbook/autoregressive/intro` +- Release notes: `https://github.com/vllm-project/vllm/releases` and + `https://github.com/sgl-project/sglang/releases` +- Performance-loop methodology (profiling, benchmark contracts): + `https://www.lmsys.org/blog/2026-07-02-agent-assisted-sglang-development` + +Go beyond this list proactively — official docs, repo issues, and reputable +benchmarks — whenever that can help the trial. Research before the first +trial and whenever a benchmark exposes a bottleneck. + +For each trial, use `dstack` tasks (see `# Task Usage`). During a trial, run +commands interactively inside the task (over SSH) and measure the +performance when needed, following `## Benchmark` below. + +If a trial needs another Docker image, stop the task, re-submit it, and +continue the same trial. This applies equally when a new trial's idea needs +a different image or serving framework: the cost of re-submitting a task is +not a reason to keep the current image or framework. You decide when a +trial is complete: better performance was achieved, or the ideas within this +trial are exhausted. + +Once a trial is completed, compile the interactive commands that produced +the final performance into a complete `dstack` task configuration with exact +commands, and log it together with the corresponding benchmark results (see +`## Benchmark` for the structure) to `trials.jsonl`. The benchmark may +be skipped in one case only: you failed to make the configuration run at +all — a failed trial. + +Each `trials.jsonl` record is one JSON line with exactly three fields: + +``` +{"task": {...}, "resources": {...}, "benchmark": {...}} +``` + +- `task`: the compiled `dstack` task configuration described above, as JSON. + Its `name` is the run name of the trial's task, and its `commands` are the + exact final commands that led to the benchmark results — the commands that + are supposed to replicate the benchmark results exactly if the task is + submitted (not `sleep infinity`). +- `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 + actual values from the latest job submission's + `job_runtime_data.offer.instance.resources` in + `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`. +- `benchmark`: the trial benchmark (see `## Benchmark` for the structure); + `null` only for a failed trial. + +Continue the trials until `max_trials` from `constraints.json` is reached or +you are entirely out of ideas on how to reach better performance. Before +stopping early, step back and think once more about what could still improve +performance within the constraints (see `# Constraints`). An early stop must +be justified in +`progress.jsonl` (see `# Progress`): report what you considered and why none +of it is worth a trial. + +Once the trials are over, pick the best trial and deploy it as a `dstack` +service to verify that it works and benchmark it finally (see +`# Final Service`). If +all is good, write the final report `final_report.json` (see +`# Final Report`). If the trial cannot be reproduced, pick the best trial +you have not verified yet and try again; proceed until a trial succeeds or +no trials remain. In that case, log the failure to `final_report.json` (see +`# Final Report`). + +## Benchmark + +During trials, run benchmarks via SSH inside the task, directly against the +serving engine: use `concurrency` from `constraints.json` and measure all +trials the same way so that their results are comparable with each other. +In trial benchmarks too, all measured requests must succeed. + +Before any benchmark — a trial one or the final one — warm the engine up by +verifying that the model works as expected: send real requests and check +the responses, including reasoning output when the model supports it. These +verification requests are never part of the measured metrics. + +Record every benchmark using the following structure and field names — +trial benchmarks in `trials.jsonl`, the final benchmark as +`final_report.json.benchmark` (values are illustrative): + +```json +{ + "tool": "vllm bench serve", + "tool_version": "0.11.0", + "command": "vllm bench serve ...", + "workload": {"api": "chat_completions", "num_requests": 16, "input_tokens": 1024, "output_tokens": 128, "concurrency": 8}, + "metrics": { + "successful_requests": 16, "failed_requests": 0, "duration_seconds": 4.0, + "total_input_tokens": 16384, "total_output_tokens": 2048, + "ttft_ms": {"mean": 110.9, "p50": 108.2, "p99": 121.6}, + "tpot_ms": {"mean": 7.5, "p50": 7.4, "p99": 8.1} + } +} +``` + +Set `tool` to the command name and subcommands without options or values, +`tool_version` to the exact version, and `command` to the secret-free +invocation. For the final benchmark, run it with streaming responses, set +`workload.concurrency` to `concurrency` from `constraints.json`, produce +every field of the structure, and calculate all metrics from the +`num_requests` measured requests only — exclude setup, health-check, and +warmup requests; all measured requests must succeed. Never invent missing +values. + +# 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 +run commands inside the task interactively, via SSH. It is important that +you follow the `/dstack-prototyping` skill when working with tasks. + +# Hardware + +Pick the hardware before the trials start. You are allowed to change the +hardware from one trial to another if this can help the outcome. + +The available hardware is defined by the allowed `dstack` fleets and their +offers (coming from the configured backends). Follow the +`/dstack-prototyping` skill on how to efficiently select offers among the +available backends or SSH fleets. Pick the offer whose hardware best fits the trial's idea. Only when several +offers fit comparably, prefer backends or SSH fleets that support idle +instances/instance volumes: later runs reuse the instance and cached model +weights, while container-based backends start clean on every re-submission. + +If the backend allows, use instance volumes to mount cache and model weights +between runs. + +When submitting a `dstack` task or service, pass exact `fleets`, `backends`, +and an intentional `resources` range based on the choice made from offers, so +the run does not land outside the intended fleet/backend/hardware. + +## Fleet Offers + +When selecting `dstack` backends/fleets and evaluating hardware that can be +used to run the model, use `dstack offer --json` and pass the fleets from +`constraints.json` with `--fleet` explicitly. If `--fleet` is not passed, +`dstack offer` can show offers that are not applicable to the allowed fleets. +Use these offers when selecting fleet, backend, and hardware. + +To classify each backend's capabilities, fetch `https://dstack.ai/docs/concepts/backends.md` and classify from the fetched document, not from memory. VM-based backends are listed under `## VM-based` (they support idle instances and instance volumes). Kubernetes backend is listed under `## Container-based`, but supports instance volumes and thus is preferred over other container-based backends. +SSH fleets can be treated as VM-based backends as they support both idle instances (its equivalent) and instance volumes. + +# Final Service + +Once the trials are over, pick the best trial that has not been verified yet +and submit its configuration as a `dstack` service. Make it work with only +minor tweaks if needed; do not change the important decisions made during +the trial. Set the service `model` name to the client-facing model name from +`constraints.json` (see `# Constraints`). `model` is required: it also enables +`dstack`'s default readiness probe — an independent check that the model can +serve requests. If the service never passes the probe, treat that as a real +failure of the configuration, not something to work around by removing `model`. + +Before the final benchmark, verify the model through the service: send real +requests using the client-facing model name and check that the model works +as it should, including reasoning output when the model supports it. This +verification also warms the service up. Only then run the final benchmark +(see `## Benchmark`). When verifying or benchmarking the service, use its +`service.url` reported by `dstack run get --json`, along with +`DSTACK_TOKEN` as the bearer token. If `service.url` is a relative path, +prepend `DSTACK_SERVER_URL` to build the absolute URL. + +During the service verification, test the context length the service +actually supports and report it as `final_report.json.context_length`. + +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 +`# Final Report`). + +When you report success, leave the final `dstack` service running: the +caller verifies the live run and stops it afterwards. + +# Secrets + +The `dstack` CLI is already configured. Do not inspect `~/.dstack/config.yml` or +print, copy, or summarize tokens, secrets, or environment variable values. + +Do not expose the value of `DSTACK_TOKEN` or the value of any environment +variable listed under `env` in `constraints.json`. + +Do not put secret values in `final_report.json` or print them. Use env +references in `final_report.json.service_yaml`; use environment variable names or redacted values in +`final_report.json.benchmark.command`. + +# Final Report + +`final_report.json` may contain only `success`, `run_id`, `run_name`, +`service_yaml`, `base`, `model`, `context_length`, `benchmark`, and +`failure_summary`. + +On success, include exactly: + +- `success`: `true` +- `run_id`: the final verified service run ID +- `run_name`: the final verified service run name +- `service_yaml`: the full YAML of the verified final service +- `base`: the base model repo, determined by the rules below +- `model`: the exact repo/path loaded by the final service command +- `context_length`: the context length verified for the final service +- `benchmark`: the final service benchmark described in `## Benchmark` + +Set `final_report.json.base` as follows: + +- If `model` in `constraints.json` has `base`, set `final_report.json.base` to + that value. +- If `model` has `repo`, inspect the repo metadata, model card, config, or + another reliable source to identify the base model repo. +- If `model.repo` is itself the base model repo, set `final_report.json.base` + to `model.repo`. +- Do not infer `final_report.json.base` only from the repo name. + +On failure, include exactly: + +- `success`: `false` +- `failure_summary`: the reason a preset could not be created and any change + required from the user or administrator + +Write the report to `final_report.json`, then submit the identical JSON object +through `StructuredOutput`. + +Verify that `final_report.json` is correct and matches the required schema. + +Stop only after `final_report.json` is written and submitted: either one +final `dstack` service was verified and benchmarked, or the trials and +unverified candidates were exhausted (see `# Trials` and `# Final Service`). + +Ending your turn stops the session even while background commands are still +running. Wait for long-running work — weight downloads, engine startup, +benchmarks, provisioning — by polling it with short blocking commands, +never by ending your turn. diff --git a/src/dstack/_internal/cli/services/presets/session.py b/src/dstack/_internal/cli/services/presets/session.py new file mode 100644 index 000000000..72ec86e0f --- /dev/null +++ b/src/dstack/_internal/cli/services/presets/session.py @@ -0,0 +1,534 @@ +"""Preset creation sessions: on-disk state, ownership, and liveness.""" + +import json +import os +import secrets +import shutil +import tempfile +import time +from contextlib import suppress +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import TYPE_CHECKING, Any, Iterator, Optional + +import psutil +import yaml +from rich.text import Text + +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.models.preset_agent import ClaudeAgentInfo +from dstack._internal.cli.utils.common import console +from dstack._internal.compat import IS_WINDOWS +from dstack._internal.core.errors import CLIError +from dstack._internal.utils.common import get_dstack_dir + +if TYPE_CHECKING: + from dstack._internal.cli.services.presets.agent import ClaudeAuth + + +_PROGRESS_FILENAME = "progress.jsonl" +_RUNS_FILENAME = "runs.jsonl" +_TRIALS_FILENAME = "trials.jsonl" +_CONSTRAINTS_FILENAME = "constraints.json" +_FINAL_REPORT_FILENAME = "final_report.json" +_SESSION_FILENAME = "session.json" +_USER_PROMPT_FILENAME = "user_prompt.md" + + +class SessionBusyError(CLIError): + """Another live process owns the session — it is following or finalizing it. + Callers that only want to view can fall back to a read-only follow.""" + + +@dataclass +class PresetAgentSession: + path: Path + debug: bool + preset_id: str = "" + # Whether progress lines echo to this process's console (a live attach), on + # top of always being recorded to agent.log. Background reconcile sets it + # False so finalizing a detached session stays silent on the read command. + echo: bool = field(default=True, repr=False) + _log_enabled: bool = field(default=True, init=False, repr=False) + + @property + def log_path(self) -> Path: + return self.path / "agent.log" + + @property + def trace_path(self) -> Path: + return self.path / "trace.jsonl" + + @property + def runs_path(self) -> Path: + return self.path / _RUNS_FILENAME + + @property + def trials_path(self) -> Path: + return self.path / _TRIALS_FILENAME + + def write_prompt(self, prompt: str) -> None: + _write_private_text(self.path / "prompt.md", prompt + "\n") + + def write_user_prompt(self, user_prompt: str) -> None: + _write_private_text(self.path / _USER_PROMPT_FILENAME, user_prompt + "\n") + + def read_user_prompt(self) -> Optional[str]: + try: + text = (self.path / _USER_PROMPT_FILENAME).read_text(encoding="utf-8").strip() + except OSError: + return None + return text or None + + def write_constraints(self, constraints_text: str) -> None: + _write_private_text(self.path / _CONSTRAINTS_FILENAME, constraints_text) + + def write_final_report(self, report_text: str) -> None: + _write_private_text(self.path / _FINAL_REPORT_FILENAME, report_text) + + def write_agent_info(self, auth: "ClaudeAuth") -> None: + from dstack._internal.cli.services.presets.agent import ( + _get_claude_auth_status, + _get_claude_version, + ) + + info = ClaudeAgentInfo.parse_obj( + { + "executable": auth.executable, + "version": _get_claude_version(auth), + "model": { + "name": auth.model, + "effort": auth.effort or "default", + }, + "auth": _get_claude_auth_status(auth), + } + ) + _write_private_text( + self.path / "agent.json", json.dumps(json.loads(info.json()), indent=2) + "\n" + ) + + def append_log(self, line: str) -> None: + if not self._log_enabled: + return + try: + with self.log_path.open("a", encoding="utf-8") as f: + f.write(line + "\n") + f.flush() + except OSError as e: + self._log_enabled = False + if self.echo: + console.print(f"[warning]Could not write agent log {self.log_path}: {e}[/]") + + def read_manifest(self) -> dict[str, Any]: + try: + manifest = json.loads((self.path / _SESSION_FILENAME).read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return {} + return manifest if isinstance(manifest, dict) else {} + + def update_manifest(self, **fields: Any) -> None: + manifest = self.read_manifest() + manifest.update(fields) + _write_private_text(self.path / _SESSION_FILENAME, json.dumps(manifest, indent=2) + "\n") + + def record_claude_session_id(self, session_id: str) -> None: + self.update_manifest(claude_session_id=session_id) + + def finish(self, status: str) -> Path: + self.update_manifest(status=status) + return self.path + + +def get_presets_dir() -> Path: + return get_dstack_dir() / "presets" + + +def create_preset_agent_session( + configuration: PresetConfiguration, + *, + debug: bool = False, +) -> PresetAgentSession: + if configuration.name is None: + raise CLIError("The service name is required to save agent output") + parent = get_presets_dir() + path: Optional[Path] = None + try: + parent.mkdir(mode=0o700, parents=True, exist_ok=True) + while True: + preset_id = secrets.token_hex(4) + path = parent / preset_id + try: + path.mkdir(mode=0o700) + except FileExistsError: + continue + break + _write_private_text(path / "agent.log", "") + manifest = { + "id": preset_id, + "status": "running", + "pid": os.getpid(), + "pid_started_at": _process_started_at(os.getpid()), + "name": configuration.name, + "model": getattr(configuration.model, "base", None) + or getattr(configuration.model, "repo", None), + "max_trials": configuration.max_trials, + "created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"), + "debug": debug, + } + _write_private_text(path / _SESSION_FILENAME, json.dumps(manifest, indent=2) + "\n") + data = json.loads(configuration.json(exclude_none=True)) + if configuration.env: + data["env"] = list(configuration.env) + else: + data.pop("env", None) + _write_private_text( + path / "preset.dstack.yml", + yaml.safe_dump(data, sort_keys=False), + ) + if debug: + _write_private_text(path / "trace.jsonl", "") + except OSError as e: + if path is not None: + shutil.rmtree(path, ignore_errors=True) + raise CLIError(f"Could not create agent output under {parent}: {e}") from e + assert path is not None + return PresetAgentSession(path=path, debug=debug, preset_id=preset_id) + + +def load_resumable_agent_session(preset_id: str) -> PresetAgentSession: + path = get_presets_dir() / preset_id + session = PresetAgentSession(path=path, debug=False, preset_id=preset_id) + manifest = session.read_manifest() + if not path.is_dir() or not manifest: + raise CLIError(f"Unknown preset: {preset_id}") + status = manifest.get("status") + if status == "success": + raise CLIError(f"Preset {preset_id} is already created; nothing to resume") + if status == "failed": + raise CLIError(f"Preset {preset_id} creation failed and cannot be resumed") + if status == "running" and session_process_alive(manifest): + raise CLIError( + f"Preset {preset_id} is still being created;" + f" follow it with dstack preset logs -f {preset_id}" + ) + if not manifest.get("claude_session_id"): + raise CLIError(f"Preset {preset_id} creation stopped before it started; create a new one") + session.debug = bool(manifest.get("debug")) + return session + + +def _pid_alive(pid: Any, started_at: Any = None) -> bool: + if not isinstance(pid, int) or pid <= 0 or not psutil.pid_exists(pid): + return False + if isinstance(started_at, (int, float)): + create_time = _process_started_at(pid) + # A recycled pid has a different start time. + if create_time is not None and abs(create_time - started_at) > 1.0: + return False + return True + + +def session_process_alive(manifest: dict[str, Any]) -> bool: + """Whether the session is still worked on: a live agent (possibly + detached) or a live CLI (possibly between agent retries).""" + if _pid_alive(manifest.get("agent_pid"), manifest.get("agent_started_at")): + return True + pid = manifest.get("pid") + if not isinstance(pid, int) or pid <= 0 or pid == os.getpid(): + return False + # Guard the CLI pid with its start time too: after an ungraceful CLI death + # the OS can recycle the pid, and a bare pid_exists() would read a dead + # session as still owned (falsely blocking reconcile / follow). + return _pid_alive(pid, manifest.get("pid_started_at")) + + +def load_attachable_agent_session(preset_id: str) -> PresetAgentSession: + path = get_presets_dir() / preset_id + session = PresetAgentSession(path=path, debug=False, preset_id=preset_id) + manifest = session.read_manifest() + if not path.is_dir() or not manifest: + raise CLIError(f"Unknown preset: {preset_id}") + status = manifest.get("status") + if status == "success": + raise CLIError(f"Preset {preset_id} is already created") + if status == "failed": + raise CLIError(f"Preset {preset_id} creation failed") + if status == "interrupted": + raise CLIError( + f"Preset {preset_id} creation was interrupted; resume it with" + f" dstack preset create -f --resume {preset_id}" + ) + pid = manifest.get("pid") + if ( + isinstance(pid, int) + and pid > 0 + and pid != os.getpid() + and _pid_alive(pid, manifest.get("pid_started_at")) + ): + raise SessionBusyError( + f"Preset {preset_id} is already being followed by another CLI (pid {pid});" + f" stop or detach it there with Ctrl+C" + ) + session.debug = bool(manifest.get("debug")) + return session + + +def load_agent_session(preset_id: str) -> PresetAgentSession: + """Loads a session of any status for read-only inspection (its log).""" + path = get_presets_dir() / preset_id + session = PresetAgentSession(path=path, debug=False, preset_id=preset_id) + if not path.is_dir() or not session.read_manifest(): + raise CLIError(f"Unknown preset: {preset_id}") + return session + + +def print_session_log(session: PresetAgentSession) -> None: + """Prints the session's redacted progress log verbatim (no markup).""" + try: + content = session.log_path.read_text(encoding="utf-8") + except OSError: + content = "" + if content.strip(): + console.print(Text(content.rstrip("\n")), soft_wrap=True) + else: + console.print(f"No log output yet for session [code]{session.preset_id}[/].") + + +def mark_session_owner( + session: PresetAgentSession, + *, + project: Optional[str] = None, + keep_service: Optional[bool] = None, + claude_model: Optional[str] = None, +) -> None: + """Records this process as the session's owner (pid + start time) and, when + given, the finalize context a later detached reconcile needs (project and + keep-service intent). `None` fields are left untouched.""" + fields: dict[str, Any] = { + "status": "running", + "pid": os.getpid(), + "pid_started_at": _process_started_at(os.getpid()), + } + if project is not None: + fields["project"] = project + if keep_service is not None: + fields["keep_service"] = keep_service + if claude_model is not None: + fields["claude_model"] = claude_model + session.update_manifest(**fields) + + +def session_report_exists(manifest: dict[str, Any]) -> bool: + """Whether the agent left a final report on disk — the durable completion + signal a detached session is finalized from.""" + workspace = manifest.get("workspace") + if not isinstance(workspace, str) or not workspace: + return False + return (Path(workspace) / "w" / _FINAL_REPORT_FILENAME).is_file() + + +def try_claim_session(session: PresetAgentSession) -> Optional[int]: + """Takes an exclusive, kernel-held lock for the duration of a session's + finalization, so concurrent readers can't both finalize it. Returns an open + file descriptor to release via `release_session_claim`, or None if another + process holds it. The kernel drops the lock if the holder dies, so there are + no stale locks to reason about.""" + try: + fd = os.open(session.path / ".reconcile.lock", os.O_CREAT | os.O_RDWR, 0o600) + except OSError: + return None + if _try_lock_fd(fd): + return fd + with suppress(OSError): + os.close(fd) + return None + + +def release_session_claim(fd: Optional[int]) -> None: + if fd is not None: + # Closing the descriptor releases the kernel lock. + with suppress(OSError): + os.close(fd) + + +def _try_lock_fd(fd: int) -> bool: + """Non-blocking exclusive lock on an open fd; True if acquired, False if + another process holds it.""" + if IS_WINDOWS: + import msvcrt + + try: + # A 1-byte range lock at offset 0 (allowed past EOF on Windows). + msvcrt.locking( # pyright: ignore[reportAttributeAccessIssue] + fd, + msvcrt.LK_NBLCK, # pyright: ignore[reportAttributeAccessIssue] + 1, + ) + return True + except OSError: + return False + import fcntl + + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + return True + except OSError: + return False + + +def claimed_session_name(manifest: dict[str, Any]) -> Optional[str]: + """The name this session holds.""" + value = manifest.get("name") + return value if isinstance(value, str) and value else None + + +def iter_agent_sessions() -> Iterator[PresetAgentSession]: + """Yields a handle for every session directory under the presets dir.""" + root = get_presets_dir() + if not root.is_dir(): + return + for path in sorted(root.iterdir()): + if path.is_dir() and not path.name.startswith((".", "models--")): + yield PresetAgentSession(path=path, debug=False, preset_id=path.name) + + +def find_session_name_claims(name: str) -> list[PresetAgentSession]: + """Sessions of any status holding `name`, including failed ones.""" + return [ + session + for session in iter_agent_sessions() + if claimed_session_name(session.read_manifest()) == name + ] + + +def resolve_session_ref(ref: str) -> str: + """A session reference may be a preset id or a claimed name.""" + if (get_presets_dir() / ref).is_dir(): + return ref + claims = find_session_name_claims(ref) + if len(claims) == 1: + return claims[0].preset_id + return ref + + +def list_agent_sessions() -> list[dict[str, Any]]: + entries = [] + for session in iter_agent_sessions(): + path = session.path + manifest = session.read_manifest() + status = manifest.get("status") + if status not in ("running", "interrupted", "success"): + continue + if status == "running" and not session_process_alive(manifest): + status = "interrupted" + entry = dict(manifest) + entry["id"] = path.name + entry["name"] = claimed_session_name(manifest) + entry["status"] = status + entry["trials"] = _summarize_session_trials(path / _TRIALS_FILENAME) + entries.append(entry) + return entries + + +def _summarize_session_trials(path: Path) -> Optional[dict[str, Any]]: + """Best-so-far summary from a session's mirrored trial records.""" + try: + lines = path.read_text(encoding="utf-8").splitlines() + except OSError: + lines = [] + count = 0 + best: Optional[dict[str, Any]] = None + for line in lines: + try: + record = json.loads(line) + except json.JSONDecodeError: + continue + if not isinstance(record, dict): + continue + # One record per trial (the agent contract); trials may share a task, + # so task names must not be deduplicated. + count += 1 + benchmark = record.get("benchmark") + if not isinstance(benchmark, dict): + continue + metrics = benchmark.get("metrics") or {} + workload = benchmark.get("workload") or {} + duration = metrics.get("duration_seconds") + tokens = metrics.get("total_output_tokens") + if not isinstance(duration, (int, float)) or duration <= 0: + continue + if not isinstance(tokens, (int, float)): + continue + tok_s = tokens / duration + if best is None or tok_s > best["tok_s"]: + resources = record.get("resources") or {} + gpu = resources.get("gpu") if isinstance(resources, dict) else None + gpu_text = None + if isinstance(gpu, dict) and gpu.get("name"): + gpu_text = str(gpu["name"]) + if gpu.get("memory"): + gpu_text += f":{gpu['memory']}" + if gpu.get("count"): + gpu_text += f":{gpu['count']}" + best = { + "tok_s": tok_s, + "concurrency": workload.get("concurrency"), + "gpu": gpu_text, + } + return {"count": count, "best": best} + + +def print_preset_progress(message: str, *, agent_session: PresetAgentSession) -> None: + timestamp = datetime.now(timezone.utc).astimezone().strftime("%Y-%m-%d %H:%M:%S") + message = message.rstrip("\r\n") + agent_session.append_log(f"[{timestamp}] {message}") + if not agent_session.echo: + return + console.print( + Text(f"[{timestamp}]", style="log.time"), + Text(message, style="log.message"), + soft_wrap=True, + ) + + +def _pid_running(pid: int) -> bool: + try: + return psutil.Process(pid).status() != psutil.STATUS_ZOMBIE + except psutil.Error: + return False + + +def _process_started_at(pid: int) -> Optional[float]: + try: + return psutil.Process(pid).create_time() + except psutil.Error: + return None + + +def _write_private_text(path: Path, content: str) -> None: + # Atomic tmp + fsync + replace (mkstemp already creates the file 0600), so + # a crash mid-write cannot leave a truncated manifest or offsets file. + fd, temporary = tempfile.mkstemp(dir=path.parent, prefix=f".{path.name}.", suffix=".tmp") + try: + with os.fdopen(fd, "w", encoding="utf-8") as f: + f.write(content) + f.flush() + os.fsync(f.fileno()) + try: + os.replace(temporary, path) + except PermissionError: + if not IS_WINDOWS: + raise + # A concurrent reader (a viewer polling the manifest) can hold the + # destination open without FILE_SHARE_DELETE; retry briefly, then + # prefer an in-place write over crashing the owner. + for _ in range(3): + time.sleep(0.01) + with suppress(PermissionError): + os.replace(temporary, path) + return + path.write_text(content, encoding="utf-8") + finally: + with suppress(FileNotFoundError): + os.unlink(temporary) diff --git a/src/dstack/_internal/cli/services/presets/store.py b/src/dstack/_internal/cli/services/presets/store.py new file mode 100644 index 000000000..6f944d71c --- /dev/null +++ b/src/dstack/_internal/cli/services/presets/store.py @@ -0,0 +1,204 @@ +import os +import shutil +import sys +import tempfile +from contextlib import suppress +from pathlib import Path +from typing import TextIO + +import yaml +from pydantic import ValidationError + +from dstack._internal.cli.models.configurations import ( + MAX_PROMPT_LENGTH, + PresetConfiguration, + PresetPromptFile, +) +from dstack._internal.cli.models.presets import Preset +from dstack._internal.cli.services.presets.presets import preset_to_data +from dstack._internal.cli.utils.common import warn +from dstack._internal.core.errors import CLIError, ConfigurationError +from dstack._internal.utils.common import get_dstack_dir + + +class PresetStore: + """Presets live at `//` — one directory per preset holding + the artifact (`preset.yaml`) next to the creation session internals. + Deleted presets are archived under `/.archive/`.""" + + def __init__(self, root: Path | None = None) -> None: + self.root = root or get_dstack_dir() / "presets" + + def list(self) -> list[Preset]: + if not self.root.exists(): + return [] + self._migrate_legacy() + presets = [] + for path in self.root.glob("*/preset.yaml"): + try: + presets.append(self._load(path)) + except CLIError as e: + # One corrupt file must not take down every read; the preset + # stays deletable by ID. stderr keeps `--json` output parseable. + warn(str(e), stderr=True) + return sorted(presets, key=lambda preset: (preset.base.lower(), preset.id)) + + def get(self, preset_id: str) -> Preset | None: + _validate_preset_id(preset_id) + if not self.root.exists(): + return None + self._migrate_legacy() + path = self.root / preset_id / "preset.yaml" + if not path.is_file(): + return None + preset = self._load(path) + if preset.id != preset_id: + raise CLIError(f"Preset file {path} does not match its path") + return preset + + def save(self, preset: Preset) -> Path: + _validate_preset_id(preset.id) + self._migrate_legacy() + directory = self.root / preset.id + directory.mkdir(parents=True, exist_ok=True) + path = directory / "preset.yaml" + content = yaml.safe_dump(preset_to_data(preset), sort_keys=False) + fd, temporary_path = tempfile.mkstemp( + dir=directory, + prefix=f".{preset.id}.", + suffix=".tmp", + ) + try: + with os.fdopen(fd, "w", encoding="utf-8") as f: + f.write(content) + f.flush() + os.fsync(f.fileno()) + os.replace(temporary_path, path) + finally: + try: + Path(temporary_path).unlink() + except FileNotFoundError: + pass + return path + + def find_by_name(self, name: str) -> Preset | None: + for preset in self.list(): + if preset.name == name: + return preset + return None + + def find_by_id_or_name(self, ref: str) -> Preset | None: + """Resolves a preset reference that may be an ID or a claimed name.""" + return self.get(ref) or self.find_by_name(ref) + + def release_name(self, name: str) -> Preset | None: + """Releases `name` from the preset holding it, keeping the preset.""" + preset = self.find_by_name(name) + if preset is None: + return None + detached = preset.copy(update={"name": None}) + self.save(detached) + return detached + + def delete(self, preset_id: str) -> bool: + # Deletes by path so a corrupt preset file is still removable. + _validate_preset_id(preset_id) + if not self.root.exists(): + return False + self._migrate_legacy() + directory = self.root / preset_id + if not (directory / "preset.yaml").is_file(): + return False + self._archive(directory) + return True + + def _archive(self, directory: Path) -> None: + archive_root = self.root / ".archive" + archive_root.mkdir(mode=0o700, parents=True, exist_ok=True) + target = archive_root / directory.name + index = 0 + while target.exists(): + index += 1 + target = archive_root / f"{directory.name}-{index}" + shutil.move(str(directory), str(target)) + + def _migrate_legacy(self) -> None: + for legacy in list(self.root.glob("models--*/*.yaml")): + target_dir = self.root / legacy.stem + target_dir.mkdir(mode=0o700, parents=True, exist_ok=True) + target = target_dir / "preset.yaml" + if target.exists(): + legacy.unlink() + else: + legacy.replace(target) + for directory in self.root.glob("models--*"): + with suppress(OSError): + directory.rmdir() + + def _load(self, path: Path) -> Preset: + try: + with path.open(encoding="utf-8") as f: + return Preset.parse_obj(yaml.safe_load(f)) + except (OSError, ValidationError, yaml.YAMLError) as e: + raise CLIError(f"Invalid preset file {path}: {e}") from e + + +def _validate_preset_id(preset_id: str) -> None: + if not preset_id or preset_id.startswith(".") or any(char in preset_id for char in "/\\"): + raise CLIError(f"Invalid preset ID: {preset_id!r}") + + +def load_preset_configuration(path: str) -> tuple[str, PresetConfiguration]: + if path == "-": + return "-", _parse_preset_configuration(sys.stdin) + configuration_path = Path(path) + if not configuration_path.is_file(): + raise ConfigurationError(f"Configuration file {path} does not exist") + try: + with configuration_path.open(encoding="utf-8") as f: + configuration = _parse_preset_configuration(f) + except OSError as e: + raise ConfigurationError(f"Failed to load configuration from {path}") from e + return str(configuration_path.resolve()), configuration + + +def _parse_preset_configuration(stream: TextIO) -> PresetConfiguration: + try: + data = yaml.safe_load(stream) + if not isinstance(data, dict): + raise ConfigurationError("Preset configuration must be a YAML object") + configuration = PresetConfiguration.parse_obj(data) + except ValidationError as e: + raise ConfigurationError(e) from e + except yaml.YAMLError as e: + raise ConfigurationError(f"Invalid preset configuration: {e}") from e + model = data.get("model") + if isinstance(model, dict) and model.get("name") is None: + key = "base" if "base" in model else "repo" + warn( + f"The nested `model.{key}` syntax is deprecated" + f" unless `model.name` is set. Use top-level `{key}:` instead" + ) + return configuration + + +def resolve_preset_prompt( + configuration: PresetConfiguration, configuration_path: str +) -> str | None: + """The resolved user prompt text; file paths are relative to the configuration file.""" + if configuration.prompt is None: + return None + if isinstance(configuration.prompt, str): + return configuration.prompt.strip() + assert isinstance(configuration.prompt, PresetPromptFile) + base = Path.cwd() if configuration_path == "-" else Path(configuration_path).parent + path = base / configuration.prompt.path + try: + text = path.read_text(encoding="utf-8").strip() + except (OSError, UnicodeDecodeError) as e: + raise ConfigurationError(f"Failed to read the prompt file {path}: {e}") from e + if not text: + raise ConfigurationError(f"The prompt file {path} is empty") + if len(text) > MAX_PROMPT_LENGTH: + raise ConfigurationError(f"The prompt file {path} exceeds {MAX_PROMPT_LENGTH} characters") + return text diff --git a/src/dstack/_internal/cli/services/presets/tail.py b/src/dstack/_internal/cli/services/presets/tail.py new file mode 100644 index 000000000..2bafbba1d --- /dev/null +++ b/src/dstack/_internal/cli/services/presets/tail.py @@ -0,0 +1,214 @@ +"""Offset-persistent tailers over the agent's stream and record files.""" + +import asyncio +import json +import threading +from contextlib import suppress +from pathlib import Path +from typing import Any, Callable, Optional, Sequence + +from dstack._internal.cli.services.presets.redaction import redact +from dstack._internal.cli.services.presets.session import ( + PresetAgentSession, + _write_private_text, + print_preset_progress, +) +from dstack._internal.cli.utils.common import console + + +class _FileLineReader: + """`readline()` over a growing file, so stream parsing survives CLI + restarts: offsets persist, and a later attach continues exactly where the + previous reader stopped.""" + + _POLL_SECONDS = 0.2 + _MAX_CHUNK = 1024 * 1024 + + def __init__( + self, + path: Path, + *, + offset_store: "_OffsetStore", + offset_key: str, + is_alive: Callable[[], bool], + ) -> None: + self._path = path + self._offset_store = offset_store + self._offset_key = offset_key + self._is_alive = is_alive + self._offset = offset_store.get(offset_key) + self._buffer = b"" + self._drained = False + + def _read_chunk(self) -> bytes: + try: + with self._path.open("rb") as f: + f.seek(self._offset) + return f.read(self._MAX_CHUNK) + except OSError: + return b"" + + async def readline(self) -> bytes: + while True: + if b"\n" in self._buffer: + line, self._buffer = self._buffer.split(b"\n", 1) + return line + b"\n" + # File IO runs off the event loop so a hung filesystem cannot + # freeze the CLI. + chunk = await asyncio.to_thread(self._read_chunk) + if chunk: + self._buffer += chunk + self._offset += len(chunk) + await asyncio.to_thread(self._offset_store.set, self._offset_key, self._offset) + continue + if not self._is_alive(): + if self._drained: + # A final partial line without a newline, then EOF. + line, self._buffer = self._buffer, b"" + return line + # One more read after death to catch the last flush. + self._drained = True + continue + await asyncio.sleep(self._POLL_SECONDS) + + +class _OffsetStore: + """Persists tailer/mirror byte offsets so resumed sessions do not repeat + output. One instance serves the whole session — every reader and mirror + shares it with disjoint keys, and the exclusive session claim guarantees no + other process writes the file.""" + + def __init__(self, path: Path) -> None: + self._path = path + # Consumers flush from worker threads; serialize the update + write. + self._lock = threading.Lock() + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + data = {} + self._data: dict[str, Any] = data if isinstance(data, dict) else {} + + def get(self, key: str) -> int: + value = self._data.get(key, 0) + return value if isinstance(value, int) and value >= 0 else 0 + + def set(self, key: str, value: int) -> None: + with self._lock: + self._data[key] = value + with suppress(OSError): + _write_private_text(self._path, json.dumps(self._data) + "\n") + + +def open_session_offsets(session: PresetAgentSession) -> _OffsetStore: + """The session's single offset store, shared by all its tailers.""" + return _OffsetStore(session.path / ".offsets.json") + + +class _ProgressTailer: + def __init__( + self, + *, + path: Path, + redacted_values: Sequence[str], + agent_session: PresetAgentSession, + offset_store: Optional[_OffsetStore] = None, + offset_key: str = "progress", + ) -> None: + self._path = path + self._redacted_values = redacted_values + self._agent_session = agent_session + self._offset_store = offset_store + self._offset_key = offset_key + self._offset = offset_store.get(offset_key) if offset_store else 0 + + async def run(self) -> None: + while True: + # File IO runs off the event loop; see _FileLineReader.readline. + await asyncio.to_thread(self.flush) + await asyncio.sleep(1) + + def flush(self) -> None: + if not self._path.exists(): + return + with self._path.open("r", encoding="utf-8", errors="replace") as f: + f.seek(self._offset) + lines = f.readlines() + self._offset = f.tell() + if lines and self._offset_store is not None: + self._offset_store.set(self._offset_key, self._offset) + for line in lines: + message = _parse_progress(line) + if message is not None: + print_preset_progress( + redact(message, self._redacted_values), + agent_session=self._agent_session, + ) + + +class _RecordMirror: + """Mirrors a workspace record file into the persistent session directory, redacted.""" + + def __init__( + self, + *, + source: Path, + target: Path, + redacted_values: Sequence[str], + offset_store: Optional[_OffsetStore] = None, + offset_key: str = "", + echo: bool = True, + ) -> None: + self._source = source + self._target = target + self._redacted_values = redacted_values + self._offset_store = offset_store + self._offset_key = offset_key + self._offset = offset_store.get(offset_key) if offset_store and offset_key else 0 + self._enabled = True + self._echo = echo + + async def run(self) -> None: + while True: + # File IO runs off the event loop; see _FileLineReader.readline. + await asyncio.to_thread(self.flush) + await asyncio.sleep(1) + + def flush(self) -> None: + if not self._enabled or not self._source.exists(): + return + with self._source.open("rb") as f: + f.seek(self._offset) + data = f.read() + # Mirror complete lines only; a partial line is kept for the next flush. + end = data.rfind(b"\n") + if end < 0: + return + chunk = data[: end + 1].decode("utf-8", errors="replace") + self._offset += end + 1 + if self._offset_store is not None and self._offset_key: + self._offset_store.set(self._offset_key, self._offset) + try: + if not self._target.exists(): + _write_private_text(self._target, "") + # newline="" writes the binary-read chunk verbatim. Without it, text mode + # re-translates "\n" to "\r\n" on Windows, doubling the "\r\n" a text-written + # source already carries ("\r\r\n"), which reads back as a blank extra line. + with self._target.open("a", encoding="utf-8", newline="") as f: + f.write(redact(chunk, self._redacted_values)) + f.flush() + except OSError as e: + self._enabled = False + if self._echo: + console.print(f"[warning]Could not mirror {self._target.name}: {e}[/]") + + +def _parse_progress(line: str) -> Optional[str]: + try: + value = json.loads(line) + except json.JSONDecodeError: + return line.strip() or None + if isinstance(value, str): + return value.strip() or None + if isinstance(value, dict) and isinstance(value.get("message"), str): + return value["message"].strip() or None + return None diff --git a/src/dstack/_internal/cli/services/endpoints/verify.py b/src/dstack/_internal/cli/services/presets/verify.py similarity index 66% rename from src/dstack/_internal/cli/services/endpoints/verify.py rename to src/dstack/_internal/cli/services/presets/verify.py index d5f2872ec..983b397a4 100644 --- a/src/dstack/_internal/cli/services/endpoints/verify.py +++ b/src/dstack/_internal/cli/services/presets/verify.py @@ -5,33 +5,38 @@ from pydantic import ValidationError -from dstack._internal.cli.models.endpoint_agent import AgentFinalReport -from dstack._internal.cli.models.endpoint_presets import ( - EndpointBenchmarkClient, - EndpointBenchmarkTarget, - EndpointPreset, - EndpointPresetValidationReplica, +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.models.preset_agent import AgentFinalReport +from dstack._internal.cli.models.presets import ( + Preset, + PresetBenchmarkClient, + PresetBenchmarkTarget, + PresetValidationReplica, ) -from dstack._internal.cli.models.endpoints import EndpointConfiguration -from dstack._internal.cli.services.endpoints.agent import ( - EndpointAgentProcessOutput, - EndpointAgentWorkspace, - redact, +from dstack._internal.cli.services.presets.agent import ( + PresetAgentProcessOutput, ) -from dstack._internal.cli.services.endpoints.presets import ( - build_endpoint_preset, +from dstack._internal.cli.services.presets.presets import ( + build_preset, resources_spec_from_instance_resources, ) +from dstack._internal.cli.services.presets.redaction import ( + redact, + redact_structure, +) +from dstack._internal.cli.services.presets.workspace import ( + PresetAgentWorkspace, +) from dstack._internal.core.errors import CLIError from dstack._internal.core.models.configurations import ServiceConfiguration from dstack._internal.core.models.envs import EnvSentinel from dstack._internal.core.models.runs import JobStatus, Run, RunStatus -def load_endpoint_agent_report( +def load_preset_agent_report( *, - output: EndpointAgentProcessOutput, - workspace: EndpointAgentWorkspace, + output: PresetAgentProcessOutput, + workspace: PresetAgentWorkspace, redacted_values: Sequence[str], ) -> AgentFinalReport: report_data = output.report_data or _load_json_object(workspace.final_report_path) @@ -42,6 +47,10 @@ def load_endpoint_agent_report( redacted_values, ) ) + # Scrub known secret values before validation: an echoed secret must never + # be persisted, but it also must not cost the whole session — the bearer + # check below still rejects unknown leaked tokens. + report_data = redact_structure(report_data, redacted_values) try: report = AgentFinalReport.parse_obj(report_data) except ValidationError as e: @@ -56,12 +65,14 @@ def load_endpoint_agent_report( return report -def build_verified_endpoint_preset( +def build_verified_preset( *, run: Run, - endpoint_configuration: EndpointConfiguration, + preset_configuration: PresetConfiguration, report: AgentFinalReport, -) -> EndpointPreset: + preset_id: Optional[str] = None, + name: Optional[str] = None, +) -> Preset: if run.id != report.run_id or run.run_spec.run_name != report.run_name: raise CLIError("Claude final report identifies a different service run") if run.status != RunStatus.RUNNING or run.service is None: @@ -69,20 +80,20 @@ def build_verified_endpoint_preset( service = run.run_spec.configuration if not isinstance(service, ServiceConfiguration) or service.model is None: raise CLIError("Claude final run is not a model service") - if service.model.name != endpoint_configuration.model.api_model_name: - raise CLIError("Claude final service model name does not match the endpoint request") + if service.model.name != preset_configuration.model.api_model_name: + raise CLIError("Claude final service model name does not match the requested model") assert report.base is not None assert report.model is not None assert report.context_length is not None assert report.benchmark is not None - if endpoint_configuration.model.allows_variant_selection: - if report.base != endpoint_configuration.model.api_model_name: - raise CLIError("Claude final report base does not match the endpoint request") - elif report.model != endpoint_configuration.model.exact_repo: + if preset_configuration.model.allows_variant_selection: + if report.base != preset_configuration.model.api_model_name: + raise CLIError("Claude final report base does not match the requested model") + elif report.model != preset_configuration.model.exact_repo: raise CLIError("Claude changed an exact model request") if ( - endpoint_configuration.context_length is not None - and report.context_length < endpoint_configuration.context_length + preset_configuration.context_length is not None + and report.context_length < preset_configuration.context_length ): raise CLIError("Claude final service does not meet the requested context length") @@ -91,30 +102,32 @@ def build_verified_endpoint_preset( ) benchmark = report.benchmark.copy( update={ - "target": EndpointBenchmarkTarget(type=target_type), - "client": EndpointBenchmarkClient(type="local"), + "target": PresetBenchmarkTarget(type=target_type), + "client": PresetBenchmarkClient(type="local"), } ) portable_service = service.copy(deep=True) - # The CLI resolved endpoint env references before submission; presets retain the references. - for key, value in endpoint_configuration.env.items(): + # The CLI resolved preset env references before submission; presets retain the references. + for key, value in preset_configuration.env.items(): if isinstance(value, EnvSentinel) and key in portable_service.env: portable_service.env[key] = value - return build_endpoint_preset( + return build_preset( + name=name, service=portable_service, validation_replicas=_get_validation_replicas(run, service), base_model=report.base, model=report.model, context_length=report.context_length, benchmark=benchmark, + preset_id=preset_id, ) def _get_validation_replicas( run: Run, service: ServiceConfiguration, -) -> list[EndpointPresetValidationReplica]: - replicas: list[EndpointPresetValidationReplica] = [] +) -> list[PresetValidationReplica]: + replicas: list[PresetValidationReplica] = [] for group in service.replica_groups: resources = [] for job in sorted(run.jobs, key=lambda job: job.job_spec.replica_num): @@ -136,7 +149,7 @@ def _get_validation_replicas( ) if not resources: raise CLIError(f"Final service replica group {group.name!r} has no running replicas") - replicas.append(EndpointPresetValidationReplica(resources=resources)) + replicas.append(PresetValidationReplica(resources=resources)) return replicas diff --git a/src/dstack/_internal/cli/services/presets/workspace.py b/src/dstack/_internal/cli/services/presets/workspace.py new file mode 100644 index 000000000..504a9baf0 --- /dev/null +++ b/src/dstack/_internal/cli/services/presets/workspace.py @@ -0,0 +1,285 @@ +"""The agent's working directory: aliasing, wrappers, and bundled skills.""" + +import json +import os +import secrets +import shutil +import subprocess +import sys +import tempfile +from contextlib import suppress +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +from dstack._internal.cli.services.presets.session import ( + _CONSTRAINTS_FILENAME, + _FINAL_REPORT_FILENAME, + _PROGRESS_FILENAME, + _RUNS_FILENAME, + _TRIALS_FILENAME, + PresetAgentSession, +) +from dstack._internal.compat import IS_WINDOWS +from dstack._internal.core.errors import CLIError + +_SKILL_NAMES = ("dstack", "dstack-prototyping") +_PROGRESS_ENV = "DSTACK_PRESET_PROGRESS_LOG" +_MAX_RUN_NAME_LENGTH = 41 +_MAX_UNIX_SOCKET_PATH_BYTES = 103 + + +@dataclass(frozen=True) +class PresetAgentWorkspace: + path: Path + dstack_home: Path + + @property + def temp_path(self) -> Path: + return self.path / ".tmp" + + @property + def bin_path(self) -> Path: + return self.path / "bin" + + @property + def progress_path(self) -> Path: + return self.path / _PROGRESS_FILENAME + + @property + def runs_path(self) -> Path: + return self.path / _RUNS_FILENAME + + @property + def trials_path(self) -> Path: + return self.path / _TRIALS_FILENAME + + @property + def constraints_path(self) -> Path: + return self.path / _CONSTRAINTS_FILENAME + + @property + def final_report_path(self) -> Path: + return self.path / _FINAL_REPORT_FILENAME + + @property + def agent_stdout_path(self) -> Path: + return self.path / ".agent-stdout.jsonl" + + @property + def agent_stderr_path(self) -> Path: + return self.path / ".agent-stderr.log" + + +def create_agent_workspace(session: PresetAgentSession) -> PresetAgentWorkspace: + real = session.path / "workspace" + try: + real.mkdir(mode=0o700) + if IS_WINDOWS: + # Windows has no Unix-socket path limit and symlinks require + # privileges, so the workspace is used directly. + alias = real + else: + alias = _create_workspace_alias(real) + _validate_control_socket_path(alias) + workspace = PresetAgentWorkspace(path=alias / "w", dstack_home=alias / "h") + _prepare_workspace(workspace) + except OSError as e: + raise CLIError(f"Could not create the agent workspace under {real}: {e}") from e + session.update_manifest(workspace=str(real), alias=str(alias)) + return workspace + + +def attach_agent_workspace(session: PresetAgentSession) -> PresetAgentWorkspace: + manifest = session.read_manifest() + real_value, alias_value = manifest.get("workspace"), manifest.get("alias") + if not real_value or not alias_value: + raise CLIError("The preset creation has no workspace to resume") + real, alias = Path(real_value), Path(alias_value) + if not real.is_dir(): + raise CLIError( + f"The preset creation workspace no longer exists: {real}. " + "Stop any leftover runs manually and create a new preset." + ) + if alias != real: + _ensure_workspace_alias(alias, real) + return PresetAgentWorkspace(path=alias / "w", dstack_home=alias / "h") + + +def remove_agent_workspace(session: PresetAgentSession) -> None: + manifest = session.read_manifest() + alias = manifest.get("alias") + workspace = manifest.get("workspace") + if alias and alias != workspace and Path(alias).is_symlink(): + with suppress(OSError): + os.unlink(alias) + if workspace: + shutil.rmtree(workspace, ignore_errors=True) + + +def scrub_workspace_token(session: PresetAgentSession) -> None: + """Removes the agent's dstack config (a live token) from a workspace kept + for resume, so an interrupted session leaves no credential on disk. Resume + re-mints it via `build_preset_agent_env`.""" + workspace = session.read_manifest().get("workspace") + if not workspace: + return + with suppress(OSError): + (Path(workspace) / "h" / ".dstack" / "config.yml").unlink() + + +def _create_workspace_alias(real: Path) -> Path: + base = _get_short_temp_dir() or tempfile.gettempdir() + while True: + alias = Path(base) / f"dpe-{secrets.token_hex(4)}" + try: + os.symlink(real, alias) + except FileExistsError: + continue + return alias + + +def _ensure_workspace_alias(alias: Path, real: Path) -> None: + if os.path.lexists(alias): + if ( + alias.is_symlink() + and os.readlink(alias) == str(real) + and (IS_WINDOWS or os.lstat(alias).st_uid == os.getuid()) + ): + return + raise CLIError( + f"The workspace alias path cannot be used safely: {alias}. " + "Remove it manually if it is yours, or start a new session." + ) + os.symlink(real, alias) + + +def _validate_control_socket_path(build_root: Path) -> None: + if IS_WINDOWS: + return + path = build_root / "h" / ".dstack" / "ssh" / f"{'x' * _MAX_RUN_NAME_LENGTH}.control.sock" + if len(os.fsencode(path)) > _MAX_UNIX_SOCKET_PATH_BYTES: + raise CLIError(f"Temporary path is too long for an SSH control socket: {build_root}") + + +def _prepare_workspace(workspace: PresetAgentWorkspace) -> None: + workspace.path.mkdir(mode=0o700, parents=True, exist_ok=False) + workspace.dstack_home.mkdir(mode=0o700) + workspace.temp_path.mkdir(mode=0o700) + for path in [ + workspace.progress_path, + workspace.runs_path, + workspace.trials_path, + ]: + path.touch() + workspace.bin_path.mkdir() + _install_python_command(workspace.bin_path, "progress", _get_progress_script()) + (workspace.dstack_home / ".ssh").mkdir(mode=0o700) + _install_dstack_wrapper(workspace.bin_path, workspace.dstack_home) + _install_home_wrapper(workspace.bin_path, "ssh", workspace.dstack_home) + _install_skills(workspace.path) + + +def _install_dstack_wrapper(bin_dir: Path, home: Path) -> None: + script = f"""#!{sys.executable} +import os +import sys + +from dstack._internal.cli.main import main + +os.environ["HOME"] = {json.dumps(str(home))} +if os.name == "nt": + os.environ["USERPROFILE"] = {json.dumps(str(home))} +sys.argv[0] = "dstack" +raise SystemExit(main()) +""" + _install_python_command(bin_dir, "dstack", script) + + +def _install_home_wrapper(bin_dir: Path, command: str, home: Path) -> None: + executable = shutil.which(command) + if executable is None: + script = f"""#!{sys.executable} +import sys + +print("Preset creation could not find the {command} executable.", file=sys.stderr) +raise SystemExit(127) +""" + else: + script = f"""#!{sys.executable} +import os +import subprocess +import sys + +os.environ["HOME"] = {json.dumps(str(home))} +if os.name == "nt": + os.environ["USERPROFILE"] = {json.dumps(str(home))} +raise SystemExit(subprocess.call([{json.dumps(executable)}, *sys.argv[1:]])) +""" + _install_python_command(bin_dir, command, script) + + +def _install_python_command(bin_dir: Path, name: str, script: str) -> None: + if not IS_WINDOWS: + path = bin_dir / name + path.write_text(script, encoding="utf-8") + path.chmod(0o755) + return + + # Avoid shadowing packages with the same name when Python adds this directory to sys.path. + script_path = bin_dir / f"_{name}.py" + script_path.write_text(script, encoding="utf-8") + command = subprocess.list2cmdline([sys.executable, str(script_path)]) + (bin_dir / f"{name}.cmd").write_text( + f"@echo off\n{command} %*\n", + encoding="utf-8", + ) + + +def _get_short_temp_dir() -> Optional[str]: + # Keep SSH control sockets below the Unix-domain socket path limit on macOS. + if IS_WINDOWS: + return None + return "/tmp" if Path("/tmp").is_dir() else None + + +def _get_progress_script() -> str: + return f"""#!{sys.executable} +import json +import os +from pathlib import Path +import sys + +message = " ".join(sys.argv[1:]).strip() +if not message and not sys.stdin.isatty(): + message = sys.stdin.read().strip() +if not message: + print("Usage: progress ", file=sys.stderr) + raise SystemExit(2) +path = Path(os.environ.get("{_PROGRESS_ENV}", "{_PROGRESS_FILENAME}")) +with path.open("a", encoding="utf-8") as f: + f.write(json.dumps({{"message": message}}, ensure_ascii=False) + "\\n") +""" + + +def _install_skills(workspace: Path) -> None: + source_dir = _get_skills_dir() + target_dir = workspace / ".claude" / "skills" + target_dir.mkdir(parents=True) + for skill_name in _SKILL_NAMES: + source = source_dir / skill_name + if not (source / "SKILL.md").is_file(): + raise CLIError(f"Missing preset agent skill: {skill_name}") + shutil.copytree(source, target_dir / skill_name) + + +def _get_skills_dir() -> Path: + source_path = Path(__file__).resolve() + candidates = ( + source_path.parent / "resources" / "skills", + source_path.parents[6] / "skills", + ) + for candidate in candidates: + if (candidate / "dstack" / "SKILL.md").is_file(): + return candidate + raise CLIError("Could not find packaged dstack skills") diff --git a/src/dstack/_internal/cli/utils/common.py b/src/dstack/_internal/cli/utils/common.py index b437a3c2c..21d5caf83 100644 --- a/src/dstack/_internal/cli/utils/common.py +++ b/src/dstack/_internal/cli/utils/common.py @@ -29,6 +29,12 @@ force_terminal=settings.CLI_RICH_FORCE_TERMINAL, ) +error_console = Console( + theme=Theme(_colors), + force_terminal=settings.CLI_RICH_FORCE_TERMINAL, + stderr=True, +) + LIVE_TABLE_REFRESH_RATE_PER_SEC = 1 LIVE_TABLE_PROVISION_INTERVAL_SECS = 2 @@ -119,11 +125,12 @@ def add_row_from_dict(table: Table, data: dict[str, Any], **kwargs): table.add_row(*row, **kwargs) -def warn(message: str): +def warn(message: str, stderr: bool = False): + """`stderr=True` keeps the warning off stdout, so `--json` output stays parseable.""" if not message.endswith("\n"): # Additional blank line for better visibility if there are more than one warning message = f"{message}\n" - console.print(f"[warning][bold]{message}[/]") + (error_console if stderr else console).print(f"[warning][bold]{message}[/]") def get_start_time(since: Optional[str]) -> Optional[datetime]: diff --git a/src/dstack/_internal/cli/utils/run.py b/src/dstack/_internal/cli/utils/run.py index d034acb20..fd90b3570 100644 --- a/src/dstack/_internal/cli/utils/run.py +++ b/src/dstack/_internal/cli/utils/run.py @@ -353,7 +353,7 @@ def _format_instance_type( def _format_run_name(run: CoreRun, show_deployment_num: bool) -> str: - parts: List[str] = [run.run_spec.run_name] + parts: List[str] = [run.run_spec.run_name or ""] if show_deployment_num: parts.append(f" [secondary]deployment={run.deployment_num}[/]") return "".join(parts) diff --git a/src/tests/_internal/cli/commands/test_endpoint.py b/src/tests/_internal/cli/commands/test_endpoint.py deleted file mode 100644 index ac843c910..000000000 --- a/src/tests/_internal/cli/commands/test_endpoint.py +++ /dev/null @@ -1,325 +0,0 @@ -import json -from io import StringIO -from types import SimpleNamespace -from unittest.mock import patch - -import pytest -from rich.console import Console -from rich.theme import Theme - -from dstack._internal.cli.services.endpoints import output as endpoint_presets_utils -from dstack._internal.cli.services.endpoints.store import EndpointPresetStore -from tests._internal.cli.common import run_dstack_cli -from tests._internal.cli.endpoint_presets import get_endpoint_preset - -pytestmark = pytest.mark.windows - - -class TestEndpointPresetLocalCommands: - @pytest.fixture(autouse=True) - def mock_ssh_client_info(self): - with patch("dstack._internal.cli.main.get_ssh_client_info"): - yield - - def test_formats_second_scale_ttft_without_scientific_notation(self): - preset = get_endpoint_preset() - ttft = preset.validations[0].benchmark.metrics.ttft_ms - ttft.mean = 8148.3 - ttft.p50 = 8151.4 - ttft.p99 = 8334.2 - - output = endpoint_presets_utils.format_endpoint_benchmark(preset, verbose=True) - - assert "TTFT 8.15s" in output - assert "TTFT mean/p50/p99=8.15/8.15/8.33s" in output - assert "e+03" not in output - - def test_preserves_benchmark_concurrency_at_narrow_width(self, monkeypatch): - output = StringIO() - monkeypatch.setattr( - endpoint_presets_utils, - "console", - Console( - file=output, - width=79, - color_system=None, - theme=Theme({"secondary": "grey58"}), - ), - ) - - endpoint_presets_utils.print_endpoint_presets([get_endpoint_preset()]) - - assert "concurrency=1" in "".join(output.getvalue().split()) - - def test_prints_created_column(self, monkeypatch): - output = StringIO() - monkeypatch.setattr( - endpoint_presets_utils, - "console", - Console( - file=output, - width=160, - color_system=None, - theme=Theme({"secondary": "grey58"}), - ), - ) - monkeypatch.setattr(endpoint_presets_utils, "pretty_date", lambda _: "2 months ago") - - endpoint_presets_utils.print_endpoint_presets([get_endpoint_preset()]) - - assert "CREATED" in output.getvalue() - assert "2 months ago" in output.getvalue() - - def test_handles_keyboard_interrupt(self, tmp_path, capsys): - configuration_path = tmp_path / "endpoint.dstack.yml" - configuration_path.write_text( - "type: endpoint\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\n" - ) - - with ( - patch("dstack.api.Client.from_config"), - patch( - "dstack._internal.cli.commands.endpoint.create_endpoint_preset", - side_effect=KeyboardInterrupt, - ), - ): - exit_code = run_dstack_cli( - ["endpoint", "preset", "create", "-f", str(configuration_path)], - home_dir=tmp_path, - ) - - assert exit_code == 0 - assert "Operation interrupted by user" in capsys.readouterr().out - - def test_lists_and_deletes_preset_without_api_client(self, tmp_path, capsys): - preset = get_endpoint_preset() - EndpointPresetStore(tmp_path / ".dstack" / "presets").save(preset) - list_output = StringIO() - - with ( - patch("dstack.api.Client.from_config") as from_config, - patch.object( - endpoint_presets_utils, - "console", - Console( - file=list_output, - width=160, - color_system=None, - theme=Theme({"secondary": "grey58"}), - ), - ), - patch.object( - endpoint_presets_utils, "pretty_date", return_value="2 months ago" - ) as pretty_date, - ): - assert run_dstack_cli(["endpoint", "preset", "list"], home_dir=tmp_path) == 0 - from_config.assert_not_called() - pretty_date.assert_called_once_with(preset.created_at) - - output = list_output.getvalue() - assert "Qwen/Qwen3.5-27B" in output - assert "preset=8f3a12c4" in output - assert "repo=community/Qwen3.5-27B-GPTQ-Int4" in output - assert "CONTEXT" in output - assert "BENCHMARK" in output - assert "32K" in output - assert "42.1" in output - assert "concurrency=1" in "".join(output.split()) - assert "tok/s" in output - assert "TTFT" in output - assert "108ms" in output - assert "A6000:48GB:1" not in output - - assert run_dstack_cli(["endpoint", "preset", "list", "-v"], home_dir=tmp_path) == 0 - verbose_output = capsys.readouterr().out - assert "hardware=" in verbose_output - assert "api=" in verbose_output - assert "n=16" in verbose_output - assert "concurrency=1" in "".join(verbose_output.split()) - assert "1K->128" in verbose_output - - with patch("dstack.api.Client.from_config") as from_config: - assert ( - run_dstack_cli( - [ - "endpoint", - "preset", - "delete", - preset.id, - "-y", - ], - home_dir=tmp_path, - ) - == 0 - ) - from_config.assert_not_called() - - assert EndpointPresetStore(tmp_path / ".dstack" / "presets").list() == [] - assert not (tmp_path / ".dstack" / "presets" / "models--Qwen--Qwen3.5-27B").exists() - - def test_gets_complete_preset_as_json_without_api_client(self, tmp_path, capsys): - preset = get_endpoint_preset() - EndpointPresetStore(tmp_path / ".dstack" / "presets").save(preset) - - with patch("dstack.api.Client.from_config") as from_config: - assert ( - run_dstack_cli( - ["endpoint", "preset", "get", preset.id, "--json"], - home_dir=tmp_path, - ) - == 0 - ) - from_config.assert_not_called() - - data = json.loads(capsys.readouterr().out) - assert data["id"] == preset.id - assert data["created_at"] == preset.created_at.isoformat() - assert data["context_length"] == 32768 - assert data["validations"][0]["benchmark"]["metrics"]["total_output_tokens"] == 2048 - - @pytest.mark.parametrize( - "args", - [ - ["endpoint", "preset", "--json"], - ["endpoint", "preset", "list", "--json"], - ], - ) - def test_lists_complete_presets_as_json(self, tmp_path, capsys, args): - preset = get_endpoint_preset() - EndpointPresetStore(tmp_path / ".dstack" / "presets").save(preset) - - assert run_dstack_cli(args, home_dir=tmp_path) == 0 - - output = json.loads(capsys.readouterr().out) - assert len(output["presets"]) == 1 - data = output["presets"][0] - assert data["id"] == preset.id - assert data["created_at"] == preset.created_at.isoformat() - assert data["context_length"] == 32768 - assert data["validations"][0]["benchmark"]["metrics"]["total_output_tokens"] == 2048 - - def test_deletes_preset_without_api_client(self, tmp_path): - preset = get_endpoint_preset() - store = EndpointPresetStore(tmp_path / ".dstack" / "presets") - store.save(preset) - store.save(preset.copy(update={"id": "01234567"})) - - with patch("dstack.api.Client.from_config") as from_config: - assert ( - run_dstack_cli( - ["endpoint", "preset", "delete", "--model", preset.base, "-y"], - home_dir=tmp_path, - ) - == 0 - ) - from_config.assert_not_called() - - assert store.list() == [] - - def test_merges_profile_configuration_and_cli_args(self, tmp_path): - (tmp_path / ".dstack").mkdir() - (tmp_path / ".dstack" / "profiles.yml").write_text( - """profiles: - - name: gpu - default: true - backends: [aws] - regions: [profile-region] - max_price: 0.25 - spot_policy: spot -""" - ) - configuration_path = tmp_path / "endpoint.dstack.yml" - configuration_path.write_text( - """type: endpoint -name: file-name -model: - base: Qwen/Qwen3.5-27B -regions: [file-region] -max_price: 0.5 -env: - - HF_TOKEN -""" - ) - preset = get_endpoint_preset() - result = SimpleNamespace( - preset=preset, - path=tmp_path / "preset.yaml", - final_run_name="qwen-build-2", - ) - - with ( - patch("dstack.api.Client.from_config"), - patch( - "dstack._internal.cli.commands.endpoint.create_endpoint_preset", - return_value=result, - ) as create, - ): - exit_code = run_dstack_cli( - [ - "endpoint", - "preset", - "create", - "-f", - str(configuration_path), - "--name", - "cli-name", - "--backend", - "gcp", - "--max-price", - "0.75", - "--fleet", - "cli-fleet", - "--debug", - ], - home_dir=tmp_path, - repo_dir=tmp_path, - ) - - assert exit_code == 0 - configuration = create.call_args.kwargs["configuration"] - assert configuration.name == "cli-name" - assert configuration.backends == ["gcp"] - assert configuration.regions == ["file-region"] - assert configuration.max_price == 0.75 - assert configuration.spot_policy.value == "spot" - assert [fleet.format() for fleet in configuration.fleets] == ["cli-fleet"] - assert create.call_args.kwargs["debug"] is True - - @pytest.mark.parametrize( - ("extra_args", "expected_preset"), - [([], "file-preset"), (["--preset", "cli-preset"], "cli-preset")], - ) - def test_apply_passes_selected_profile_and_preset(self, tmp_path, extra_args, expected_preset): - (tmp_path / ".dstack").mkdir() - (tmp_path / ".dstack" / "profiles.yml").write_text( - "profiles:\n - name: gpu\n max_price: 0.5\n" - ) - configuration_path = tmp_path / "endpoint.dstack.yml" - configuration_path.write_text( - "type: endpoint\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\npreset: file-preset\n" - ) - - with ( - patch("dstack.api.Client.from_config"), - patch("dstack._internal.cli.commands.endpoint.apply_endpoint_preset") as apply, - ): - args = [ - "endpoint", - "preset", - "apply", - "-f", - str(configuration_path), - "--profile", - "gpu", - *extra_args, - ] - exit_code = run_dstack_cli( - args, - home_dir=tmp_path, - repo_dir=tmp_path, - ) - - assert exit_code == 0 - assert apply.call_args.kwargs["profile_name"] == "gpu" - assert apply.call_args.kwargs["preset_id"] == expected_preset - assert apply.call_args.kwargs["configuration"].max_price == 0.5 diff --git a/src/tests/_internal/cli/commands/test_preset.py b/src/tests/_internal/cli/commands/test_preset.py new file mode 100644 index 000000000..5f8c17c65 --- /dev/null +++ b/src/tests/_internal/cli/commands/test_preset.py @@ -0,0 +1,444 @@ +import json +from contextlib import contextmanager +from io import StringIO +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from dstack._internal.cli.services.presets import output as presets_utils +from dstack._internal.cli.services.presets.store import PresetStore +from tests._internal.cli.common import plain_console, run_dstack_cli +from tests._internal.cli.preset_factories import get_preset + +pytestmark = pytest.mark.windows + + +@contextmanager +def _patched_create_preset(**create_kwargs): + """Patch the create pipeline; yields the create_preset mock.""" + with ( + patch("dstack.api.Client.from_config"), + patch( + "dstack._internal.cli.commands.preset.plan_preset", + return_value=("fleet-a",), + ), + patch("dstack._internal.cli.commands.preset.create_preset", **create_kwargs) as create, + ): + yield create + + +@pytest.fixture(autouse=True) +def mock_ssh_client_info(): + # Keeps every CLI invocation in this module off the real ssh binary. + with patch("dstack._internal.cli.main.get_ssh_client_info"): + yield + + +class TestPresetLocalCommands: + def test_handles_keyboard_interrupt(self, tmp_path, capsys): + configuration_path = tmp_path / "preset.dstack.yml" + configuration_path.write_text( + "type: preset\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\nmax_trials: 1\n" + ) + + with _patched_create_preset(side_effect=KeyboardInterrupt): + exit_code = run_dstack_cli( + ["preset", "create", "-y", "-f", str(configuration_path)], + home_dir=tmp_path, + ) + + assert exit_code == 0 + # create reports its own Ctrl+C outcome (Detached / interrupted) via the + # interrupt handler, so the command layer stays silent — no generic line. + assert "Operation interrupted by user" not in capsys.readouterr().out + + def test_create_ends_quietly_when_stopped_from_another_cli(self, tmp_path, capsys): + from dstack._internal.cli.services.presets.create import CreationStopped + + configuration_path = tmp_path / "preset.dstack.yml" + configuration_path.write_text( + "type: preset\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\nmax_trials: 1\n" + ) + + with _patched_create_preset(side_effect=CreationStopped): + exit_code = run_dstack_cli( + ["preset", "create", "-y", "-f", str(configuration_path)], + home_dir=tmp_path, + ) + + # The stopping CLI reported the interruption; the owner just ends. + assert exit_code == 0 + assert "Traceback" not in capsys.readouterr().err + + def test_create_requires_max_trials(self, tmp_path, capsys): + configuration_path = tmp_path / "preset.dstack.yml" + configuration_path.write_text("type: preset\nname: qwen\nbase: Qwen/Qwen3.5-27B\n") + + with _patched_create_preset() as create: + exit_code = run_dstack_cli( + ["preset", "create", "-y", "-f", str(configuration_path)], + home_dir=tmp_path, + repo_dir=tmp_path, + ) + + assert exit_code == 1 + captured = capsys.readouterr() + assert "max_trials is required" in captured.out + captured.err + create.assert_not_called() + + def _list_output(self, tmp_path, args, *, created_at): + output = StringIO() + with ( + patch("dstack.api.Client.from_config") as from_config, + patch.object(presets_utils, "console", plain_console(output)), + patch.object(presets_utils, "pretty_date", return_value="2 months ago") as pretty_date, + ): + assert run_dstack_cli(args, home_dir=tmp_path) == 0 + from_config.assert_not_called() + pretty_date.assert_called_once_with(created_at) + return output.getvalue() + + def test_lists_presets_without_api_client(self, tmp_path): + preset = get_preset() + PresetStore(tmp_path / ".dstack" / "presets").save(preset) + + output = self._list_output(tmp_path, ["preset", "list"], created_at=preset.created_at) + + assert "Qwen/Qwen3.5-27B" in output + assert "8f3a12c4" in output + # The repo row is shown only in verbose mode. + assert "repo=community/Qwen3.5-27B-GPTQ-Int4" not in output + # The context column is shown only in verbose mode. + assert "CONTEXT" not in output + assert "BENCHMARK" in output + assert "32K" not in output + assert "42.1" in output + assert "con=1" in "".join(output.split()) + assert "tok/s" in output + assert "TTFT" in output + assert "108ms" in output + assert "A6000:48GB:1" not in output + + def test_verbose_list_adds_repo_and_context(self, tmp_path): + preset = get_preset() + PresetStore(tmp_path / ".dstack" / "presets").save(preset) + + joined_verbose = "".join( + self._list_output( + tmp_path, ["preset", "list", "-v"], created_at=preset.created_at + ).split() + ) + + # Verbose adds only the repo and the ctx= benchmark prefix. + assert "repo=community/Qwen3.5-27B-GPTQ-Int4" in joined_verbose + assert "ctx=32K" in joined_verbose + assert "con=1" in joined_verbose + assert "hardware=" not in joined_verbose + + def test_deletes_preset_without_api_client(self, tmp_path): + preset = get_preset() + PresetStore(tmp_path / ".dstack" / "presets").save(preset) + + with patch("dstack.api.Client.from_config") as from_config: + assert run_dstack_cli(["preset", "delete", preset.id, "-y"], home_dir=tmp_path) == 0 + from_config.assert_not_called() + + assert PresetStore(tmp_path / ".dstack" / "presets").list() == [] + + def test_gets_complete_preset_as_json_without_api_client(self, tmp_path, capsys): + preset = get_preset() + PresetStore(tmp_path / ".dstack" / "presets").save(preset) + + with patch("dstack.api.Client.from_config") as from_config: + assert ( + run_dstack_cli( + ["preset", "get", preset.id, "--json"], + home_dir=tmp_path, + ) + == 0 + ) + from_config.assert_not_called() + + data = json.loads(capsys.readouterr().out) + assert data["id"] == preset.id + assert data["created_at"] == preset.created_at.isoformat() + assert data["context_length"] == 32768 + assert data["validations"][0]["benchmark"]["metrics"]["total_output_tokens"] == 2048 + + @pytest.mark.parametrize( + "args", + [ + ["preset", "--json"], + ["preset", "list", "--json"], + ], + ) + def test_lists_complete_presets_as_json(self, tmp_path, capsys, args): + preset = get_preset() + PresetStore(tmp_path / ".dstack" / "presets").save(preset) + + assert run_dstack_cli(args, home_dir=tmp_path) == 0 + + output = json.loads(capsys.readouterr().out) + assert len(output["presets"]) == 1 + data = output["presets"][0] + assert data["id"] == preset.id + assert data["created_at"] == preset.created_at.isoformat() + assert data["context_length"] == 32768 + assert data["validations"][0]["benchmark"]["metrics"]["total_output_tokens"] == 2048 + + @pytest.mark.parametrize("flag_attribute", [("--base", "base"), ("--repo", "model")]) + def test_deletes_all_presets_of_model_keeping_others_without_api_client( + self, tmp_path, flag_attribute + ): + flag, attribute = flag_attribute + preset = get_preset() + store = PresetStore(tmp_path / ".dstack" / "presets") + store.save(preset) + store.save(preset.copy(update={"id": "01234567"})) + # A preset of a different model must survive the delete. + store.save( + preset.copy(update={"id": "89abcdef", "base": "meta/Llama-4", "model": "meta/Llama-4"}) + ) + + with patch("dstack.api.Client.from_config") as from_config: + assert ( + run_dstack_cli( + ["preset", "delete", flag, getattr(preset, attribute), "-y"], + home_dir=tmp_path, + ) + == 0 + ) + from_config.assert_not_called() + + assert [remaining.id for remaining in store.list()] == ["89abcdef"] + + @pytest.mark.parametrize("flag_attribute", [("--base", "base"), ("--repo", "model")]) + def test_lists_presets_filtered_by_model(self, tmp_path, capsys, flag_attribute): + flag, attribute = flag_attribute + preset = get_preset() + store = PresetStore(tmp_path / ".dstack" / "presets") + store.save(preset) + store.save( + preset.copy(update={"id": "01234567", "base": "meta/Llama-4", "model": "meta/Llama-4"}) + ) + + args = ["preset", "list", "--json", flag, getattr(preset, attribute)] + assert run_dstack_cli(args, home_dir=tmp_path) == 0 + + output = json.loads(capsys.readouterr().out) + assert [entry["id"] for entry in output["presets"]] == [preset.id] + + def test_corrupt_preset_stays_deletable_and_keeps_json_parseable(self, tmp_path, capsys): + preset = get_preset() + store = PresetStore(tmp_path / ".dstack" / "presets") + store.save(preset) + corrupt_dir = tmp_path / ".dstack" / "presets" / "deadbeef" + corrupt_dir.mkdir(parents=True) + (corrupt_dir / "preset.yaml").write_text("{not valid yaml") + + # The corrupt-file warning goes to stderr, so --json stdout stays parseable. + assert run_dstack_cli(["preset", "list", "--json"], home_dir=tmp_path) == 0 + captured = capsys.readouterr() + output = json.loads(captured.out) + assert [entry["id"] for entry in output["presets"]] == [preset.id] + assert "Invalid preset file" in captured.err + + # And the corrupt preset is still removable through the CLI. + assert run_dstack_cli(["preset", "delete", "deadbeef", "-y"], home_dir=tmp_path) == 0 + assert not corrupt_dir.exists() + assert [remaining.id for remaining in store.list()] == [preset.id] + + def test_merges_profile_configuration_and_cli_args(self, tmp_path): + (tmp_path / ".dstack").mkdir() + (tmp_path / ".dstack" / "profiles.yml").write_text( + """profiles: + - name: gpu + default: true + backends: [aws] + regions: [profile-region] + max_price: 0.25 + spot_policy: spot +""" + ) + configuration_path = tmp_path / "preset.dstack.yml" + configuration_path.write_text( + """type: preset +name: file-name +model: + base: Qwen/Qwen3.5-27B +regions: [file-region] +max_price: 0.5 +max_trials: 1 +env: + - HF_TOKEN +""" + ) + preset = get_preset() + result = SimpleNamespace( + preset=preset, + path=tmp_path / "preset.yaml", + final_run_name="qwen-build-2", + ) + + with _patched_create_preset(return_value=result) as create: + exit_code = run_dstack_cli( + [ + "preset", + "create", + "-y", + "-f", + str(configuration_path), + "--name", + "cli-name", + "--backend", + "gcp", + "--max-price", + "0.75", + "--fleet", + "cli-fleet", + "--debug", + ], + home_dir=tmp_path, + repo_dir=tmp_path, + ) + + assert exit_code == 0 + configuration = create.call_args.kwargs["configuration"] + assert configuration.name == "cli-name" + assert configuration.backends == ["gcp"] + assert configuration.regions == ["file-region"] + assert configuration.max_price == 0.75 + assert configuration.spot_policy.value == "spot" + assert [fleet.format() for fleet in configuration.fleets] == ["cli-fleet"] + assert create.call_args.kwargs["debug"] is True + + def test_apply_passes_selected_profile_and_preset_id(self, tmp_path): + extra_args = ["--id", "cli-preset"] + (tmp_path / ".dstack").mkdir() + (tmp_path / ".dstack" / "profiles.yml").write_text( + "profiles:\n - name: gpu\n max_price: 0.5\n" + ) + configuration_path = tmp_path / "preset.dstack.yml" + configuration_path.write_text( + "type: preset\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\n" + ) + + with ( + patch("dstack.api.Client.from_config"), + patch("dstack._internal.cli.commands.preset.apply_preset") as apply, + ): + args = [ + "preset", + "apply", + "-f", + str(configuration_path), + "--profile", + "gpu", + *extra_args, + ] + exit_code = run_dstack_cli( + args, + home_dir=tmp_path, + repo_dir=tmp_path, + ) + + assert exit_code == 0 + assert apply.call_args.kwargs["profile_name"] == "gpu" + assert apply.call_args.kwargs["preset_id"] == "cli-preset" + assert apply.call_args.kwargs["configuration"].max_price == 0.5 + + def test_apply_requires_preset_id(self, tmp_path, capsys): + configuration_path = tmp_path / "preset.dstack.yml" + configuration_path.write_text( + "type: preset\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\n" + ) + + with patch("dstack._internal.cli.commands.preset.apply_preset") as apply: + exit_code = run_dstack_cli( + ["preset", "apply", "-f", str(configuration_path)], + home_dir=tmp_path, + repo_dir=tmp_path, + ) + + assert exit_code == 2 + assert "--id" in capsys.readouterr().err + apply.assert_not_called() + + +class TestPresetNameClaims: + def test_create_detaches_the_name_from_the_old_preset(self, tmp_path): + preset = get_preset().copy(update={"name": "qwen"}) + store = PresetStore(tmp_path / ".dstack" / "presets") + store.save(preset) + configuration_path = tmp_path / "preset.dstack.yml" + configuration_path.write_text( + "type: preset\nname: qwen\nbase: Qwen/Qwen3.5-27B\nmax_trials: 1\n" + ) + result = SimpleNamespace( + preset=preset, path=tmp_path / "preset.yaml", final_run_name="qwen-1" + ) + + with _patched_create_preset(return_value=result) as create: + exit_code = run_dstack_cli( + ["preset", "create", "-f", str(configuration_path), "-y"], + home_dir=tmp_path, + repo_dir=tmp_path, + ) + + assert exit_code == 0 + create.assert_called_once() + assert store.get(preset.id).name is None + + def test_create_without_confirmation_exits_before_creating(self, tmp_path): + preset = get_preset().copy(update={"name": "qwen"}) + store = PresetStore(tmp_path / ".dstack" / "presets") + store.save(preset) + configuration_path = tmp_path / "preset.dstack.yml" + configuration_path.write_text( + "type: preset\nname: qwen\nbase: Qwen/Qwen3.5-27B\nmax_trials: 1\n" + ) + + with ( + _patched_create_preset() as create, + patch("dstack._internal.cli.commands.preset.confirm_ask", return_value=False), + ): + exit_code = run_dstack_cli( + ["preset", "create", "-f", str(configuration_path)], + home_dir=tmp_path, + repo_dir=tmp_path, + ) + + assert exit_code == 0 + create.assert_not_called() + assert store.get(preset.id).name == "qwen" + + def test_get_and_delete_resolve_names(self, tmp_path, capsys): + preset = get_preset().copy(update={"name": "qwen"}) + PresetStore(tmp_path / ".dstack" / "presets").save(preset) + + assert run_dstack_cli(["preset", "get", "qwen", "--json"], home_dir=tmp_path) == 0 + assert json.loads(capsys.readouterr().out)["id"] == preset.id + + assert run_dstack_cli(["preset", "delete", "qwen", "-y"], home_dir=tmp_path) == 0 + assert PresetStore(tmp_path / ".dstack" / "presets").list() == [] + + def test_create_always_asks_even_without_a_name_conflict(self, tmp_path): + configuration_path = tmp_path / "preset.dstack.yml" + configuration_path.write_text("type: preset\nbase: Qwen/Qwen3.5-27B\nmax_trials: 1\n") + + with ( + _patched_create_preset() as create, + patch( + "dstack._internal.cli.commands.preset.confirm_ask", return_value=False + ) as confirm, + ): + exit_code = run_dstack_cli( + ["preset", "create", "-f", str(configuration_path)], + home_dir=tmp_path, + repo_dir=tmp_path, + ) + + assert exit_code == 0 + confirm.assert_called_once_with("Create the preset?") + create.assert_not_called() diff --git a/src/tests/_internal/cli/common.py b/src/tests/_internal/cli/common.py index b5f1dd915..4e4442ba3 100644 --- a/src/tests/_internal/cli/common.py +++ b/src/tests/_internal/cli/common.py @@ -1,12 +1,25 @@ import os from pathlib import Path -from typing import List, Optional +from typing import IO, List, Optional from unittest.mock import patch +from rich.console import Console +from rich.theme import Theme + from dstack._internal.cli.main import main from dstack._internal.compat import IS_WINDOWS +def plain_console(file: IO[str], *, width: int = 250) -> Console: + """A plain-text console for asserting on CLI table output.""" + return Console( + file=file, + width=width, + color_system=None, + theme=Theme({"secondary": "grey58"}), + ) + + def run_dstack_cli( cli_args: List[str], home_dir: Optional[Path] = None, diff --git a/src/tests/_internal/cli/models/test_configurations.py b/src/tests/_internal/cli/models/test_configurations.py new file mode 100644 index 000000000..81b17146d --- /dev/null +++ b/src/tests/_internal/cli/models/test_configurations.py @@ -0,0 +1,84 @@ +import pytest +from pydantic import ValidationError + +from dstack._internal.cli.models.configurations import ( + PresetConfiguration, + PresetModelBase, + PresetModelRepo, +) + +pytestmark = pytest.mark.windows + + +class TestPresetConfiguration: + def test_schema_documents_supported_input(self): + assert all( + field.field_info.description for field in PresetConfiguration.__fields__.values() + ) + assert all(field.field_info.description for field in PresetModelBase.__fields__.values()) + assert all(field.field_info.description for field in PresetModelRepo.__fields__.values()) + assert {"type": "string"} in PresetConfiguration.schema()["properties"]["model"]["anyOf"] + + def test_parses_string_as_exact_repo(self): + configuration = PresetConfiguration(model="Qwen/Qwen3.5-27B") + + assert isinstance(configuration.model, PresetModelRepo) + assert configuration.model.exact_repo == "Qwen/Qwen3.5-27B" + assert configuration.model.api_model_name == "Qwen/Qwen3.5-27B" + assert not configuration.model.allows_variant_selection + + def test_parses_base_model(self): + configuration = PresetConfiguration(model={"base": "Qwen/Qwen3.5-27B"}) + + assert isinstance(configuration.model, PresetModelBase) + assert configuration.model.exact_repo is None + assert configuration.model.api_model_name == "Qwen/Qwen3.5-27B" + assert configuration.model.allows_variant_selection + + def test_parses_exact_repo_with_client_facing_name(self): + configuration = PresetConfiguration( + model={ + "repo": "community/Qwen3.5-27B-GPTQ-Int4", + "name": "Qwen/Qwen3.5-27B", + } + ) + + assert configuration.model.exact_repo == "community/Qwen3.5-27B-GPTQ-Int4" + assert configuration.model.api_model_name == "Qwen/Qwen3.5-27B" + + def test_rejects_ambiguous_model_object(self): + with pytest.raises(ValidationError): + PresetConfiguration(model={"base": "Qwen/base", "repo": "Qwen/repo"}) + + def test_parses_top_level_base_shorthand(self): + configuration = PresetConfiguration(base="Qwen/Qwen3.5-27B") + + assert isinstance(configuration.model, PresetModelBase) + assert configuration.model.api_model_name == "Qwen/Qwen3.5-27B" + assert configuration.base is None + + def test_parses_top_level_repo_shorthand(self): + configuration = PresetConfiguration(repo="community/Qwen3.5-27B-GPTQ-Int4") + + assert isinstance(configuration.model, PresetModelRepo) + assert configuration.model.exact_repo == "community/Qwen3.5-27B-GPTQ-Int4" + assert configuration.repo is None + + def test_shorthand_round_trips_through_dict(self): + configuration = PresetConfiguration(base="Qwen/Qwen3.5-27B") + + round_tripped = PresetConfiguration.parse_obj(configuration.dict()) + + assert round_tripped.model == configuration.model + + def test_rejects_combined_base_and_repo_shorthand(self): + with pytest.raises(ValidationError): + PresetConfiguration(base="Qwen/base", repo="Qwen/repo") + + def test_rejects_shorthand_combined_with_model(self): + with pytest.raises(ValidationError): + PresetConfiguration(base="Qwen/base", model={"repo": "Qwen/repo"}) + + def test_requires_model(self): + with pytest.raises(ValidationError): + PresetConfiguration() diff --git a/src/tests/_internal/cli/models/test_endpoints.py b/src/tests/_internal/cli/models/test_endpoints.py deleted file mode 100644 index a507a1a01..000000000 --- a/src/tests/_internal/cli/models/test_endpoints.py +++ /dev/null @@ -1,51 +0,0 @@ -import pytest -from pydantic import ValidationError - -from dstack._internal.cli.models.endpoints import ( - EndpointConfiguration, - EndpointModelBase, - EndpointModelRepo, -) - -pytestmark = pytest.mark.windows - - -class TestEndpointConfiguration: - def test_schema_documents_supported_input(self): - assert all( - field.field_info.description for field in EndpointConfiguration.__fields__.values() - ) - assert all(field.field_info.description for field in EndpointModelBase.__fields__.values()) - assert all(field.field_info.description for field in EndpointModelRepo.__fields__.values()) - assert {"type": "string"} in EndpointConfiguration.schema()["properties"]["model"]["anyOf"] - - def test_parses_string_as_exact_repo(self): - configuration = EndpointConfiguration(model="Qwen/Qwen3.5-27B") - - assert isinstance(configuration.model, EndpointModelRepo) - assert configuration.model.exact_repo == "Qwen/Qwen3.5-27B" - assert configuration.model.api_model_name == "Qwen/Qwen3.5-27B" - assert not configuration.model.allows_variant_selection - - def test_parses_base_model(self): - configuration = EndpointConfiguration(model={"base": "Qwen/Qwen3.5-27B"}) - - assert isinstance(configuration.model, EndpointModelBase) - assert configuration.model.exact_repo is None - assert configuration.model.api_model_name == "Qwen/Qwen3.5-27B" - assert configuration.model.allows_variant_selection - - def test_parses_exact_repo_with_client_facing_name(self): - configuration = EndpointConfiguration( - model={ - "repo": "community/Qwen3.5-27B-GPTQ-Int4", - "name": "Qwen/Qwen3.5-27B", - } - ) - - assert configuration.model.exact_repo == "community/Qwen3.5-27B-GPTQ-Int4" - assert configuration.model.api_model_name == "Qwen/Qwen3.5-27B" - - def test_rejects_ambiguous_model_object(self): - with pytest.raises(ValidationError): - EndpointConfiguration(model={"base": "Qwen/base", "repo": "Qwen/repo"}) diff --git a/src/tests/_internal/cli/models/test_presets.py b/src/tests/_internal/cli/models/test_presets.py new file mode 100644 index 000000000..d71e7eba8 --- /dev/null +++ b/src/tests/_internal/cli/models/test_presets.py @@ -0,0 +1,52 @@ +import pytest +from pydantic import ValidationError + +from dstack._internal.cli.models.preset_agent import AGENT_FINAL_REPORT_JSON_SCHEMA +from dstack._internal.cli.models.presets import ( + PresetBenchmark, + PresetBenchmarkLatency, + PresetBenchmarkMetrics, + PresetBenchmarkWorkload, +) +from tests._internal.cli.preset_factories import get_preset_benchmark + +pytestmark = pytest.mark.windows + + +class TestPresetBenchmark: + def test_agent_schema_matches_benchmark_model(self): + schema = AGENT_FINAL_REPORT_JSON_SCHEMA["properties"]["benchmark"] + assert set(schema["properties"]) == set(PresetBenchmark.__fields__) - { + "target", + "client", + } + assert set(schema["required"]) == set(schema["properties"]) + workload_schema = schema["properties"]["workload"] + metrics_schema = schema["properties"]["metrics"] + assert set(workload_schema["properties"]) == set(PresetBenchmarkWorkload.__fields__) + assert set(workload_schema["required"]) == set(workload_schema["properties"]) + assert set(metrics_schema["properties"]) == set(PresetBenchmarkMetrics.__fields__) + assert set(metrics_schema["required"]) == set(metrics_schema["properties"]) + assert set(metrics_schema["properties"]["ttft_ms"]["properties"]) == set( + PresetBenchmarkLatency.__fields__ + ) + + @pytest.mark.parametrize( + ("field", "value", "error"), + [ + ("failed_requests", 1, "must not include failed requests"), + ("successful_requests", 15, "must match workload.num_requests"), + ], + ) + def test_rejects_inconsistent_successful_metrics(self, field, value, error): + data = get_preset_benchmark().dict() + data["metrics"][field] = value + with pytest.raises(ValidationError, match=error): + PresetBenchmark.parse_obj(data) + + def test_rejects_tool_specific_metrics(self): + data = get_preset_benchmark().dict() + data["metrics"]["tool_specific"] = 1 + + with pytest.raises(ValidationError, match="extra fields not permitted"): + PresetBenchmark.parse_obj(data) diff --git a/src/tests/_internal/cli/endpoint_presets.py b/src/tests/_internal/cli/preset_factories.py similarity index 81% rename from src/tests/_internal/cli/endpoint_presets.py rename to src/tests/_internal/cli/preset_factories.py index fe654857e..3104447d7 100644 --- a/src/tests/_internal/cli/endpoint_presets.py +++ b/src/tests/_internal/cli/preset_factories.py @@ -2,14 +2,14 @@ from types import SimpleNamespace from uuid import uuid4 -from dstack._internal.cli.models.endpoint_agent import AgentFinalReport -from dstack._internal.cli.models.endpoint_presets import ( - EndpointBenchmark, - EndpointBenchmarkClient, - EndpointBenchmarkTarget, - EndpointPreset, - EndpointPresetValidation, - EndpointPresetValidationReplica, +from dstack._internal.cli.models.preset_agent import AgentFinalReport +from dstack._internal.cli.models.presets import ( + Preset, + PresetBenchmark, + PresetBenchmarkClient, + PresetBenchmarkTarget, + PresetValidation, + PresetValidationReplica, ) from dstack._internal.core.models.configurations import ServiceConfiguration from dstack._internal.core.models.instances import Disk, Gpu, Resources @@ -17,8 +17,8 @@ from dstack._internal.core.models.runs import JobStatus, Run, RunStatus, ServiceSpec -def get_endpoint_benchmark(*, verified: bool = True) -> EndpointBenchmark: - benchmark = EndpointBenchmark( +def get_preset_benchmark(*, verified: bool = True) -> PresetBenchmark: + benchmark = PresetBenchmark( tool="vllm bench serve", tool_version="0.11.0", command="vllm bench serve --base-url $SERVICE_URL", @@ -43,17 +43,17 @@ def get_endpoint_benchmark(*, verified: bool = True) -> EndpointBenchmark: return benchmark return benchmark.copy( update={ - "target": EndpointBenchmarkTarget(type="server-proxy"), - "client": EndpointBenchmarkClient(type="local"), + "target": PresetBenchmarkTarget(type="server-proxy"), + "client": PresetBenchmarkClient(type="local"), } ) -def get_endpoint_preset( +def get_preset( *, preset_id: str = "8f3a12c4", context_length: int = 32768, -) -> EndpointPreset: +) -> Preset: resources = ResourcesSpec.parse_obj( { "cpu": "16", @@ -62,7 +62,7 @@ def get_endpoint_preset( "gpu": {"name": "A6000", "memory": "48GB", "count": 1}, } ) - return EndpointPreset( + return Preset( base="Qwen/Qwen3.5-27B", id=preset_id, model="community/Qwen3.5-27B-GPTQ-Int4", @@ -79,9 +79,9 @@ def get_endpoint_preset( } ), validations=[ - EndpointPresetValidation( - replicas=[EndpointPresetValidationReplica(resources=[resources])], - benchmark=get_endpoint_benchmark(), + PresetValidation( + replicas=[PresetValidationReplica(resources=[resources])], + benchmark=get_preset_benchmark(), ) ], ) @@ -136,7 +136,7 @@ def get_running_service_run() -> Run: ) -def get_successful_endpoint_report(run: Run) -> AgentFinalReport: +def get_successful_preset_report(run: Run) -> AgentFinalReport: return AgentFinalReport( success=True, run_id=run.id, @@ -145,5 +145,5 @@ def get_successful_endpoint_report(run: Run) -> AgentFinalReport: base="Qwen/Qwen3.5-27B", model="community/Qwen3.5-27B-GPTQ-Int4", context_length=32768, - benchmark=get_endpoint_benchmark(verified=False), + benchmark=get_preset_benchmark(verified=False), ) diff --git a/src/tests/_internal/cli/services/endpoints/test_agent.py b/src/tests/_internal/cli/services/endpoints/test_agent.py deleted file mode 100644 index c3c13ef52..000000000 --- a/src/tests/_internal/cli/services/endpoints/test_agent.py +++ /dev/null @@ -1,393 +0,0 @@ -import asyncio -import json -import os -import shutil -import subprocess -import sys -from types import SimpleNamespace - -import psutil -import pytest -import yaml - -from dstack._internal.cli.models.endpoints import EndpointConfiguration -from dstack._internal.cli.services.endpoints.agent import ( - ClaudeAuth, - EndpointAgentSession, - EndpointAgentWorkspace, - _build_claude_command, - _create_agent_session_directory, - _prepare_subprocess_command, - _ProgressTailer, - _terminate_process, - build_endpoint_agent_env, - contains_redacted_value, - create_endpoint_agent_session, - endpoint_agent_workspace, - get_claude_auth, - print_endpoint_progress, - redact, - run_endpoint_agent, -) -from dstack._internal.compat import IS_WINDOWS -from dstack._internal.core.errors import CLIError -from dstack._internal.core.services.configs import ConfigManager - -pytestmark = pytest.mark.windows - - -def _claude_auth(*, api_key: str | None = "anthropic-secret", effort=None) -> ClaudeAuth: - return ClaudeAuth( - api_key=api_key, - executable="claude", - effort=effort, - model="claude-test", - ) - - -class TestClaudeAuth: - def test_uses_api_key_when_set(self, monkeypatch): - monkeypatch.setenv("DSTACK_AGENT_ANTHROPIC_API_KEY", "key") - monkeypatch.setattr(shutil, "which", lambda _: "/usr/bin/claude") - - auth = get_claude_auth() - - assert auth.api_key == "key" - - def test_uses_existing_auth_when_api_key_is_unset(self, monkeypatch): - monkeypatch.delenv("DSTACK_AGENT_ANTHROPIC_API_KEY", raising=False) - monkeypatch.setattr(shutil, "which", lambda _: "/usr/bin/claude") - - auth = get_claude_auth() - - assert auth.api_key is None - - @pytest.mark.parametrize("api_key", ["key", None]) - def test_builds_command_for_selected_auth_mode(self, api_key): - command = _build_claude_command(auth=_claude_auth(api_key=api_key, effort="high")) - - assert ("--bare" in command) is (api_key is not None) - assert ("--setting-sources" in command) is (api_key is None) - assert command[command.index("--effort") + 1] == "high" - - @pytest.mark.windows_only - def test_runs_windows_batch_launcher(self, tmp_path): - script = tmp_path / "fake-claude.cmd" - script.write_text("@echo off\necho batch-ok\n") - - result = subprocess.run( - _prepare_subprocess_command([str(script)]), - check=False, - capture_output=True, - text=True, - ) - - assert result.returncode == 0 - assert result.stdout.strip() == "batch-ok" - - -class TestAgentIsolation: - def test_inherits_only_required_environment(self, tmp_path, monkeypatch): - monkeypatch.setenv("PATH", "/usr/bin") - monkeypatch.setenv("HOME", "/home/test") - monkeypatch.setenv("UNRELATED_SECRET", "must-not-be-inherited") - api = SimpleNamespace( - project="main", - client=SimpleNamespace(base_url="http://127.0.0.1:3000"), - ) - - env = build_endpoint_agent_env( - api=api, - endpoint_env={"HF_TOKEN": "hf-secret"}, - auth=_claude_auth(), - workspace=EndpointAgentWorkspace( - path=tmp_path, - dstack_home=tmp_path / "home", - ), - token="dstack-secret", - ) - - assert env["DSTACK_SERVER_URL"] == "http://127.0.0.1:3000" - assert env["DSTACK_PROJECT"] == "main" - assert env["DSTACK_TOKEN"] == "dstack-secret" - assert env["HF_TOKEN"] == "hf-secret" - assert env["ANTHROPIC_API_KEY"] == "anthropic-secret" - assert env["HOME"] == str(tmp_path / "home") - assert "UNRELATED_SECRET" not in env - project = ConfigManager(tmp_path / "home" / ".dstack").get_project_config() - assert project is not None - assert project.name == "main" - assert project.url == "http://127.0.0.1:3000" - assert project.token == "dstack-secret" - - def test_creates_private_cli_home_and_dstack_wrapper(self): - with endpoint_agent_workspace() as workspace: - if not IS_WINDOWS: - assert (workspace.dstack_home / ".ssh").stat().st_mode & 0o777 == 0o700 - assert not (workspace.dstack_home / ".dstack" / "config.yml").exists() - dstack_command = shutil.which("dstack", path=str(workspace.bin_path)) - assert dstack_command is not None - result = subprocess.run( - [dstack_command, "--help"], - check=False, - capture_output=True, - text=True, - ) - assert result.returncode == 0 - assert "Usage: dstack" in result.stdout - - def test_keeps_control_socket_path_bounded(self): - with endpoint_agent_workspace() as workspace: - if not IS_WINDOWS: - socket_path = ( - workspace.dstack_home / ".dstack" / "ssh" / f"{'x' * 41}.control.sock" - ) - assert len(os.fsencode(socket_path)) <= 103 - - def test_detects_known_secret_in_generated_artifact(self): - assert contains_redacted_value( - {"commands": ["serve --token secret-token"]}, - ("secret-token",), - ) - - -class TestAgentSession: - def test_saves_log_and_debug_files(self, tmp_path, monkeypatch, capsys): - monkeypatch.setenv("HOME", str(tmp_path)) - monkeypatch.setenv("USERPROFILE", str(tmp_path)) - configuration = EndpointConfiguration( - name="qwen", - model={"base": "Qwen/Qwen3.5-27B"}, - max_price=0.5, - env=["HF_TOKEN", "TOKENIZERS_PARALLELISM=false"], - ) - - session = create_endpoint_agent_session(configuration) - - assert session.path.parent == tmp_path / ".dstack" / "agent" / "qwen" - assert session.path.name.endswith("-running") - assert {path.name for path in session.path.iterdir()} == {"agent.log"} - print_endpoint_progress("creating preset", agent_session=session) - assert "creating preset" in session.log_path.read_text() - assert "creating preset" in capsys.readouterr().out - if not IS_WINDOWS: - assert session.path.stat().st_mode & 0o777 == 0o700 - assert session.log_path.stat().st_mode & 0o777 == 0o600 - - debug_session = create_endpoint_agent_session(configuration, debug=True) - data = yaml.safe_load((debug_session.path / "endpoint.dstack.yml").read_text()) - - assert {path.name for path in debug_session.path.iterdir()} == { - "agent.log", - "endpoint.dstack.yml", - "trace.jsonl", - } - assert data["max_price"] == 0.5 - assert data["env"] == ["HF_TOKEN", "TOKENIZERS_PARALLELISM"] - assert "false" not in (debug_session.path / "endpoint.dstack.yml").read_text() - success_path = debug_session.finish("8f3a12c4") - assert success_path.name.endswith("-8f3a12c4") - - failed_session = create_endpoint_agent_session(configuration) - failed_path = failed_session.finish() - assert failed_path.name.endswith("-failed") - - def test_avoids_existing_session_directories(self, tmp_path): - timestamp = "20260714-120000-000000Z" - (tmp_path / f"{timestamp}-running").mkdir() - - path = _create_agent_session_directory(tmp_path, timestamp) - - assert path.name == f"{timestamp}-1-running" - - def test_finish_does_not_replace_existing_directory(self, tmp_path): - running = tmp_path / "20260714-120000-000000Z-running" - running.mkdir() - (running / "agent.log").touch() - existing = tmp_path / "20260714-120000-000000Z-failed" - existing.mkdir() - (existing / "keep").touch() - session = EndpointAgentSession( - path=running, - timestamp="20260714-120000-000000Z", - debug=False, - ) - - path = session.finish() - - assert path.name == "20260714-120000-000000Z-failed-1" - assert (existing / "keep").is_file() - - def test_reports_invalid_existing_parent(self, tmp_path, monkeypatch): - monkeypatch.setenv("HOME", str(tmp_path)) - monkeypatch.setenv("USERPROFILE", str(tmp_path)) - parent = tmp_path / ".dstack" / "agent" - parent.mkdir(parents=True) - (parent / "qwen").write_text("not a directory") - - with pytest.raises(CLIError, match="Could not create agent output"): - create_endpoint_agent_session( - EndpointConfiguration(name="qwen", model={"base": "Qwen/Qwen3.5-27B"}) - ) - - def test_log_write_failure_warns_once(self, tmp_path, capsys): - path = tmp_path / "agent-running" - path.mkdir() - (path / "agent.log").touch() - session = EndpointAgentSession( - path=path, - timestamp="20260714-120000-000000Z", - debug=False, - ) - shutil.rmtree(path) - - session.append_log("first") - session.append_log("second") - - assert capsys.readouterr().out.count("Could not write agent log") == 1 - - -class TestRedaction: - def test_does_not_replace_short_values_inside_diagnostics(self): - assert redact("DEBUG=1; enabled=false", ("1", "false")) == "DEBUG=1; enabled=false" - assert redact("false", ("false",)) == "[redacted]" - assert redact("token=secret-token", ("secret-token",)) == "token=[redacted]" - - -class TestAgentOutput: - @pytest.mark.asyncio - async def test_sends_prompt_and_redacts_raw_output(self, tmp_path, monkeypatch, capsys): - script = tmp_path / "fake_claude.py" - script.write_text( - """import json -import sys - -prompt = sys.stdin.read() -print(json.dumps({ - "type": "result", - "is_error": True, - "result": "bad secret-token", - "secret-token": "must be redacted", - "structured_output": {"prompt": prompt}, -})) -""" - ) - (tmp_path / "progress.jsonl").touch() - monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.agent._build_claude_command", - lambda **_: [sys.executable, str(script)], - ) - - workspace = EndpointAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home") - session_path = tmp_path / "debug-running" - session_path.mkdir() - (session_path / "agent.log").touch() - (session_path / "trace.jsonl").touch() - agent_session = EndpointAgentSession( - path=session_path, - timestamp="20260714-120000Z", - debug=True, - ) - output = await run_endpoint_agent( - prompt="full endpoint prompt", - env=os.environ.copy(), - workspace=workspace, - auth=_claude_auth(), - redacted_values=("secret-token",), - agent_session=agent_session, - ) - - assert output.report_data == {"prompt": "full endpoint prompt"} - assert output.error == "bad [redacted]" - trace = [json.loads(line) for line in agent_session.trace_path.read_text().splitlines()] - assert len(trace) == 1 - assert trace[0]["timestamp"].endswith("Z") - assert trace[0]["stream"] == "stdout" - assert trace[0]["event"]["result"] == "bad [redacted]" - assert "[redacted]" in trace[0]["event"] - assert capsys.readouterr().out == "" - - @pytest.mark.asyncio - async def test_accepts_stream_event_larger_than_64_kib(self, tmp_path, monkeypatch): - script = tmp_path / "fake_claude.py" - script.write_text( - """import json - -print(json.dumps({ - "type": "result", - "structured_output": {"value": "x" * (128 * 1024)}, -})) -""" - ) - (tmp_path / "progress.jsonl").touch() - session_path = tmp_path / "agent-running" - session_path.mkdir() - (session_path / "agent.log").touch() - monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.agent._build_claude_command", - lambda **_: [sys.executable, str(script)], - ) - - output = await run_endpoint_agent( - prompt="prompt", - env=os.environ.copy(), - workspace=EndpointAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home"), - auth=_claude_auth(), - redacted_values=(), - agent_session=EndpointAgentSession( - path=session_path, - timestamp="20260714-120000Z", - debug=False, - ), - ) - - assert output.report_data is not None - assert len(output.report_data["value"]) == 128 * 1024 - - def test_progress_stream_prints_only_redacted_messages(self, tmp_path, capsys): - progress_path = tmp_path / "progress.jsonl" - progress_path.write_text('{"message":"using secret-token"}\n') - session_path = tmp_path / "agent-running" - session_path.mkdir() - (session_path / "agent.log").touch() - agent_session = EndpointAgentSession( - path=session_path, - timestamp="20260714-120000Z", - debug=False, - ) - - _ProgressTailer( - path=progress_path, - redacted_values=("secret-token",), - agent_session=agent_session, - ).flush() - - output = capsys.readouterr().out - assert "using [redacted]" in output - assert "secret-token" not in output - log = agent_session.log_path.read_text() - assert "using [redacted]" in log - assert "secret-token" not in log - - -class TestProcessCleanup: - @pytest.mark.windows_only - @pytest.mark.asyncio - async def test_terminates_windows_process_tree(self): - proc = await asyncio.create_subprocess_exec( - sys.executable, - "-c", - ( - "import subprocess,sys,time; " - "p=subprocess.Popen([sys.executable,'-c','import time; time.sleep(60)']); " - "print(p.pid,flush=True); time.sleep(60)" - ), - stdout=asyncio.subprocess.PIPE, - ) - assert proc.stdout is not None - child_pid = int((await proc.stdout.readline()).decode()) - - await _terminate_process(proc) - - assert proc.returncode is not None - assert not psutil.pid_exists(child_pid) diff --git a/src/tests/_internal/cli/services/endpoints/test_apply.py b/src/tests/_internal/cli/services/endpoints/test_apply.py deleted file mode 100644 index c9962ffd3..000000000 --- a/src/tests/_internal/cli/services/endpoints/test_apply.py +++ /dev/null @@ -1,151 +0,0 @@ -from types import SimpleNamespace -from unittest.mock import Mock - -import pytest - -from dstack._internal.cli.models.endpoints import EndpointConfiguration -from dstack._internal.cli.services.endpoints.apply import ( - _build_service, - _get_matching_presets, - _select_plan, - apply_endpoint_preset, -) -from dstack._internal.core.models.instances import InstanceAvailability -from tests._internal.cli.endpoint_presets import get_endpoint_preset - -pytestmark = pytest.mark.windows - - -class TestGetMatchingPresets: - def test_matches_base_model_context_and_preset(self): - presets = [ - get_endpoint_preset(preset_id="small", context_length=4096), - get_endpoint_preset(preset_id="large", context_length=32768), - ] - configuration = EndpointConfiguration( - name="qwen", - model={"base": "Qwen/Qwen3.5-27B"}, - context_length=8192, - ) - - assert _get_matching_presets(presets, configuration=configuration, preset_id=None) == [ - presets[1] - ] - assert _get_matching_presets(presets, configuration=configuration, preset_id="large") == [ - presets[1] - ] - assert not _get_matching_presets(presets, configuration=configuration, preset_id="small") - - def test_exact_request_matches_repo_and_client_facing_name(self): - matching = get_endpoint_preset(preset_id="matching") - configuration = EndpointConfiguration( - name="qwen", - model={ - "repo": "community/Qwen3.5-27B-GPTQ-Int4", - "name": "Qwen/Qwen3.5-27B", - }, - ) - - assert _get_matching_presets([matching], configuration=configuration, preset_id=None) == [ - matching - ] - assert not _get_matching_presets( - [matching.copy(update={"model": "other/repo"})], - configuration=configuration, - preset_id=None, - ) - - -class TestBuildService: - def test_applies_endpoint_name_env_gateway_and_constraints(self): - configuration = EndpointConfiguration( - name="qwen-production", - model={"base": "Qwen/Qwen3.5-27B"}, - gateway="inference", - env={"HF_TOKEN": "token"}, - fleets=["gpu-fleet"], - max_price=1, - ) - - service = _build_service(configuration, get_endpoint_preset()) - - assert service.name == "qwen-production" - assert service.gateway == "inference" - assert service.env["HF_TOKEN"] == "token" - assert [fleet.format() for fleet in service.fleets] == ["gpu-fleet"] - assert service.max_price == 1 - - -class TestSelectPlan: - def test_selects_first_preset_with_available_offer(self): - presets = [ - get_endpoint_preset(preset_id="unavailable"), - get_endpoint_preset(preset_id="available"), - get_endpoint_preset(preset_id="never-probed"), - ] - configurator = Mock() - plans = [ - (_plan(InstanceAvailability.NOT_AVAILABLE), Mock()), - (_plan(InstanceAvailability.AVAILABLE), Mock()), - ] - configurator.get_plan.side_effect = plans - service_args = SimpleNamespace(profile=None) - - selected = _select_plan( - configuration=EndpointConfiguration(name="qwen", model={"base": "Qwen/Qwen3.5-27B"}), - configuration_path="endpoint.dstack.yml", - presets=presets, - configurator=configurator, - service_args=service_args, - ) - - assert selected.preset.id == "available" - assert selected.run_plan is plans[1][0] - assert selected.repo is plans[1][1] - assert configurator.get_plan.call_count == 2 - - def test_applies_the_selected_plan(self, monkeypatch): - preset = get_endpoint_preset() - run_plan = _plan(InstanceAvailability.AVAILABLE) - repo = Mock() - service_args = SimpleNamespace(profile="gpu") - configurator = Mock() - configurator.get_parser.return_value.parse_args.return_value = service_args - configurator.get_plan.return_value = (run_plan, repo) - monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.apply.ServiceConfigurator", - lambda api_client: configurator, - ) - command_args = SimpleNamespace() - - apply_endpoint_preset( - api=Mock(), - configuration=EndpointConfiguration( - name="qwen", - model={"base": "Qwen/Qwen3.5-27B"}, - ), - configuration_path="endpoint.dstack.yml", - preset_id=None, - profile_name="gpu", - command_args=command_args, - store=Mock(list=Mock(return_value=[preset])), - ) - - assert service_args.profile == "gpu" - configurator.apply_plan.assert_called_once_with( - run_plan=run_plan, - repo=repo, - command_args=command_args, - configurator_args=service_args, - plan_properties={ - "Model": "Qwen/Qwen3.5-27B ([secondary]base[/])", - "Preset": "8f3a12c4 " - "([secondary]context=32K, concurrency=1 42.1 tok/s TTFT 108ms[/])", - }, - ) - - -def _plan(availability: InstanceAvailability): - return SimpleNamespace( - job_plans=[SimpleNamespace(offers=[SimpleNamespace(availability=availability)])] - ) diff --git a/src/tests/_internal/cli/services/endpoints/test_create.py b/src/tests/_internal/cli/services/endpoints/test_create.py deleted file mode 100644 index 99edded5d..000000000 --- a/src/tests/_internal/cli/services/endpoints/test_create.py +++ /dev/null @@ -1,425 +0,0 @@ -import asyncio -import json -import uuid -from types import SimpleNamespace - -import pytest - -from dstack._internal.cli.models.endpoints import EndpointConfiguration -from dstack._internal.cli.services.endpoints.agent import ( - ClaudeAuth, - EndpointAgentProcessOutput, - EndpointAgentSession, - EndpointAgentWorkspace, - print_endpoint_progress, -) -from dstack._internal.cli.services.endpoints.create import ( - EndpointPresetCreateResult, - _build_prompt, - _cleanup_runs, - _create_endpoint_preset, - _get_build_name, - create_endpoint_preset, -) -from dstack._internal.cli.services.endpoints.store import EndpointPresetStore -from dstack._internal.core.errors import CLIError -from dstack._internal.core.models.envs import EnvSentinel -from dstack._internal.core.models.runs import Run, RunStatus -from tests._internal.cli.endpoint_presets import ( - get_endpoint_preset, - get_running_service_run, - get_successful_endpoint_report, -) - -pytestmark = pytest.mark.windows - - -def _claude_auth() -> ClaudeAuth: - return ClaudeAuth( - api_key="anthropic-secret", - executable="claude", - effort=None, - model="claude-test", - ) - - -@pytest.fixture -def creation_context(tmp_path, monkeypatch): - run = get_running_service_run() - run_apis = _FakeRunAPIs(run) - api = SimpleNamespace( - project="main", - runs=run_apis, - client=SimpleNamespace( - _token="dstack-secret", - base_url="http://127.0.0.1:3000", - runs=run_apis, - ), - ) - configuration = EndpointConfiguration( - name="qwen-build", - model={"base": "Qwen/Qwen3.5-27B"}, - context_length=8192, - fleets=["gpu-fleet"], - env={"LICENSE": "license-secret", "TOKENIZERS_PARALLELISM": "false"}, - ) - source_configuration = EndpointConfiguration( - name="qwen-build", - model={"base": "Qwen/Qwen3.5-27B"}, - context_length=8192, - fleets=["gpu-fleet"], - env=["LICENSE", "TOKENIZERS_PARALLELISM=false"], - ) - monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.get_claude_auth", - _claude_auth, - ) - monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create._get_build_name", - lambda _: "qwen-build", - ) - return SimpleNamespace( - api=api, - configuration=configuration, - source_configuration=source_configuration, - run=run, - run_apis=run_apis, - store=EndpointPresetStore(tmp_path / "presets"), - ) - - -class TestCreateEndpointPreset: - def test_saves_agent_log_without_debug(self, tmp_path, monkeypatch, capsys): - monkeypatch.setenv("HOME", str(tmp_path)) - monkeypatch.setenv("USERPROFILE", str(tmp_path)) - preset = get_endpoint_preset() - - async def create(**kwargs): - print_endpoint_progress("testing preset", agent_session=kwargs["agent_session"]) - return EndpointPresetCreateResult( - preset=preset, - path=tmp_path / "preset.yaml", - final_run_id=uuid.uuid4(), - final_run_name="qwen-build-2", - ) - - monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create._create_endpoint_preset", - create, - ) - - create_endpoint_preset( - api=SimpleNamespace(), - configuration=EndpointConfiguration( - name="qwen", - model={"base": "Qwen/Qwen3.5-27B"}, - ), - store=EndpointPresetStore(tmp_path / "presets"), - ) - - paths = list((tmp_path / ".dstack" / "agent" / "qwen").iterdir()) - assert len(paths) == 1 - assert paths[0].name.endswith(f"-{preset.id}") - assert {path.name for path in paths[0].iterdir()} == {"agent.log"} - assert "testing preset" in (paths[0] / "agent.log").read_text() - output = capsys.readouterr().out.replace("\n", "") - assert f"Agent log saved to {paths[0] / 'agent.log'}" in output - - def test_debug_finalization_error_does_not_mask_success(self, tmp_path, monkeypatch, capsys): - monkeypatch.setenv("HOME", str(tmp_path)) - monkeypatch.setenv("USERPROFILE", str(tmp_path)) - monkeypatch.setenv("HF_TOKEN", "hf-secret") - preset = get_endpoint_preset() - - async def create(**kwargs): - assert kwargs["configuration"].env.as_dict() == { - "HF_TOKEN": "hf-secret", - "TOKENIZERS_PARALLELISM": "false", - } - assert isinstance(kwargs["source_configuration"].env["HF_TOKEN"], EnvSentinel) - assert kwargs["source_configuration"].env["TOKENIZERS_PARALLELISM"] == "false" - kwargs["agent_session"].write_prompt("test prompt") - return EndpointPresetCreateResult( - preset=preset, - path=tmp_path / "preset.yaml", - final_run_id=uuid.uuid4(), - final_run_name="qwen-build-2", - ) - - def fail_finish(self, preset_id=None): - raise OSError("rename failed") - - monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create._create_endpoint_preset", - create, - ) - monkeypatch.setattr(EndpointAgentSession, "finish", fail_finish) - - result = create_endpoint_preset( - api=SimpleNamespace(), - configuration=EndpointConfiguration( - name="qwen", - model={"base": "Qwen/Qwen3.5-27B"}, - env=["HF_TOKEN", "TOKENIZERS_PARALLELISM=false"], - ), - store=EndpointPresetStore(tmp_path / "presets"), - debug=True, - ) - - paths = list((tmp_path / ".dstack" / "agent" / "qwen").iterdir()) - assert len(paths) == 1 - assert paths[0].name.endswith("-running") - assert result.preset == preset - assert {path.name for path in paths[0].iterdir()} == { - "endpoint.dstack.yml", - "agent.log", - "prompt.md", - "trace.jsonl", - } - assert "hf-secret" not in (paths[0] / "endpoint.dstack.yml").read_text() - assert "Files remain at" in capsys.readouterr().out - - def test_debug_finalization_does_not_mask_creation_error(self, tmp_path, monkeypatch): - monkeypatch.setenv("HOME", str(tmp_path)) - monkeypatch.setenv("USERPROFILE", str(tmp_path)) - - async def create(**kwargs): - raise RuntimeError("creation failed") - - def fail_finish(self, preset_id=None): - raise OSError("rename failed") - - monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create._create_endpoint_preset", - create, - ) - monkeypatch.setattr(EndpointAgentSession, "finish", fail_finish) - - with pytest.raises(RuntimeError, match="creation failed"): - create_endpoint_preset( - api=SimpleNamespace(), - configuration=EndpointConfiguration( - name="qwen", - model={"base": "Qwen/Qwen3.5-27B"}, - ), - store=EndpointPresetStore(tmp_path / "presets"), - debug=True, - ) - - @pytest.mark.asyncio - async def test_checks_active_fleets_before_claude_auth(self, tmp_path, monkeypatch): - api = SimpleNamespace( - project="main", - client=SimpleNamespace(fleets=SimpleNamespace(list=lambda *args, **kwargs: [])), - ) - monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.get_claude_auth", - lambda: pytest.fail("Claude auth must not be checked without an active fleet"), - ) - - with pytest.raises(CLIError, match="no active fleets"): - await _create_endpoint_preset( - api=api, - configuration=EndpointConfiguration( - name="qwen-build", - model={"base": "Qwen/Qwen3.5-27B"}, - ), - store=EndpointPresetStore(tmp_path / "presets"), - agent_session=_agent_session(tmp_path), - ) - - @pytest.mark.asyncio - async def test_cleans_up_runs_when_cancelled(self, creation_context, monkeypatch, tmp_path): - async def run_agent(**_): - raise asyncio.CancelledError - - cleanup_calls = [] - - async def cleanup_runs(**kwargs): - cleanup_calls.append(kwargs) - - monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.run_endpoint_agent", - run_agent, - ) - monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create._cleanup_runs", - cleanup_runs, - ) - - with pytest.raises(asyncio.CancelledError): - await _create_endpoint_preset( - api=creation_context.api, - configuration=creation_context.configuration, - source_configuration=creation_context.source_configuration, - store=creation_context.store, - agent_session=_agent_session(tmp_path), - ) - - assert len(cleanup_calls) == 1 - assert cleanup_calls[0]["build_name"] == "qwen-build" - - @pytest.mark.parametrize( - ("keep_service", "stopped_names"), - [(False, ["qwen-build-2"]), (True, [])], - ) - @pytest.mark.asyncio - async def test_saves_preset_and_cleans_up_runs( - self, creation_context, monkeypatch, keep_service, stopped_names, tmp_path - ): - session_path = tmp_path / "debug-running" - session_path.mkdir() - (session_path / "agent.log").touch() - (session_path / "trace.jsonl").touch() - agent_session = EndpointAgentSession( - path=session_path, - timestamp="20260714-120000Z", - debug=True, - ) - - async def run_agent(**kwargs): - assert kwargs["agent_session"] is agent_session - assert (session_path / "prompt.md").is_file() - return EndpointAgentProcessOutput( - report_data=json.loads(get_successful_endpoint_report(creation_context.run).json()) - ) - - monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.run_endpoint_agent", - run_agent, - ) - result = await _create_endpoint_preset( - api=creation_context.api, - configuration=creation_context.configuration, - source_configuration=creation_context.source_configuration, - store=creation_context.store, - keep_service=keep_service, - agent_session=agent_session, - ) - - assert result.preset.base == "Qwen/Qwen3.5-27B" - assert result.path.is_file() - assert creation_context.store.list() == [result.preset] - assert "license-secret" not in result.path.read_text() - assert result.preset.service.env["TOKENIZERS_PARALLELISM"] == "false" - assert creation_context.run_apis.stopped_names == stopped_names - - -class TestBuildName: - def test_requires_name_and_keeps_generated_prefix_bounded(self, monkeypatch): - with pytest.raises(CLIError, match="Endpoint name is required"): - _get_build_name(None) - monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.secrets.token_hex", - lambda _: "a1b2c3", - ) - - build_name = _get_build_name("qwen-endpoint-with-a-name-that-is-forty-one") - - assert len(f"{build_name}-99999") <= 41 - - -class TestBuildPrompt: - def test_distinguishes_base_from_exact_model(self): - base_prompt = _build_prompt( - configuration=EndpointConfiguration( - name="qwen", - model={"base": "Qwen/Qwen3.5-27B"}, - context_length=8192, - ), - build_name="qwen-build", - allowed_fleets=("gpu-fleet",), - ) - exact_prompt = _build_prompt( - configuration=EndpointConfiguration( - name="qwen", - model={ - "repo": "community/Qwen3.5-27B-GPTQ-Int4", - "name": "Qwen/Qwen3.5-27B", - }, - ), - build_name="qwen-build", - allowed_fleets=("gpu-fleet",), - ) - - assert "- base_model: Qwen/Qwen3.5-27B" in base_prompt - assert "- model_repo:" not in base_prompt - assert "- context_length: 8192" in base_prompt - assert "- model_repo: community/Qwen3.5-27B-GPTQ-Int4" in exact_prompt - assert "- base_model:" not in exact_prompt - assert "- context_length:" not in exact_prompt - - -class TestCleanupRuns: - @pytest.mark.asyncio - async def test_stops_only_recorded_build_runs(self, tmp_path, monkeypatch): - (tmp_path / "submissions.jsonl").write_text( - '{"name":"qwen-build-1"}\n{"name":"unrelated-run"}\n{"name":"qwen-build-2"}\n' - ) - runs = _FakeRuns() - api = SimpleNamespace( - runs=runs, - project="main", - client=SimpleNamespace(runs=runs), - ) - - async def no_sleep(_): - return None - - monkeypatch.setattr(asyncio, "sleep", no_sleep) - await _cleanup_runs( - api=api, - build_name="qwen-build", - workspace=EndpointAgentWorkspace( - path=tmp_path, - dstack_home=tmp_path / "home", - ), - final_run_name="qwen-build-2", - keep_final_service=True, - agent_session=_agent_session(tmp_path), - ) - - assert runs.stopped_names == ["qwen-build-1"] - - -def _agent_session(tmp_path, *, debug: bool = False) -> EndpointAgentSession: - path = tmp_path / "agent-running" - path.mkdir() - (path / "agent.log").touch() - if debug: - (path / "trace.jsonl").touch() - return EndpointAgentSession( - path=path, - timestamp="20260714-120000Z", - debug=debug, - ) - - -class _FakeRuns: - def __init__(self): - self.stopped_names: list[str] = [] - - def get(self, name): - status = RunStatus.TERMINATED if name in self.stopped_names else RunStatus.RUNNING - return SimpleNamespace(status=status) - - def stop(self, project, names, abort): - assert project == "main" - assert abort is False - self.stopped_names.extend(names) - - -class _FakeRunAPIs: - def __init__(self, run: Run): - self.run = run - self.stopped_names: list[str] = [] - - def get(self, *args): - name = args[-1] - return self.run if name == self.run.run_spec.run_name else None - - def stop(self, project, names, abort): - assert project == "main" - assert abort is False - self.stopped_names.extend(names) - self.run.status = RunStatus.TERMINATED diff --git a/src/tests/_internal/cli/services/endpoints/test_store.py b/src/tests/_internal/cli/services/endpoints/test_store.py deleted file mode 100644 index 6088ec8cd..000000000 --- a/src/tests/_internal/cli/services/endpoints/test_store.py +++ /dev/null @@ -1,127 +0,0 @@ -from pathlib import Path - -import pytest -import yaml -from pydantic import ValidationError - -from dstack._internal.cli.models.endpoint_agent import AGENT_FINAL_REPORT_JSON_SCHEMA -from dstack._internal.cli.models.endpoint_presets import ( - EndpointBenchmark, - EndpointBenchmarkLatency, - EndpointBenchmarkMetrics, - EndpointBenchmarkWorkload, -) -from dstack._internal.cli.services.endpoints.store import EndpointPresetStore -from dstack._internal.core.errors import CLIError -from dstack._internal.core.models.envs import EnvSentinel -from tests._internal.cli.endpoint_presets import ( - get_endpoint_benchmark, - get_endpoint_preset, -) - -pytestmark = pytest.mark.windows - - -class TestEndpointBenchmark: - def test_agent_schema_matches_benchmark_model(self): - schema = AGENT_FINAL_REPORT_JSON_SCHEMA["properties"]["benchmark"] - assert set(schema["properties"]) == set(EndpointBenchmark.__fields__) - { - "target", - "client", - } - assert set(schema["required"]) == set(schema["properties"]) - workload_schema = schema["properties"]["workload"] - metrics_schema = schema["properties"]["metrics"] - assert set(workload_schema["properties"]) == set(EndpointBenchmarkWorkload.__fields__) - assert set(workload_schema["required"]) == set(workload_schema["properties"]) - assert set(metrics_schema["properties"]) == set(EndpointBenchmarkMetrics.__fields__) - assert set(metrics_schema["required"]) == set(metrics_schema["properties"]) - assert set(metrics_schema["properties"]["ttft_ms"]["properties"]) == set( - EndpointBenchmarkLatency.__fields__ - ) - - @pytest.mark.parametrize( - ("field", "value", "error"), - [ - ("failed_requests", 1, "must not include failed requests"), - ("successful_requests", 15, "must match workload.num_requests"), - ], - ) - def test_rejects_inconsistent_successful_metrics(self, field, value, error): - data = get_endpoint_benchmark().dict() - data["metrics"][field] = value - with pytest.raises(ValidationError, match=error): - EndpointBenchmark.parse_obj(data) - - def test_rejects_tool_specific_metrics(self): - data = get_endpoint_benchmark().dict() - data["metrics"]["tool_specific"] = 1 - - with pytest.raises(ValidationError, match="extra fields not permitted"): - EndpointBenchmark.parse_obj(data) - - -class TestEndpointPresetStore: - def test_saves_and_lists_self_contained_preset(self, tmp_path: Path): - store = EndpointPresetStore(tmp_path / "presets") - preset = get_endpoint_preset() - - path = store.save(preset) - - assert path == (tmp_path / "presets" / "models--Qwen--Qwen3.5-27B" / "8f3a12c4.yaml") - data = yaml.safe_load(path.read_text()) - assert data["base"] == preset.base - assert data["id"] == preset.id - assert data["model"] == preset.model - assert data["created_at"] == preset.created_at.isoformat() - assert "presets" not in data - assert store.list() == [preset] - assert store.get(preset.id) == preset - assert not list(path.parent.glob("*.tmp")) - - def test_replaces_same_preset_id_atomically(self, tmp_path: Path): - store = EndpointPresetStore(tmp_path / "presets") - preset = get_endpoint_preset() - store.save(preset) - - updated = preset.copy(update={"context_length": 16384}) - store.save(updated) - - assert store.get(updated.id) == updated - - def test_rejects_duplicate_preset_id(self, tmp_path: Path): - store = EndpointPresetStore(tmp_path / "presets") - preset = get_endpoint_preset() - store.save(preset) - store.save(preset.copy(update={"base": "Qwen/Another-Model"})) - - with pytest.raises(CLIError, match="is not unique"): - store.get(preset.id) - - def test_rejects_preset_without_successful_benchmark(self, tmp_path: Path): - store = EndpointPresetStore(tmp_path / "presets") - path = store.save(get_endpoint_preset()) - data = yaml.safe_load(path.read_text()) - data["validations"][0].pop("benchmark") - path.write_text(yaml.safe_dump(data, sort_keys=False)) - - with pytest.raises(CLIError, match="benchmark"): - store.list() - - def test_preserves_literal_env_values(self, tmp_path: Path): - store = EndpointPresetStore(tmp_path / "presets") - preset = get_endpoint_preset() - preset.service.env.update( - { - "TOKENIZERS_PARALLELISM": "false", - "MODEL_LABEL": "monkey", - "HF_TOKEN": EnvSentinel(key="HF_TOKEN"), - } - ) - - path = store.save(preset) - env = yaml.safe_load(path.read_text())["service"]["env"] - - assert "TOKENIZERS_PARALLELISM=false" in env - assert "MODEL_LABEL=monkey" in env - assert "HF_TOKEN" in env diff --git a/src/tests/_internal/cli/services/endpoints/test_verify.py b/src/tests/_internal/cli/services/endpoints/test_verify.py deleted file mode 100644 index 29452b4b8..000000000 --- a/src/tests/_internal/cli/services/endpoints/test_verify.py +++ /dev/null @@ -1,82 +0,0 @@ -from datetime import datetime, timezone -from unittest.mock import patch - -import pytest -from pydantic import ValidationError - -from dstack._internal.cli.models.endpoint_agent import AgentFinalReport -from dstack._internal.cli.models.endpoints import EndpointConfiguration -from dstack._internal.cli.services.endpoints.verify import ( - build_verified_endpoint_preset, -) -from dstack._internal.core.errors import CLIError -from dstack._internal.core.models.envs import EnvSentinel -from dstack._internal.core.models.profiles import ProfileParams -from tests._internal.cli.endpoint_presets import ( - get_running_service_run, - get_successful_endpoint_report, -) - -pytestmark = pytest.mark.windows - - -class TestBuildVerifiedEndpointPreset: - def test_successful_report_requires_benchmark(self): - run = get_running_service_run() - data = get_successful_endpoint_report(run).dict() - data.pop("benchmark") - - with pytest.raises(ValidationError, match="benchmark"): - AgentFinalReport.parse_obj(data) - - def test_builds_portable_self_contained_preset(self): - run = get_running_service_run() - created_at = datetime(2026, 1, 2, 3, 4, tzinfo=timezone.utc) - - with patch( - "dstack._internal.cli.services.endpoints.presets.get_current_datetime", - return_value=created_at, - ): - preset = build_verified_endpoint_preset( - run=run, - endpoint_configuration=EndpointConfiguration( - name="qwen-build", - model={"base": "Qwen/Qwen3.5-27B"}, - context_length=8192, - gateway="benchmark-gateway", - env=["LICENSE", "TOKENIZERS_PARALLELISM=false"], - ), - report=get_successful_endpoint_report(run), - ) - - assert preset.base == "Qwen/Qwen3.5-27B" - assert preset.model == "community/Qwen3.5-27B-GPTQ-Int4" - assert preset.context_length == 32768 - assert preset.created_at == created_at - assert preset.service.name is None - assert preset.service.gateway is None - assert all(getattr(preset.service, field) is None for field in ProfileParams.__fields__) - assert isinstance(preset.service.env["LICENSE"], EnvSentinel) - assert preset.service.env["TOKENIZERS_PARALLELISM"] == "false" - assert preset.service.resources.gpu.vendor.value == "nvidia" - validation = preset.validations[0] - assert validation.replicas[0].resources[0].gpu.name == ["A6000"] - assert validation.benchmark.target.type == "server-proxy" - assert validation.benchmark.client.type == "local" - - def test_rejects_variant_for_exact_model_request(self): - run = get_running_service_run() - report = get_successful_endpoint_report(run).copy(update={"model": "other/model"}) - - with pytest.raises(CLIError, match="changed an exact model request"): - build_verified_endpoint_preset( - run=run, - endpoint_configuration=EndpointConfiguration( - name="qwen-build", - model={ - "repo": "community/Qwen3.5-27B-GPTQ-Int4", - "name": "Qwen/Qwen3.5-27B", - }, - ), - report=report, - ) diff --git a/src/tests/_internal/cli/services/endpoints/__init__.py b/src/tests/_internal/cli/services/presets/__init__.py similarity index 100% rename from src/tests/_internal/cli/services/endpoints/__init__.py rename to src/tests/_internal/cli/services/presets/__init__.py diff --git a/src/tests/_internal/cli/services/presets/test_agent.py b/src/tests/_internal/cli/services/presets/test_agent.py new file mode 100644 index 000000000..9e96c4412 --- /dev/null +++ b/src/tests/_internal/cli/services/presets/test_agent.py @@ -0,0 +1,969 @@ +import asyncio +import json +import os +import shutil +import signal +import subprocess +import sys +from contextlib import suppress +from pathlib import Path +from types import SimpleNamespace + +import psutil +import pytest +import yaml + +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.services.presets.agent import ( + ClaudeAuth, + _build_claude_command, + _prepare_subprocess_command, + _terminate_process, + build_preset_agent_env, + get_claude_auth, + run_preset_agent, +) +from dstack._internal.cli.services.presets.redaction import ( + contains_redacted_value, + redact, +) +from dstack._internal.cli.services.presets.session import ( + PresetAgentSession, + _summarize_session_trials, + create_preset_agent_session, + load_resumable_agent_session, + print_preset_progress, +) +from dstack._internal.cli.services.presets.tail import ( + _ProgressTailer, + _RecordMirror, +) +from dstack._internal.cli.services.presets.workspace import ( + PresetAgentWorkspace, + attach_agent_workspace, + create_agent_workspace, + remove_agent_workspace, +) +from dstack._internal.compat import IS_WINDOWS +from dstack._internal.core.errors import CLIError +from dstack._internal.core.services.configs import ConfigManager + +pytestmark = pytest.mark.windows + + +def _claude_auth(*, api_key: str | None = "anthropic-secret", effort=None) -> ClaudeAuth: + return ClaudeAuth( + api_key=api_key, + executable="claude", + effort=effort, + model="claude-test", + ) + + +class TestClaudeAuth: + @pytest.mark.parametrize("api_key_env", ["key", None]) + def test_uses_api_key_only_when_env_is_set(self, monkeypatch, api_key_env): + if api_key_env is None: + monkeypatch.delenv("DSTACK_AGENT_ANTHROPIC_API_KEY", raising=False) + else: + monkeypatch.setenv("DSTACK_AGENT_ANTHROPIC_API_KEY", api_key_env) + monkeypatch.setattr(shutil, "which", lambda _: "/usr/bin/claude") + + auth = get_claude_auth() + + assert auth.api_key == api_key_env + + @pytest.mark.parametrize("api_key", ["key", None]) + def test_builds_command_for_selected_auth_mode(self, api_key): + command = _build_claude_command(auth=_claude_auth(api_key=api_key, effort="high")) + + assert ("--bare" in command) is (api_key is not None) + assert ("--setting-sources" in command) is (api_key is None) + assert command[command.index("--effort") + 1] == "high" + + @pytest.mark.windows_only + def test_runs_windows_batch_launcher(self, tmp_path): + script = tmp_path / "fake-claude.cmd" + script.write_text("@echo off\necho batch-ok\n") + + result = subprocess.run( + _prepare_subprocess_command([str(script)]), + check=False, + capture_output=True, + text=True, + ) + + assert result.returncode == 0 + assert result.stdout.strip() == "batch-ok" + + +class TestAgentIsolation: + def test_inherits_only_required_environment(self, tmp_path, monkeypatch): + monkeypatch.setenv("PATH", "/usr/bin") + monkeypatch.setenv("HOME", "/home/test") + monkeypatch.setenv("UNRELATED_SECRET", "must-not-be-inherited") + api = SimpleNamespace( + project="main", + client=SimpleNamespace(base_url="http://127.0.0.1:3000"), + ) + + env = build_preset_agent_env( + api=api, + preset_env={"HF_TOKEN": "hf-secret"}, + auth=_claude_auth(), + workspace=PresetAgentWorkspace( + path=tmp_path, + dstack_home=tmp_path / "home", + ), + token="dstack-secret", + ) + + assert env["DSTACK_SERVER_URL"] == "http://127.0.0.1:3000" + assert env["DSTACK_PROJECT"] == "main" + assert env["DSTACK_TOKEN"] == "dstack-secret" + assert env["HF_TOKEN"] == "hf-secret" + assert env["ANTHROPIC_API_KEY"] == "anthropic-secret" + assert env["HOME"] == str(tmp_path / "home") + assert "UNRELATED_SECRET" not in env + project = ConfigManager(tmp_path / "home" / ".dstack").get_project_config() + assert project is not None + assert project.name == "main" + assert project.url == "http://127.0.0.1:3000" + assert project.token == "dstack-secret" + + def test_creates_private_cli_home_and_dstack_wrapper(self, tmp_path): + workspace = _session_workspace(tmp_path) + if not IS_WINDOWS: + assert (workspace.dstack_home / ".ssh").stat().st_mode & 0o777 == 0o700 + assert not (workspace.dstack_home / ".dstack" / "config.yml").exists() + dstack_command = shutil.which("dstack", path=str(workspace.bin_path)) + assert dstack_command is not None + result = subprocess.run( + [dstack_command, "--help"], + check=False, + capture_output=True, + text=True, + ) + assert result.returncode == 0 + assert "Usage: dstack" in result.stdout + + def test_keeps_control_socket_path_bounded(self, tmp_path): + workspace = _session_workspace(tmp_path) + if not IS_WINDOWS: + socket_path = workspace.dstack_home / ".dstack" / "ssh" / f"{'x' * 41}.control.sock" + assert len(os.fsencode(socket_path)) <= 103 + + def test_detects_known_secret_in_generated_artifact(self): + assert contains_redacted_value( + {"commands": ["serve --token secret-token"]}, + ("secret-token",), + ) + + +def _session_workspace(tmp_path): + session_dir = tmp_path / "session-under-test" + session_dir.mkdir() + session = PresetAgentSession(path=session_dir, debug=False, preset_id="abcd1234") + return create_agent_workspace(session) + + +class TestAgentSession: + def _configuration(self) -> PresetConfiguration: + return PresetConfiguration( + name="qwen", + model={"base": "Qwen/Qwen3.5-27B"}, + max_price=0.5, + env=["HF_TOKEN", "TOKENIZERS_PARALLELISM=false"], + ) + + def _home(self, tmp_path, monkeypatch) -> None: + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + + def test_creates_private_session_with_log_and_manifest(self, tmp_path, monkeypatch, capsys): + self._home(tmp_path, monkeypatch) + + session = create_preset_agent_session(self._configuration()) + + assert session.path.parent == tmp_path / ".dstack" / "presets" + assert session.path.name == session.preset_id + assert len(session.preset_id) == 8 + assert {path.name for path in session.path.iterdir()} == { + "agent.log", + "session.json", + "preset.dstack.yml", + } + manifest = json.loads((session.path / "session.json").read_text()) + assert manifest["id"] == session.preset_id + assert manifest["status"] == "running" + assert manifest["name"] == "qwen" + assert manifest["model"] == "Qwen/Qwen3.5-27B" + assert manifest["pid"] == os.getpid() + print_preset_progress("creating preset", agent_session=session) + assert "creating preset" in session.log_path.read_text() + assert "creating preset" in capsys.readouterr().out + if not IS_WINDOWS: + assert session.path.stat().st_mode & 0o777 == 0o700 + assert session.log_path.stat().st_mode & 0o777 == 0o600 + + def test_debug_session_saves_scrubbed_configuration_and_trace(self, tmp_path, monkeypatch): + self._home(tmp_path, monkeypatch) + + debug_session = create_preset_agent_session(self._configuration(), debug=True) + + data = yaml.safe_load((debug_session.path / "preset.dstack.yml").read_text()) + assert {path.name for path in debug_session.path.iterdir()} == { + "agent.log", + "preset.dstack.yml", + "session.json", + "trace.jsonl", + } + assert data["max_price"] == 0.5 + assert data["env"] == ["HF_TOKEN", "TOKENIZERS_PARALLELISM"] + assert "false" not in (debug_session.path / "preset.dstack.yml").read_text() + + @pytest.mark.parametrize("status", ["success", "failed"]) + def test_finish_records_terminal_status(self, tmp_path, monkeypatch, status): + self._home(tmp_path, monkeypatch) + session = create_preset_agent_session(self._configuration()) + + finished_path = session.finish(status) + + assert finished_path == session.path + assert json.loads((session.path / "session.json").read_text())["status"] == status + + def test_finish_writes_status_in_place(self, tmp_path): + session_dir = tmp_path / "20260714-120000-000000Z" + session_dir.mkdir() + (session_dir / "agent.log").touch() + session = PresetAgentSession( + path=session_dir, + debug=False, + ) + + path = session.finish("failed") + + assert path == session_dir + assert json.loads((session_dir / "session.json").read_text()) == {"status": "failed"} + + def test_reports_invalid_existing_parent(self, tmp_path, monkeypatch): + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + (tmp_path / ".dstack").mkdir(parents=True) + (tmp_path / ".dstack" / "presets").write_text("not a directory") + + with pytest.raises(CLIError, match="Could not create agent output"): + create_preset_agent_session( + PresetConfiguration(name="qwen", model={"base": "Qwen/Qwen3.5-27B"}) + ) + + def test_log_write_failure_warns_once(self, tmp_path, capsys): + path = tmp_path / "agent-running" + path.mkdir() + (path / "agent.log").touch() + session = PresetAgentSession( + path=path, + debug=False, + ) + shutil.rmtree(path) + + session.append_log("first") + session.append_log("second") + + assert capsys.readouterr().out.count("Could not write agent log") == 1 + + +class TestRedaction: + def test_does_not_replace_short_values_inside_diagnostics(self): + assert redact("DEBUG=1; enabled=false", ("1", "false")) == "DEBUG=1; enabled=false" + assert redact("false", ("false",)) == "[redacted]" + assert redact("token=secret-token", ("secret-token",)) == "token=[redacted]" + + +class TestAgentOutput: + @pytest.mark.asyncio + async def test_sends_prompt_and_redacts_raw_output(self, tmp_path, monkeypatch, capsys): + script = tmp_path / "fake_claude.py" + script.write_text( + """import json +import sys + +prompt = sys.stdin.read() +print(json.dumps({ + "type": "result", + "is_error": True, + "result": "bad secret-token", + "secret-token": "must be redacted", + "structured_output": {"prompt": prompt}, +})) +""" + ) + (tmp_path / "progress.jsonl").touch() + monkeypatch.setattr( + "dstack._internal.cli.services.presets.agent._build_claude_command", + lambda **_: [sys.executable, str(script)], + ) + + workspace = PresetAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home") + session_path = tmp_path / "debug-running" + session_path.mkdir() + (session_path / "agent.log").touch() + (session_path / "trace.jsonl").touch() + agent_session = PresetAgentSession( + path=session_path, + debug=True, + ) + output = await run_preset_agent( + prompt="full preset prompt", + env=os.environ.copy(), + workspace=workspace, + auth=_claude_auth(), + redacted_values=("secret-token",), + agent_session=agent_session, + ) + + assert output.report_data == {"prompt": "full preset prompt"} + assert output.error == "bad [redacted]" + trace = [json.loads(line) for line in agent_session.trace_path.read_text().splitlines()] + assert len(trace) == 1 + assert trace[0]["timestamp"].endswith("Z") + assert trace[0]["stream"] == "stdout" + assert trace[0]["event"]["result"] == "bad [redacted]" + assert "[redacted]" in trace[0]["event"] + assert capsys.readouterr().out == "" + + @pytest.mark.asyncio + async def test_accepts_stream_event_larger_than_64_kib(self, tmp_path, monkeypatch): + script = tmp_path / "fake_claude.py" + script.write_text( + """import json + +print(json.dumps({ + "type": "result", + "structured_output": {"value": "x" * (128 * 1024)}, +})) +""" + ) + (tmp_path / "progress.jsonl").touch() + session_path = tmp_path / "agent-running" + session_path.mkdir() + (session_path / "agent.log").touch() + monkeypatch.setattr( + "dstack._internal.cli.services.presets.agent._build_claude_command", + lambda **_: [sys.executable, str(script)], + ) + + output = await run_preset_agent( + prompt="prompt", + env=os.environ.copy(), + workspace=PresetAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home"), + auth=_claude_auth(), + redacted_values=(), + agent_session=PresetAgentSession( + path=session_path, + debug=False, + ), + ) + + assert output.report_data is not None + assert len(output.report_data["value"]) == 128 * 1024 + + def test_progress_stream_prints_only_redacted_messages(self, tmp_path, capsys): + progress_path = tmp_path / "progress.jsonl" + progress_path.write_text('{"message":"using secret-token"}\n') + session_path = tmp_path / "agent-running" + session_path.mkdir() + (session_path / "agent.log").touch() + agent_session = PresetAgentSession( + path=session_path, + debug=False, + ) + + _ProgressTailer( + path=progress_path, + redacted_values=("secret-token",), + agent_session=agent_session, + ).flush() + + output = capsys.readouterr().out + assert "using [redacted]" in output + assert "secret-token" not in output + log = agent_session.log_path.read_text() + assert "using [redacted]" in log + assert "secret-token" not in log + + +class TestProcessCleanup: + @pytest.mark.windows_only + @pytest.mark.asyncio + async def test_terminates_windows_process_tree(self): + proc = await asyncio.create_subprocess_exec( + sys.executable, + "-c", + ( + "import subprocess,sys,time; " + "p=subprocess.Popen([sys.executable,'-c','import time; time.sleep(60)']); " + "print(p.pid,flush=True); time.sleep(60)" + ), + stdout=asyncio.subprocess.PIPE, + ) + assert proc.stdout is not None + child_pid = int((await proc.stdout.readline()).decode()) + + await _terminate_process(proc) + + assert proc.returncode is not None + assert not psutil.pid_exists(child_pid) + + +class TestRecordMirror: + def test_mirrors_complete_lines_redacted(self, tmp_path): + source = tmp_path / "runs.jsonl" + target = tmp_path / "mirror" / "runs.jsonl" + target.parent.mkdir() + mirror = _RecordMirror(source=source, target=target, redacted_values=["dstack-secret"]) + + source.write_text( + '{"name":"run-1","note":"dstack-secret"}\n{"name":"run-2"', encoding="utf-8" + ) + mirror.flush() + + assert target.read_text() == '{"name":"run-1","note":"[redacted]"}\n' + + with source.open("a", encoding="utf-8") as f: + f.write("}\n") + mirror.flush() + + assert target.read_text().splitlines() == [ + '{"name":"run-1","note":"[redacted]"}', + '{"name":"run-2"}', + ] + + def test_missing_source_is_no_op(self, tmp_path): + mirror = _RecordMirror( + source=tmp_path / "absent.jsonl", + target=tmp_path / "target.jsonl", + redacted_values=[], + ) + + mirror.flush() + + assert not (tmp_path / "target.jsonl").exists() + + +class TestWriteAgentInfo: + def test_writes_model_params_and_auth(self, tmp_path, monkeypatch): + monkeypatch.setattr( + "dstack._internal.cli.services.presets.agent._get_claude_version", + lambda auth: "2.1.0 (Claude Code)", + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.agent._get_claude_auth_status", + lambda auth: {"authMethod": "claude.ai", "loggedIn": True}, + ) + session_dir = tmp_path / "session" + session_dir.mkdir() + session = PresetAgentSession(path=session_dir, debug=True) + + session.write_agent_info( + ClaudeAuth(api_key=None, executable="claude", effort=None, model="claude-opus-4-8") + ) + + assert json.loads((session_dir / "agent.json").read_text()) == { + "executable": "claude", + "version": "2.1.0 (Claude Code)", + "model": {"name": "claude-opus-4-8", "effort": "default"}, + "auth": {"authMethod": "claude.ai", "loggedIn": True}, + } + + +def _subprocess_env() -> dict[str, str]: + # A minimal realistic agent env: build_preset_agent_env never produces an + # empty dict, and env={} crashes CreateProcess on Windows 3.10, which lacks + # the cpython gh-105436 fix. + names = ("PATH", "SYSTEMROOT", "SYSTEMDRIVE", "COMSPEC", "TEMP", "TMP") + return {name: value for name in names if (value := os.environ.get(name))} + + +def _write_fake_claude(tmp_path, script_body: str) -> Path: + script = tmp_path / "fake_claude.py" + script.write_text(script_body) + return script + + +def _agent_setup(tmp_path): + (tmp_path / "progress.jsonl").touch() + workspace = PresetAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home") + session_path = tmp_path / "session" + session_path.mkdir() + (session_path / "agent.log").touch() + agent_session = PresetAgentSession( + path=session_path, + debug=False, + ) + return workspace, agent_session + + +def _patch_claude_command(monkeypatch, script): + monkeypatch.setattr( + "dstack._internal.cli.services.presets.agent._build_claude_command", + lambda **kwargs: [sys.executable, str(script)] + + (["--resume", kwargs["resume_session_id"]] if kwargs.get("resume_session_id") else []), + ) + + +class TestConnectionResume: + @pytest.mark.asyncio + async def test_resumes_after_connection_error(self, tmp_path, monkeypatch, capsys): + script = _write_fake_claude( + tmp_path, + """import json +import sys +from pathlib import Path + +args = sys.argv[1:] +with Path("calls.jsonl").open("a") as f: + f.write(json.dumps({"args": args, "prompt": sys.stdin.read()}) + "\\n") +if "--resume" in args: + print(json.dumps({ + "type": "result", + "session_id": "sid-123", + "structured_output": {"resumed": True}, + })) +else: + print(json.dumps({"type": "system", "subtype": "init", "session_id": "sid-123"})) + print(json.dumps({ + "type": "result", + "is_error": True, + "result": "API Error: Unable to connect to API (ENOTFOUND)", + })) + sys.exit(1) +""", + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.agent._RESUME_DELAYS_SECONDS", (0,) + ) + _patch_claude_command(monkeypatch, script) + workspace, agent_session = _agent_setup(tmp_path) + + output = await run_preset_agent( + prompt="system prompt", + env=_subprocess_env(), + workspace=workspace, + auth=ClaudeAuth(api_key=None, executable="claude", effort=None, model="m"), + redacted_values=(), + agent_session=agent_session, + ) + + assert output.report_data == {"resumed": True} + calls = [json.loads(line) for line in (tmp_path / "calls.jsonl").read_text().splitlines()] + assert len(calls) == 2 + assert calls[0]["args"] == [] + assert calls[0]["prompt"] == "system prompt" + assert calls[1]["args"] == ["--resume", "sid-123"] + assert "The previous agent process was interrupted" in calls[1]["prompt"] + assert "resuming" in capsys.readouterr().out + + @pytest.mark.asyncio + async def test_resumes_on_any_unreported_death(self, tmp_path, monkeypatch): + script = _write_fake_claude( + tmp_path, + """import json +import sys +from pathlib import Path + +with Path("calls.jsonl").open("a") as f: + f.write("call\\n") +if "--resume" in sys.argv[1:]: + print(json.dumps({"type": "result", "structured_output": {"recovered": True}})) +else: + print(json.dumps({"type": "system", "subtype": "init", "session_id": "sid-123"})) + print(json.dumps({"type": "result", "is_error": True, "result": "model refused"})) + sys.exit(1) +""", + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.agent._RESUME_DELAYS_SECONDS", (0,) + ) + _patch_claude_command(monkeypatch, script) + workspace, agent_session = _agent_setup(tmp_path) + + output = await run_preset_agent( + prompt="system prompt", + env=_subprocess_env(), + workspace=workspace, + auth=ClaudeAuth(api_key=None, executable="claude", effort=None, model="m"), + redacted_values=(), + agent_session=agent_session, + ) + + assert output.report_data == {"recovered": True} + assert len((tmp_path / "calls.jsonl").read_text().splitlines()) == 2 + + +class TestConnectionRetryExhaustion: + @pytest.mark.asyncio + async def test_gives_up_after_repeated_no_progress_failures(self, tmp_path, monkeypatch): + script = _write_fake_claude( + tmp_path, + """import json +import sys +from pathlib import Path + +with Path("calls.jsonl").open("a") as f: + f.write("call\\n") +print(json.dumps({"type": "system", "subtype": "init", "session_id": "sid-123"})) +print(json.dumps({ + "type": "result", + "is_error": True, + "result": "API Error: Unable to connect to API (ENOTFOUND)", +})) +sys.exit(1) +""", + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.agent._RESUME_DELAYS_SECONDS", (0, 0) + ) + _patch_claude_command(monkeypatch, script) + workspace, agent_session = _agent_setup(tmp_path) + + output = await run_preset_agent( + prompt="system prompt", + env=_subprocess_env(), + workspace=workspace, + auth=ClaudeAuth(api_key=None, executable="claude", effort=None, model="m"), + redacted_values=(), + agent_session=agent_session, + ) + + assert output.report_data is None + assert output.error is not None and "Unable to connect" in output.error + # Initial attempt + one retry per configured delay, then give up. + assert len((tmp_path / "calls.jsonl").read_text().splitlines()) == 3 + + @pytest.mark.asyncio + async def test_does_not_resurrect_an_externally_stopped_agent(self, tmp_path, monkeypatch): + script = _write_fake_claude( + tmp_path, + """import json +import sys +from pathlib import Path + +with Path("calls.jsonl").open("a") as f: + f.write("call\\n") +print(json.dumps({"type": "system", "subtype": "init", "session_id": "sid-123"})) +print(json.dumps({ + "type": "result", + "is_error": True, + "result": "API Error: Unable to connect to API (ENOTFOUND)", +})) +sys.exit(1) +""", + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.agent._RESUME_DELAYS_SECONDS", (0, 0) + ) + _patch_claude_command(monkeypatch, script) + workspace, agent_session = _agent_setup(tmp_path) + # Another CLI recorded a stop: the death must read as a decision, not an + # outage to retry. + agent_session.update_manifest(status="interrupted") + + output = await run_preset_agent( + prompt="system prompt", + env=_subprocess_env(), + workspace=workspace, + auth=ClaudeAuth(api_key=None, executable="claude", effort=None, model="m"), + redacted_values=(), + agent_session=agent_session, + ) + + assert output.report_data is None + assert len((tmp_path / "calls.jsonl").read_text().splitlines()) == 1 + + +class TestWorkspaceLifecycle: + def _session(self, tmp_path): + session_dir = tmp_path / "sessions" / "ab12cd34" + session_dir.mkdir(parents=True) + return PresetAgentSession(path=session_dir, debug=False, preset_id="ab12cd34") + + @pytest.mark.skipif(IS_WINDOWS, reason="workspace alias symlinks are POSIX-only") + def test_create_attach_and_remove(self, tmp_path): + session = self._session(tmp_path) + + workspace = create_agent_workspace(session) + manifest = session.read_manifest() + alias = Path(manifest["alias"]) + assert Path(manifest["workspace"]) == session.path / "workspace" + assert alias.is_symlink() + (workspace.path / "note.txt").write_text("kept", encoding="utf-8") + + os.unlink(alias) + attached = attach_agent_workspace(session) + assert alias.is_symlink() + assert (attached.path / "note.txt").read_text() == "kept" + + remove_agent_workspace(session) + assert not os.path.lexists(alias) + assert not (session.path / "workspace").exists() + + @pytest.mark.skipif(IS_WINDOWS, reason="workspace alias symlinks are POSIX-only") + def test_attach_refuses_occupied_alias(self, tmp_path): + session = self._session(tmp_path) + create_agent_workspace(session) + manifest = session.read_manifest() + alias = Path(manifest["alias"]) + os.unlink(alias) + alias.mkdir() + try: + with pytest.raises(CLIError, match="cannot be used safely"): + attach_agent_workspace(session) + finally: + alias.rmdir() + + def test_attach_fails_when_workspace_is_gone(self, tmp_path): + session = self._session(tmp_path) + create_agent_workspace(session) + remove_agent_workspace(session) + with pytest.raises(CLIError, match="no longer exists"): + attach_agent_workspace(session) + + +class TestOffsetPersistence: + def test_mirror_does_not_duplicate_after_restart(self, tmp_path): + from dstack._internal.cli.services.presets.tail import _OffsetStore + + source = tmp_path / "runs.jsonl" + target = tmp_path / "mirror.jsonl" + state = tmp_path / ".offsets.json" + source.write_text('{"name":"one"}\n', encoding="utf-8") + + mirror = _RecordMirror( + source=source, + target=target, + redacted_values=(), + offset_store=_OffsetStore(state), + offset_key="runs", + ) + mirror.flush() + + with source.open("a", encoding="utf-8") as f: + f.write('{"name":"two"}\n') + restarted = _RecordMirror( + source=source, + target=target, + redacted_values=(), + offset_store=_OffsetStore(state), + offset_key="runs", + ) + restarted.flush() + + assert target.read_text().splitlines() == ['{"name":"one"}', '{"name":"two"}'] + + +class TestLoadResumableSession: + def _write_session(self, tmp_path, monkeypatch, manifest): + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + path = tmp_path / ".dstack" / "presets" / manifest["id"] + path.mkdir(parents=True) + (path / "session.json").write_text(json.dumps(manifest), encoding="utf-8") + return path + + def test_loads_interrupted_session(self, tmp_path, monkeypatch): + self._write_session( + tmp_path, + monkeypatch, + { + "id": "ab12cd34", + "status": "interrupted", + "claude_session_id": "sid-1", + "debug": True, + "created_at": "2026-07-20T10:00:00+00:00", + }, + ) + + session = load_resumable_agent_session("ab12cd34") + + assert session.preset_id == "ab12cd34" + assert session.debug is True + + def test_treats_dead_running_session_as_resumable(self, tmp_path, monkeypatch): + self._write_session( + tmp_path, + monkeypatch, + { + "id": "ab12cd34", + "status": "running", + "pid": 4242, + "claude_session_id": "sid-1", + }, + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.session.psutil.pid_exists", + lambda pid: False, + ) + + assert load_resumable_agent_session("ab12cd34").preset_id == "ab12cd34" + + @pytest.mark.parametrize( + ("manifest", "match"), + [ + pytest.param(None, "Unknown preset", id="unknown-preset"), + pytest.param( + {"id": "aa000001", "status": "success"}, "nothing to resume", id="success" + ), + pytest.param({"id": "aa000002", "status": "failed"}, "cannot be resumed", id="failed"), + pytest.param( + { + "id": "aa000003", + "status": "running", + "pid": 4242, + "claude_session_id": "sid-1", + }, + "still being created", + id="still-running", + ), + pytest.param( + {"id": "aa000004", "status": "interrupted"}, + "stopped before it started", + id="interrupted-before-start", + ), + ], + ) + def test_refusals(self, tmp_path, monkeypatch, manifest, match): + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + if manifest is not None: + self._write_session(tmp_path, monkeypatch, manifest) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.session.psutil.pid_exists", + lambda pid: True, + ) + preset_id = manifest["id"] if manifest is not None else "00000000" + + with pytest.raises(CLIError, match=match): + load_resumable_agent_session(preset_id) + + +class TestSummarizeSessionTrials: + def test_counts_records_even_when_trials_share_a_task(self, tmp_path): + record = { + "task": {"name": "qwen-ab12cd34-1"}, + "resources": {"gpu": {"name": "A40", "memory": "48GB", "count": 1}}, + "benchmark": { + "workload": {"concurrency": 8}, + "metrics": {"duration_seconds": 10.0, "total_output_tokens": 20000}, + }, + } + rerun = json.loads(json.dumps(record)) + rerun["benchmark"]["metrics"]["total_output_tokens"] = 23000 + second_trial = json.loads(json.dumps(record)) + second_trial["task"]["name"] = "qwen-ab12cd34-2" + nameless = {"benchmark": {"workload": {}, "metrics": {}}} + path = tmp_path / "trials.jsonl" + path.write_text( + "\n".join(json.dumps(entry) for entry in [record, rerun, second_trial, nameless]) + ) + + summary = _summarize_session_trials(path) + + # 4 records = 4 trials: one long-lived task commonly hosts several + # trials, so shared task names must not collapse the count. + assert summary["count"] == 4 + assert summary["best"] == {"tok_s": 2300.0, "concurrency": 8, "gpu": "A40:48GB:1"} + + +class TestFileLineReader: + @pytest.mark.asyncio + async def test_reads_lines_and_continues_from_persisted_offset(self, tmp_path): + from dstack._internal.cli.services.presets.tail import _FileLineReader, _OffsetStore + + stream = tmp_path / "stdout.jsonl" + state = tmp_path / ".offsets.json" + stream.write_bytes(b"first\nsecond\n") + alive = True + + reader = _FileLineReader( + stream, + offset_store=_OffsetStore(state), + offset_key="agent_stdout", + is_alive=lambda: alive, + ) + assert await reader.readline() == b"first\n" + assert await reader.readline() == b"second\n" + + # A new reader (a later attach) continues where the previous stopped. + with stream.open("ab") as f: + f.write(b"third\npartial") + alive = False + attached = _FileLineReader( + stream, + offset_store=_OffsetStore(state), + offset_key="agent_stdout", + is_alive=lambda: alive, + ) + assert await attached.readline() == b"third\n" + assert await attached.readline() == b"partial" + assert await attached.readline() == b"" + + +class TestStopOrDetach: + @pytest.mark.skipif(IS_WINDOWS, reason="exercises POSIX process groups") + def test_detach_keeps_the_agent_and_stop_terminates_it(self, tmp_path, monkeypatch, capsys): + from dstack._internal.cli.services.presets import create as create_module + from dstack._internal.cli.services.presets.create import _stop_or_detach_agent_session + + session_dir = tmp_path / "ab12cd34" + session_dir.mkdir() + (session_dir / "agent.log").touch() + session = PresetAgentSession(path=session_dir, debug=False, preset_id="ab12cd34") + agent = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(300)"], start_new_session=True + ) + try: + session.update_manifest(status="running", agent_pid=agent.pid) + + monkeypatch.setattr(create_module, "confirm_ask", lambda *_: False) + _stop_or_detach_agent_session(session) + assert agent.poll() is None + assert session.read_manifest()["status"] == "running" + assert "Detached" in capsys.readouterr().out + + monkeypatch.setattr(create_module, "confirm_ask", lambda *_: True) + _stop_or_detach_agent_session(session) + psutil.Process(agent.pid).wait(timeout=10) + assert session.read_manifest()["status"] == "interrupted" + finally: + with suppress(OSError): + os.killpg(agent.pid, signal.SIGKILL) + + +class TestOffsetStoreSharing: + def test_shared_store_keeps_every_writers_keys(self, tmp_path): + import threading + + from dstack._internal.cli.services.presets.tail import _OffsetStore + + state = tmp_path / ".offsets.json" + store = _OffsetStore(state) + + # One store serves the whole session; readers and mirrors write + # disjoint keys from worker threads. + def advance(key: str) -> None: + for offset in range(1, 51): + store.set(key, offset) + + threads = [ + threading.Thread(target=advance, args=(key,)) + for key in ("agent_stdout", "agent_stderr", "runs", "trials") + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + reloaded = _OffsetStore(state) + for key in ("agent_stdout", "agent_stderr", "runs", "trials"): + assert reloaded.get(key) == 50 diff --git a/src/tests/_internal/cli/services/presets/test_apply.py b/src/tests/_internal/cli/services/presets/test_apply.py new file mode 100644 index 000000000..3eae3cf69 --- /dev/null +++ b/src/tests/_internal/cli/services/presets/test_apply.py @@ -0,0 +1,135 @@ +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.services.presets.apply import ( + _build_service, + _validate_preset_matches, + apply_preset, +) +from dstack._internal.core.errors import CLIError +from dstack._internal.core.models.instances import InstanceAvailability +from tests._internal.cli.preset_factories import get_preset + +pytestmark = pytest.mark.windows + + +class TestValidatePresetMatches: + def test_accepts_matching_base_model_and_context(self): + preset = get_preset(preset_id="large", context_length=32768) + configuration = PresetConfiguration( + name="qwen", + model={"base": "Qwen/Qwen3.5-27B"}, + context_length=8192, + ) + + _validate_preset_matches(preset, configuration=configuration) + + def test_rejects_insufficient_context(self): + preset = get_preset(preset_id="small", context_length=4096) + configuration = PresetConfiguration( + name="qwen", + model={"base": "Qwen/Qwen3.5-27B"}, + context_length=8192, + ) + + with pytest.raises(CLIError, match="context length"): + _validate_preset_matches(preset, configuration=configuration) + + def test_exact_request_matches_repo_and_client_facing_name(self): + matching = get_preset(preset_id="matching") + configuration = PresetConfiguration( + name="qwen", + model={ + "repo": "community/Qwen3.5-27B-GPTQ-Int4", + "name": "Qwen/Qwen3.5-27B", + }, + ) + + _validate_preset_matches(matching, configuration=configuration) + with pytest.raises(CLIError, match="does not serve repo"): + _validate_preset_matches( + matching.copy(update={"model": "other/repo"}), + configuration=configuration, + ) + + +class TestBuildService: + def test_applies_preset_name_env_gateway_and_constraints(self): + configuration = PresetConfiguration( + name="qwen-production", + model={"base": "Qwen/Qwen3.5-27B"}, + gateway="inference", + env={"HF_TOKEN": "token"}, + fleets=["gpu-fleet"], + max_price=1, + ) + + service = _build_service(configuration, get_preset()) + + assert service.name == "qwen-production" + assert service.gateway == "inference" + assert service.env["HF_TOKEN"] == "token" + assert [fleet.format() for fleet in service.fleets] == ["gpu-fleet"] + assert service.max_price == 1 + + +class TestApplyPreset: + def test_rejects_unknown_preset(self): + with pytest.raises(CLIError, match="does not exist"): + apply_preset( + api=Mock(), + configuration=PresetConfiguration(name="qwen", model={"base": "Qwen/Qwen3.5-27B"}), + configuration_path="preset.dstack.yml", + preset_id="ee55ff66", + profile_name=None, + command_args=SimpleNamespace(), + store=Mock(get=Mock(return_value=None)), + ) + + def test_applies_the_referenced_preset(self, monkeypatch): + preset = get_preset() + run_plan = _plan(InstanceAvailability.AVAILABLE) + repo = Mock() + service_args = SimpleNamespace(profile="gpu") + configurator = Mock() + configurator.get_parser.return_value.parse_args.return_value = service_args + configurator.get_plan.return_value = (run_plan, repo) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.apply.ServiceConfigurator", + lambda api_client: configurator, + ) + command_args = SimpleNamespace() + + apply_preset( + api=Mock(), + configuration=PresetConfiguration( + name="qwen", + model={"base": "Qwen/Qwen3.5-27B"}, + ), + configuration_path="preset.dstack.yml", + preset_id="8f3a12c4", + profile_name="gpu", + command_args=command_args, + store=Mock(get=Mock(return_value=preset)), + ) + + assert service_args.profile == "gpu" + configurator.apply_plan.assert_called_once_with( + run_plan=run_plan, + repo=repo, + command_args=command_args, + configurator_args=service_args, + plan_properties={ + "Model": "Qwen/Qwen3.5-27B ([secondary]base[/])", + "Preset": "8f3a12c4 ([secondary]ctx=32K con=1 42.1 tok/s TTFT 108ms[/])", + }, + ) + + +def _plan(availability: InstanceAvailability): + return SimpleNamespace( + job_plans=[SimpleNamespace(offers=[SimpleNamespace(availability=availability)])] + ) diff --git a/src/tests/_internal/cli/services/presets/test_create.py b/src/tests/_internal/cli/services/presets/test_create.py new file mode 100644 index 000000000..e87e589f5 --- /dev/null +++ b/src/tests/_internal/cli/services/presets/test_create.py @@ -0,0 +1,1126 @@ +import asyncio +import json +import os +import uuid +from types import SimpleNamespace + +import pytest + +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.services.presets.agent import ( + ClaudeAuth, + PresetAgentProcessOutput, +) +from dstack._internal.cli.services.presets.create import ( + PresetCreateResult, + SessionBusyError, + _build_constraints, + _cleanup_runs, + _create_preset, + _get_build_name, + _print_fleet_offers, + _save_final_report_copy, + _stop_active_session_runs, + _suspend_agent_session, + create_preset, + follow_preset, + reconcile_detached_sessions, + stop_preset_session, +) +from dstack._internal.cli.services.presets.session import ( + PresetAgentSession, + load_agent_session, + mark_session_owner, + print_preset_progress, + print_session_log, + release_session_claim, + session_process_alive, + try_claim_session, +) +from dstack._internal.cli.services.presets.store import PresetStore +from dstack._internal.cli.services.presets.workspace import ( + PresetAgentWorkspace, + create_agent_workspace, + remove_agent_workspace, +) +from dstack._internal.core.errors import CLIError +from dstack._internal.core.models.envs import EnvSentinel +from dstack._internal.core.models.runs import Run, RunStatus +from tests._internal.cli.preset_factories import ( + get_preset, + get_running_service_run, + get_successful_preset_report, +) + +pytestmark = pytest.mark.windows + + +def _claude_auth() -> ClaudeAuth: + return ClaudeAuth( + api_key="anthropic-secret", + executable="claude", + effort=None, + model="claude-test", + ) + + +def _session_dirs(tmp_path): + return [ + path + for path in (tmp_path / ".dstack" / "presets").iterdir() + if path.is_dir() and not path.name.startswith(".") + ] + + +@pytest.fixture +def creation_context(tmp_path, monkeypatch): + run = get_running_service_run() + run_apis = _FakeRunAPIs(run) + api = SimpleNamespace( + project="main", + runs=run_apis, + client=SimpleNamespace( + _token="dstack-secret", + base_url="http://127.0.0.1:3000", + runs=run_apis, + ), + ) + configuration = PresetConfiguration( + name="qwen-build", + model={"base": "Qwen/Qwen3.5-27B"}, + context_length=8192, + max_trials=1, + fleets=["gpu-fleet"], + env={"LICENSE": "license-secret", "TOKENIZERS_PARALLELISM": "false"}, + ) + source_configuration = PresetConfiguration( + name="qwen-build", + model={"base": "Qwen/Qwen3.5-27B"}, + context_length=8192, + max_trials=1, + fleets=["gpu-fleet"], + env=["LICENSE", "TOKENIZERS_PARALLELISM=false"], + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.get_claude_auth", + _claude_auth, + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create._get_build_name", + lambda *_: "qwen-build", + ) + return SimpleNamespace( + api=api, + configuration=configuration, + source_configuration=source_configuration, + run=run, + run_apis=run_apis, + store=PresetStore(tmp_path / "presets"), + ) + + +class TestCreatePreset: + def test_saves_agent_log_without_debug(self, tmp_path, monkeypatch): + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + preset = get_preset() + + async def create(**kwargs): + print_preset_progress("testing preset", agent_session=kwargs["agent_session"]) + return PresetCreateResult( + preset=preset, + path=tmp_path / "preset.yaml", + final_run_id=uuid.uuid4(), + final_run_name="qwen-build-2", + ) + + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create._create_preset", + create, + ) + + create_preset( + api=SimpleNamespace(), + configuration=PresetConfiguration( + name="qwen", + model={"base": "Qwen/Qwen3.5-27B"}, + ), + store=PresetStore(tmp_path / "presets"), + ) + + paths = _session_dirs(tmp_path) + assert len(paths) == 1 + assert {path.name for path in paths[0].iterdir()} == { + "agent.log", + "session.json", + "preset.dstack.yml", + } + manifest = json.loads((paths[0] / "session.json").read_text()) + assert manifest["status"] == "success" + assert manifest["id"] == paths[0].name + assert "testing preset" in (paths[0] / "agent.log").read_text() + + def test_debug_finalization_error_does_not_mask_success(self, tmp_path, monkeypatch, capsys): + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + monkeypatch.setenv("HF_TOKEN", "hf-secret") + preset = get_preset() + + async def create(**kwargs): + assert kwargs["configuration"].env.as_dict() == { + "HF_TOKEN": "hf-secret", + "TOKENIZERS_PARALLELISM": "false", + } + assert isinstance(kwargs["source_configuration"].env["HF_TOKEN"], EnvSentinel) + assert kwargs["source_configuration"].env["TOKENIZERS_PARALLELISM"] == "false" + kwargs["agent_session"].write_prompt("test prompt") + return PresetCreateResult( + preset=preset, + path=tmp_path / "preset.yaml", + final_run_id=uuid.uuid4(), + final_run_name="qwen-build-2", + ) + + def fail_finish(self, preset_id=None): + raise OSError("rename failed") + + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create._create_preset", + create, + ) + monkeypatch.setattr(PresetAgentSession, "finish", fail_finish) + + result = create_preset( + api=SimpleNamespace(), + configuration=PresetConfiguration( + name="qwen", + model={"base": "Qwen/Qwen3.5-27B"}, + env=["HF_TOKEN", "TOKENIZERS_PARALLELISM=false"], + ), + store=PresetStore(tmp_path / "presets"), + debug=True, + ) + + paths = _session_dirs(tmp_path) + assert len(paths) == 1 + assert result.preset == preset + assert {path.name for path in paths[0].iterdir()} == { + "preset.dstack.yml", + "agent.log", + "prompt.md", + "session.json", + "trace.jsonl", + } + assert json.loads((paths[0] / "session.json").read_text())["status"] == "running" + assert "hf-secret" not in (paths[0] / "preset.dstack.yml").read_text() + assert "Files remain at" in capsys.readouterr().out + + def test_debug_finalization_does_not_mask_creation_error(self, tmp_path, monkeypatch): + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + + async def create(**kwargs): + raise RuntimeError("creation failed") + + def fail_finish(self, preset_id=None): + raise OSError("rename failed") + + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create._create_preset", + create, + ) + monkeypatch.setattr(PresetAgentSession, "finish", fail_finish) + + with pytest.raises(RuntimeError, match="creation failed"): + create_preset( + api=SimpleNamespace(), + configuration=PresetConfiguration( + name="qwen", + model={"base": "Qwen/Qwen3.5-27B"}, + ), + store=PresetStore(tmp_path / "presets"), + debug=True, + ) + + @pytest.mark.asyncio + async def test_checks_active_fleets_before_claude_auth(self, tmp_path, monkeypatch): + api = SimpleNamespace( + project="main", + client=SimpleNamespace(fleets=SimpleNamespace(list=lambda *args, **kwargs: [])), + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.get_claude_auth", + lambda: pytest.fail("Claude auth must not be checked without an active fleet"), + ) + + with pytest.raises(CLIError, match="no fleets"): + await _create_preset( + api=api, + configuration=PresetConfiguration( + name="qwen-build", + model={"base": "Qwen/Qwen3.5-27B"}, + ), + store=PresetStore(tmp_path / "presets"), + agent_session=_agent_session(tmp_path), + ) + + @pytest.mark.asyncio + async def test_skips_cleanup_when_cancelled(self, creation_context, monkeypatch, tmp_path): + async def run_agent(**_): + raise asyncio.CancelledError + + cleanup_calls = [] + + async def cleanup_runs(**kwargs): + cleanup_calls.append(kwargs) + + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.run_preset_agent", + run_agent, + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create._cleanup_runs", + cleanup_runs, + ) + + with pytest.raises(asyncio.CancelledError): + await _create_preset( + api=creation_context.api, + configuration=creation_context.configuration, + source_configuration=creation_context.source_configuration, + store=creation_context.store, + agent_session=_agent_session(tmp_path), + ) + + assert cleanup_calls == [] + + @pytest.mark.parametrize( + ("keep_service", "stopped_names"), + [(False, ["qwen-build-2"]), (True, [])], + ) + @pytest.mark.asyncio + async def test_saves_preset_and_cleans_up_runs( + self, creation_context, monkeypatch, keep_service, stopped_names, tmp_path + ): + session_path = tmp_path / "debug-running" + session_path.mkdir() + (session_path / "agent.log").touch() + (session_path / "trace.jsonl").touch() + agent_session = PresetAgentSession( + path=session_path, + debug=True, + ) + + async def run_agent(**kwargs): + assert kwargs["agent_session"] is agent_session + assert (session_path / "prompt.md").is_file() + return PresetAgentProcessOutput( + report_data=json.loads(get_successful_preset_report(creation_context.run).json()) + ) + + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.run_preset_agent", + run_agent, + ) + result = await _create_preset( + api=creation_context.api, + configuration=creation_context.configuration, + source_configuration=creation_context.source_configuration, + store=creation_context.store, + keep_service=keep_service, + build_name="qwen-build", + agent_session=agent_session, + ) + + assert result.preset.base == "Qwen/Qwen3.5-27B" + assert result.path.is_file() + assert creation_context.store.list() == [result.preset] + assert "license-secret" not in result.path.read_text() + assert result.preset.service.env["TOKENIZERS_PARALLELISM"] == "false" + assert creation_context.run_apis.stopped_names == stopped_names + + +class TestBuildName: + def test_derives_slug_for_nameless_and_keeps_prefix_bounded(self): + assert _get_build_name(None, "Qwen/Qwen3.5-27B", "a1b2c3d4") == "qwen3-5-27b-a1b2c3d4" + + build_name = _get_build_name( + "qwen-preset-with-a-name-that-is-forty-one", "Qwen/Qwen3.5-27B", "a1b2c3d4" + ) + + assert build_name.endswith("-a1b2c3d4") + assert len(f"{build_name}-99999") <= 41 + + +class TestCleanupRuns: + @pytest.mark.asyncio + async def test_stops_only_recorded_build_runs(self, tmp_path, monkeypatch): + (tmp_path / "runs.jsonl").write_text( + '{"name":"qwen-build-1"}\n{"name":"unrelated-run"}\n{"name":"qwen-build-2"}\n' + ) + runs = _FakeRuns() + api = SimpleNamespace( + runs=runs, + project="main", + client=SimpleNamespace(runs=runs), + ) + + async def no_sleep(_): + return None + + monkeypatch.setattr(asyncio, "sleep", no_sleep) + await _cleanup_runs( + api=api, + build_name="qwen-build", + workspace=PresetAgentWorkspace( + path=tmp_path, + dstack_home=tmp_path / "home", + ), + final_run_name="qwen-build-2", + keep_final_service=True, + agent_session=_agent_session(tmp_path), + ) + + assert runs.stopped_names == ["qwen-build-1"] + + +def _agent_session(tmp_path, *, debug: bool = False) -> PresetAgentSession: + path = tmp_path / "agent-running" + path.mkdir() + (path / "agent.log").touch() + if debug: + (path / "trace.jsonl").touch() + return PresetAgentSession( + path=path, + debug=debug, + preset_id="ab12cd34", + ) + + +class _FakeRuns: + def __init__(self): + self.stopped_names: list[str] = [] + + def get(self, name): + status = RunStatus.TERMINATED if name in self.stopped_names else RunStatus.RUNNING + return SimpleNamespace(status=status) + + def stop(self, project, names, abort): + assert project == "main" + assert abort is False + self.stopped_names.extend(names) + + +class _FakeRunAPIs: + def __init__(self, run: Run): + self.run = run + self.stopped_names: list[str] = [] + + def get(self, *args): + name = args[-1] + return self.run if name == self.run.run_spec.run_name else None + + def stop(self, project, names, abort): + assert project == "main" + assert abort is False + self.stopped_names.extend(names) + self.run.status = RunStatus.TERMINATED + + +class TestBuildConstraints: + def test_renders_all_fields_with_explicit_nulls_and_defaults(self): + configuration = PresetConfiguration( + name="qwen", + model={"base": "Qwen/Qwen3-32B"}, + max_trials=3, + env=["HF_TOKEN"], + ) + + text = _build_constraints( + configuration=configuration, + build_name="qwen-abc123", + allowed_fleets=("gpu-fleet",), + ) + + assert text.endswith("\n") + assert json.loads(text) == { + "run_name_prefix": "qwen-abc123", + "model": {"base": "Qwen/Qwen3-32B"}, + "context_length": None, + "max_trials": 3, + "concurrency": 8, + "fleets": ["gpu-fleet"], + "env": ["HF_TOKEN"], + } + + def test_renders_configured_values(self): + configuration = PresetConfiguration( + name="qwen", + model={"repo": "Qwen/Qwen3-32B-AWQ", "name": "qwen3"}, + context_length=32768, + max_trials=10, + concurrency=16, + ) + + data = json.loads( + _build_constraints( + configuration=configuration, + build_name="qwen-abc123", + allowed_fleets=("a", "b"), + ) + ) + + assert data["model"] == {"repo": "Qwen/Qwen3-32B-AWQ", "name": "qwen3"} + assert data["context_length"] == 32768 + assert data["max_trials"] == 10 + assert data["concurrency"] == 16 + assert data["fleets"] == ["a", "b"] + + +class TestSaveFinalReportCopy: + def test_copies_report_redacted(self, tmp_path): + workspace = PresetAgentWorkspace(path=tmp_path / "w", dstack_home=tmp_path / "h") + workspace.path.mkdir() + workspace.final_report_path.write_text( + '{"success": true, "note": "token dstack-secret"}', encoding="utf-8" + ) + session = _agent_session(tmp_path, debug=True) + + _save_final_report_copy( + workspace=workspace, + agent_session=session, + redacted_values=["dstack-secret"], + ) + + copied = (session.path / "final_report.json").read_text() + assert "dstack-secret" not in copied + assert "[redacted]" in copied + + def test_missing_report_is_no_op(self, tmp_path): + workspace = PresetAgentWorkspace(path=tmp_path / "w", dstack_home=tmp_path / "h") + workspace.path.mkdir() + session = _agent_session(tmp_path, debug=True) + + _save_final_report_copy( + workspace=workspace, + agent_session=session, + redacted_values=["dstack-secret"], + ) + + assert not (session.path / "final_report.json").exists() + + +class TestInterruptAndResume: + def test_interrupt_suspends_session(self, tmp_path, monkeypatch, capsys): + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + + async def create(**kwargs): + raise KeyboardInterrupt + + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create._create_preset", + create, + ) + + with pytest.raises(KeyboardInterrupt): + create_preset( + api=SimpleNamespace(), + configuration=PresetConfiguration(name="qwen", model={"base": "Qwen/Qwen3.5-27B"}), + store=PresetStore(tmp_path / "presets"), + ) + + sessions = _session_dirs(tmp_path) + assert len(sessions) == 1 + manifest = json.loads((sessions[0] / "session.json").read_text()) + assert manifest["status"] == "interrupted" + output = capsys.readouterr().out + assert "--resume" in output + assert sessions[0].name in output + + def test_suspend_scrubs_workspace_token(self, tmp_path, capsys): + session_dir = tmp_path / "ab12cd34" + session_dir.mkdir() + (session_dir / "agent.log").touch() + session = PresetAgentSession(path=session_dir, debug=False, preset_id="ab12cd34") + workspace_root = tmp_path / "workspace" + config_dir = workspace_root / "h" / ".dstack" + config_dir.mkdir(parents=True) + (config_dir / "config.yml").write_text("projects: []\n") + (workspace_root / "w").mkdir() + (workspace_root / "w" / "constraints.json").write_text("{}") + session.update_manifest(workspace=str(workspace_root)) + + _suspend_agent_session(session) + + # The live credential is gone; the rest of the workspace stays resumable. + assert not (config_dir / "config.yml").exists() + assert (workspace_root / "w" / "constraints.json").exists() + assert session.read_manifest()["status"] == "interrupted" + + @pytest.mark.asyncio + async def test_resume_uses_saved_claude_session(self, creation_context, monkeypatch, tmp_path): + session_dir = tmp_path / "sessions" / "fe98dc76" + session_dir.mkdir(parents=True) + (session_dir / "agent.log").touch() + agent_session = PresetAgentSession(path=session_dir, debug=False, preset_id="fe98dc76") + workspace = create_agent_workspace(agent_session) + workspace.constraints_path.write_text( + '{"run_name_prefix": "qwen-build"}', encoding="utf-8" + ) + agent_session.update_manifest(claude_session_id="sid-xyz", claude_model="claude-pinned") + captured = {} + + async def run_agent(**kwargs): + captured.update(kwargs) + return PresetAgentProcessOutput( + report_data=json.loads(get_successful_preset_report(creation_context.run).json()) + ) + + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.run_preset_agent", + run_agent, + ) + + result = await _create_preset( + api=creation_context.api, + configuration=creation_context.configuration, + source_configuration=creation_context.source_configuration, + store=creation_context.store, + agent_session=agent_session, + resume=True, + ) + + assert captured["initial_resume_session_id"] == "sid-xyz" + assert captured["auth"].model == "claude-pinned" + assert result.preset.id == "fe98dc76" + assert (session_dir / "workspace").is_dir() + remove_agent_workspace(agent_session) + + @pytest.mark.asyncio + async def test_pins_user_prompt_on_create(self, creation_context, monkeypatch, tmp_path): + session_dir = tmp_path / "ab34ef12" + session_dir.mkdir() + (session_dir / "agent.log").touch() + agent_session = PresetAgentSession(path=session_dir, debug=False, preset_id="ab34ef12") + captured = {} + + async def run_agent(**kwargs): + captured.update(kwargs) + return PresetAgentProcessOutput( + report_data=json.loads(get_successful_preset_report(creation_context.run).json()) + ) + + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.run_preset_agent", + run_agent, + ) + + await _create_preset( + api=creation_context.api, + configuration=creation_context.configuration, + source_configuration=creation_context.source_configuration, + store=creation_context.store, + build_name="qwen-build", + agent_session=agent_session, + user_prompt="Optimize for RAG traffic.", + ) + + assert agent_session.read_user_prompt() == "Optimize for RAG traffic." + assert "## Additional instructions" in captured["prompt"] + assert "Optimize for RAG traffic." in captured["prompt"] + + @pytest.mark.asyncio + async def test_resume_keeps_the_pinned_user_prompt( + self, creation_context, monkeypatch, tmp_path, capsys + ): + session_dir = tmp_path / "ab34ef12" + session_dir.mkdir() + (session_dir / "agent.log").touch() + agent_session = PresetAgentSession(path=session_dir, debug=False, preset_id="ab34ef12") + workspace = create_agent_workspace(agent_session) + workspace.constraints_path.write_text( + '{"run_name_prefix": "qwen-build"}', encoding="utf-8" + ) + agent_session.update_manifest(claude_session_id="sid-abc") + agent_session.write_user_prompt("Optimize for RAG traffic.") + captured = {} + + async def run_agent(**kwargs): + captured.update(kwargs) + return PresetAgentProcessOutput( + report_data=json.loads(get_successful_preset_report(creation_context.run).json()) + ) + + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.run_preset_agent", + run_agent, + ) + + await _create_preset( + api=creation_context.api, + configuration=creation_context.configuration, + source_configuration=creation_context.source_configuration, + store=creation_context.store, + agent_session=agent_session, + resume=True, + user_prompt="A different prompt.", + ) + + # The session keeps its original prompt; the new one is ignored with a warning. + assert "Optimize for RAG traffic." in captured["prompt"] + assert "A different prompt." not in captured["prompt"] + assert "keepsitsoriginalprompt" in "".join(capsys.readouterr().out.split()) + remove_agent_workspace(agent_session) + + +class TestFleetOffersPreview: + def test_no_offers_shows_the_shared_warning_without_failing(self, capsys): + plan = SimpleNamespace( + project_name="main", + user="admin", + job_plans=[SimpleNamespace(offers=[], total_offers=0, max_price=None)], + ) + api = SimpleNamespace( + project="main", + client=SimpleNamespace(runs=SimpleNamespace(get_plan=lambda *a, **k: plan)), + ) + + _print_fleet_offers(api, ("arm-fleet",)) + + out = capsys.readouterr().out + assert "arm-fleet" in out + assert "No matching instance offers available" in out + + +class TestSessionLog: + def _session(self, tmp_path, preset_id: str, status: str, log: str) -> PresetAgentSession: + session_dir = tmp_path / preset_id + session_dir.mkdir() + (session_dir / "agent.log").write_text(log) + session = PresetAgentSession(path=session_dir, debug=False, preset_id=preset_id) + session.update_manifest(status=status) + return session + + def test_load_agent_session_reads_any_status(self, tmp_path, monkeypatch): + self._session(tmp_path, "dead0000", "failed", "[t] boom\n") + monkeypatch.setattr( + "dstack._internal.cli.services.presets.session.get_presets_dir", + lambda: tmp_path, + ) + # A failed session is off-limits to follow/resume, but its log is readable. + session = load_agent_session("dead0000") + assert session.preset_id == "dead0000" + with pytest.raises(CLIError, match="Unknown preset"): + load_agent_session("nope0000") + + def test_print_session_log_dumps_log_verbatim(self, tmp_path, monkeypatch, capsys): + session = self._session( + tmp_path, "abcd0000", "success", "[t] trial 1 done\n[t] saved preset\n" + ) + print_session_log(session) + out = capsys.readouterr().out + assert "trial 1 done" in out + assert "saved preset" in out + + def test_print_session_log_notes_empty_log(self, tmp_path, capsys): + session = self._session(tmp_path, "empty000", "running", "") + print_session_log(session) + assert "No log output yet" in capsys.readouterr().out + + +class TestFollowPreset: + def _detached_session(self, tmp_path, configuration_yaml: str) -> PresetAgentSession: + session_dir = tmp_path / "ab12cd34" + session_dir.mkdir() + (session_dir / "agent.log").touch() + (session_dir / "preset.dstack.yml").write_text(configuration_yaml) + session = PresetAgentSession(path=session_dir, debug=False, preset_id="ab12cd34") + workspace = create_agent_workspace(session) + workspace.constraints_path.write_text('{"run_name_prefix": "qwen-build"}') + session.update_manifest(status="running", agent_pid=987654321) + return session + + def test_finalizes_a_detached_session(self, creation_context, monkeypatch, tmp_path): + session = self._detached_session( + tmp_path, "type: preset\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\n" + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.session.get_presets_dir", + lambda: tmp_path, + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.load_attachable_agent_session", + lambda preset_id: session, + ) + + async def fake_attach(**kwargs): + return PresetAgentProcessOutput( + report_data=json.loads(get_successful_preset_report(creation_context.run).json()) + ) + + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.attach_preset_agent", + fake_attach, + ) + + result = follow_preset( + api=creation_context.api, + store=creation_context.store, + preset_id="ab12cd34", + ) + + assert result.preset.id == "ab12cd34" + assert creation_context.store.get("ab12cd34") is not None + assert session.read_manifest()["status"] == "success" + + def test_agent_death_without_report_suspends_instead_of_failing( + self, creation_context, monkeypatch, tmp_path + ): + session = self._detached_session( + tmp_path, "type: preset\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\n" + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.load_attachable_agent_session", + lambda preset_id: session, + ) + + async def fake_attach(**kwargs): + return PresetAgentProcessOutput(error="agent died") + + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.attach_preset_agent", + fake_attach, + ) + + with pytest.raises(CLIError, match="agent died"): + follow_preset( + api=creation_context.api, + store=creation_context.store, + preset_id="ab12cd34", + ) + + assert session.read_manifest()["status"] == "interrupted" + + def test_backs_off_when_claim_is_held(self, creation_context, monkeypatch, tmp_path): + session = self._detached_session( + tmp_path, "type: preset\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\n" + ) + # Another live holder owns the finalize lock (reconcile or logs -f). + held = try_claim_session(session) + assert held is not None + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.load_attachable_agent_session", + lambda preset_id: session, + ) + # follow must refuse rather than double-finalize; the session is untouched. + with pytest.raises(SessionBusyError): + follow_preset( + api=creation_context.api, + store=creation_context.store, + preset_id="ab12cd34", + ) + assert session.read_manifest()["status"] == "running" + release_session_claim(held) + + +class TestStopPresetSession: + def _session_dir(self, tmp_path, manifest: dict): + session_dir = tmp_path / ".dstack" / "presets" / manifest["id"] + session_dir.mkdir(parents=True) + (session_dir / "agent.log").touch() + (session_dir / "session.json").write_text(json.dumps(manifest)) + return session_dir + + def _patch_root(self, monkeypatch, tmp_path): + monkeypatch.setattr( + "dstack._internal.cli.services.presets.session.get_presets_dir", + lambda: tmp_path / ".dstack" / "presets", + ) + + @pytest.mark.parametrize( + ("status", "message"), + [ + ("success", "is already created"), + ("failed", "creation failed"), + ("interrupted", "creation was interrupted"), + ], + ) + def test_reports_terminal_states_without_stopping( + self, tmp_path, monkeypatch, capsys, status, message + ): + self._patch_root(monkeypatch, tmp_path) + self._session_dir(tmp_path, {"id": "ab12cd34", "status": status}) + + stop_preset_session(SimpleNamespace(), "ab12cd34") + + assert message in capsys.readouterr().out + + def test_finalizes_completed_session_and_reports_created(self, tmp_path, monkeypatch, capsys): + self._patch_root(monkeypatch, tmp_path) + workspace = tmp_path / "workspace" + (workspace / "w").mkdir(parents=True) + (workspace / "w" / "final_report.json").write_text("{}") + self._session_dir( + tmp_path, + { + "id": "ab12cd34", + "status": "running", + "workspace": str(workspace), + "keep_service": True, + }, + ) + calls = [] + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.follow_preset", + lambda **kwargs: calls.append(kwargs), + ) + + stop_preset_session(SimpleNamespace(), "ab12cd34") + + # Finalized silently like reconcile, honoring the persisted keep-service. + assert len(calls) == 1 + assert calls[0]["preset_id"] == "ab12cd34" + assert calls[0]["keep_service"] is True + assert calls[0]["echo"] is False + assert "is already created" in capsys.readouterr().out + + def test_stop_wins_over_a_live_owner(self, tmp_path, monkeypatch, capsys): + self._patch_root(monkeypatch, tmp_path) + session_dir = self._session_dir( + tmp_path, + { + "id": "ab12cd34", + "status": "running", + # A live owner: this very process. + "agent_pid": os.getpid(), + "agent_started_at": None, + }, + ) + order = [] + monkeypatch.setattr( + "dstack._internal.cli.services.presets.session._pid_alive", + lambda pid, started_at=None: True, + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.terminate_agent_process", + lambda manifest: order.append("terminate"), + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create._stop_active_session_runs", + lambda api, session: order.append("stop_runs"), + ) + + def record_finish(session, status): + order.append(f"finish:{status}") + session.finish(status) + + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create._finish_agent_session", + record_finish, + ) + + stop_preset_session(SimpleNamespace(), "ab12cd34") + + # The intent is recorded before the agent dies, so a live owner's retry + # loop can never resurrect it. + assert order[0] == "finish:interrupted" + assert "terminate" in order and "stop_runs" in order + manifest = json.loads((session_dir / "session.json").read_text()) + assert manifest["status"] == "interrupted" + out = capsys.readouterr().out + assert "creation interrupted" in out + + +class TestStopActiveSessionRuns: + def _session(self, tmp_path) -> PresetAgentSession: + session_dir = tmp_path / "ab12cd34" + session_dir.mkdir() + (session_dir / "runs.jsonl").write_text( + '{"name":"qwen-build-1","id":"a"}\n{"name":"qwen-build-2","id":"b"}\n' + ) + return PresetAgentSession(path=session_dir, debug=False, preset_id="ab12cd34") + + def _api(self, statuses: dict, stopped: list) -> SimpleNamespace: + def get(project, name): + return SimpleNamespace(status=statuses[name]) + + def stop(project, names, abort): + stopped.extend(names) + + return SimpleNamespace( + project="main", + client=SimpleNamespace(runs=SimpleNamespace(get=get, stop=stop)), + ) + + @pytest.mark.parametrize( + ("statuses", "expected_stopped"), + [ + pytest.param( + {"qwen-build-1": RunStatus.DONE, "qwen-build-2": RunStatus.RUNNING}, + ["qwen-build-2"], + id="stops-only-non-terminal", + ), + pytest.param( + {"qwen-build-1": RunStatus.DONE, "qwen-build-2": RunStatus.DONE}, + [], + id="all-terminal", + ), + ], + ) + def test_stops_only_non_terminal_runs(self, tmp_path, statuses, expected_stopped): + stopped: list = [] + api = self._api(statuses, stopped) + + # No per-run prompt: active runs are stopped automatically (like dstack stop). + _stop_active_session_runs(api, self._session(tmp_path)) + assert stopped == expected_stopped + + +class TestReconcileDetachedSessions: + def _session_dir( + self, + tmp_path, + preset_id="dead0001", + *, + status="running", + project="main", + with_report=True, + keep_service=False, + owner_alive=False, + ): + session_dir = tmp_path / preset_id + (session_dir / "workspace" / "w").mkdir(parents=True) + manifest = { + "id": preset_id, + "status": status, + "keep_service": keep_service, + # A dead pid unless a live owner is requested below. + "pid": 987654321, + "pid_started_at": 0.0, + "workspace": str(session_dir / "workspace"), + } + if project is not None: + manifest["project"] = project + if owner_alive: + # A live pid with no recorded start time reads as an active owner. + manifest["agent_pid"] = os.getpid() + (session_dir / "session.json").write_text(json.dumps(manifest)) + if with_report: + (session_dir / "workspace" / "w" / "final_report.json").write_text("{}") + return session_dir + + def _patch(self, monkeypatch, tmp_path, follow): + # reconcile iterates via session.iter_agent_sessions -> session.get_presets_dir. + monkeypatch.setattr( + "dstack._internal.cli.services.presets.session.get_presets_dir", + lambda: tmp_path, + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.Client", + SimpleNamespace(from_config=lambda project_name=None: SimpleNamespace()), + ) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.follow_preset", + follow, + ) + + def _recording_follow(self, calls): + def follow(**kwargs): + calls.append(kwargs) + return SimpleNamespace( + preset=SimpleNamespace(id=kwargs["preset_id"], base="Qwen/Qwen3.5-27B") + ) + + return follow + + def test_finalizes_eligible_detached_session(self, tmp_path, monkeypatch): + self._session_dir(tmp_path, keep_service=True) + calls: list = [] + self._patch(monkeypatch, tmp_path, self._recording_follow(calls)) + reconcile_detached_sessions(PresetStore(tmp_path / "store")) + assert len(calls) == 1 + assert calls[0]["preset_id"] == "dead0001" + # Honors persisted keep-service; non-interactive, non-blocking, and + # silent (follow_preset itself takes the finalize claim). + assert calls[0]["keep_service"] is True + assert calls[0]["wait_for_run_stop"] is False + assert calls[0]["echo"] is False + + @pytest.mark.parametrize( + "kwargs", + [ + {"status": "success"}, + {"status": "failed"}, + {"with_report": False}, + {"project": None}, + {"owner_alive": True}, + # `interrupted` is `stop`'s job now, not reconcile's — with or without a report. + {"status": "interrupted"}, + {"status": "interrupted", "with_report": False}, + ], + ) + def test_skips_ineligible(self, tmp_path, monkeypatch, kwargs): + self._session_dir(tmp_path, **kwargs) + calls: list = [] + self._patch(monkeypatch, tmp_path, self._recording_follow(calls)) + reconcile_detached_sessions(PresetStore(tmp_path / "store")) + assert calls == [] + + def test_never_raises_when_finalize_fails(self, tmp_path, monkeypatch): + self._session_dir(tmp_path) + + def boom(**kwargs): + raise RuntimeError("finalize blew up") + + self._patch(monkeypatch, tmp_path, boom) + # A read command must never fail because reconcile did. + reconcile_detached_sessions(PresetStore(tmp_path / "store")) + + +class TestSessionClaim: + def _session(self, tmp_path): + (tmp_path / "sess").mkdir() + return PresetAgentSession(path=tmp_path / "sess", debug=False, preset_id="sess") + + def test_claim_is_exclusive_and_releasable(self, tmp_path): + session = self._session(tmp_path) + first = try_claim_session(session) + assert first is not None + assert try_claim_session(session) is None # the kernel lock is held + release_session_claim(first) + again = try_claim_session(session) + assert again is not None + release_session_claim(again) + + def test_claim_acquires_when_lock_file_is_unheld(self, tmp_path): + session = self._session(tmp_path) + # A leftover lock file from a crashed run holds no kernel lock: the file's + # presence must not block a new claim (no stale-lock reasoning needed). + (session.path / ".reconcile.lock").write_text("stale") + fd = try_claim_session(session) + assert fd is not None + release_session_claim(fd) + + +class TestSessionProcessAlive: + def test_recycled_pid_with_stale_start_time_is_not_alive(self): + # A live pid whose recorded start time does not match — the pid was recycled. + assert session_process_alive({"agent_pid": os.getpid(), "agent_started_at": 0.0}) is False + + def test_dead_pids_are_not_alive(self): + assert session_process_alive({"pid": 987654321, "pid_started_at": 0.0}) is False + assert session_process_alive({}) is False + + +class TestMarkSessionOwner: + def test_persists_finalize_context(self, tmp_path): + (tmp_path / "s").mkdir() + session = PresetAgentSession(path=tmp_path / "s", debug=False, preset_id="s") + session.update_manifest(status="running") + mark_session_owner(session, project="main", keep_service=True) + manifest = session.read_manifest() + assert manifest["project"] == "main" + assert manifest["keep_service"] is True + assert manifest["pid"] == os.getpid() + assert "pid_started_at" in manifest diff --git a/src/tests/_internal/cli/services/presets/test_output.py b/src/tests/_internal/cli/services/presets/test_output.py new file mode 100644 index 000000000..7e76dbe06 --- /dev/null +++ b/src/tests/_internal/cli/services/presets/test_output.py @@ -0,0 +1,171 @@ +from datetime import timedelta +from io import StringIO + +import pytest +from rich.table import Table + +from dstack._internal.cli.services.presets import output as output_module +from dstack._internal.cli.services.presets.output import _add_session, _format_number +from tests._internal.cli.common import plain_console +from tests._internal.cli.preset_factories import get_preset + +pytestmark = pytest.mark.windows + + +class TestFormatNumber: + def test_large_values_avoid_scientific_notation(self): + assert _format_number(1723.4) == "1723" + assert _format_number(999.6) == "1000" + + def test_small_values_keep_three_significant_digits(self): + assert _format_number(384.8) == "385" + assert _format_number(15.92) == "15.9" + assert _format_number(0.99) == "0.99" + + +class TestFormatPresetBenchmark: + def test_formats_second_scale_ttft_without_scientific_notation(self): + preset = get_preset() + ttft = preset.validations[0].benchmark.metrics.ttft_ms + ttft.mean = 8148.3 + ttft.p50 = 8151.4 + ttft.p99 = 8334.2 + + output = output_module.format_preset_benchmark(preset, verbose=True) + + assert output.startswith("ctx=32K ") + assert "TTFT 8.15s" in output + assert "e+03" not in output + + +class TestPrintPresets: + def test_preserves_benchmark_concurrency_at_narrow_width(self, monkeypatch): + output = StringIO() + monkeypatch.setattr(output_module, "console", plain_console(output, width=79)) + + output_module.print_presets([get_preset()]) + + assert "con=1" in "".join(output.getvalue().split()) + + def test_prints_submitted_column(self, monkeypatch): + output = StringIO() + monkeypatch.setattr(output_module, "console", plain_console(output, width=160)) + monkeypatch.setattr(output_module, "pretty_date", lambda _: "2 months ago") + + output_module.print_presets([get_preset()]) + + assert "SUBMITTED" in output.getvalue() + assert "2 months ago" in output.getvalue() + + +def _session_row(session: dict) -> dict: + table = Table(box=None) + for column in ("BASE", "ID", "GPU", "BENCHMARK", "STATUS", "SUBMITTED"): + table.add_column(column) + _add_session(table, session) + return { + column.header: "".join(str(cell) for cell in column._cells) for column in table.columns + } + + +class TestSessionRow: + def test_shows_progress_after_status_and_best_benchmark(self): + row = _session_row( + { + "id": "c7e18d52", + "status": "running", + "max_trials": 3, + "trials": { + "count": 2, + "best": {"tok_s": 2339.0, "concurrency": 8, "gpu": "A40:48GB:1"}, + }, + } + ) + + assert row["STATUS"] == "[bold sea_green3]clauding[/] [secondary](2/3)[/]" + assert row["BENCHMARK"] == "best trial: con=8 2339 tok/s" + assert row["GPU"] == "A40:48GB:1" + + def test_shows_zero_progress_without_benchmark(self): + row = _session_row( + {"id": "ab12cd34", "status": "running", "max_trials": 3, "trials": {"count": 0}} + ) + + assert row["STATUS"] == "[bold sea_green3]clauding[/] [secondary](0/3)[/]" + assert row["BENCHMARK"] == "" + + def test_omits_progress_without_trials_data(self): + row = _session_row({"id": "ab12cd34", "status": "interrupted"}) + + assert row["STATUS"] == "[bold gold1]interrupted[/]" + + def test_counts_without_max_trials(self): + row = _session_row({"id": "ab12cd34", "status": "interrupted", "trials": {"count": 2}}) + + assert row["STATUS"] == "[bold gold1]interrupted[/] [secondary](2)[/]" + + +class TestGroupOrdering: + def test_sorts_presets_and_sessions_newest_first(self, monkeypatch): + buffer = StringIO() + monkeypatch.setattr(output_module, "console", plain_console(buffer, width=200)) + old = get_preset() + new = old.copy(update={"id": "11aa22bb", "created_at": old.created_at + timedelta(days=2)}) + sessions = [ + { + "id": "aaaaaaaa", + "status": "interrupted", + "model": old.base, + "created_at": "2026-07-01T00:00:00+00:00", + }, + { + "id": "bbbbbbbb", + "status": "interrupted", + "model": old.base, + "created_at": "2026-07-02T00:00:00+00:00", + }, + ] + + output_module.print_presets([old, new], sessions=sessions) + + text = buffer.getvalue() + assert text.index("11aa22bb") < text.index(old.id) + assert text.index(old.id) < text.index("bbbbbbbb") < text.index("aaaaaaaa") + + +class TestDoneProgress: + def test_completed_creation_decorates_preset_row_without_extra_session_row(self, monkeypatch): + buffer = StringIO() + monkeypatch.setattr(output_module, "console", plain_console(buffer, width=200)) + preset = get_preset() + sessions = [ + { + "id": preset.id, + "status": "success", + "model": preset.base, + "max_trials": 4, + "trials": {"count": 3}, + } + ] + + output_module.print_presets([preset], sessions=sessions) + + text = buffer.getvalue() + assert "verified (3/4)" in text + assert text.count(preset.id) == 1 + + +class TestVerifyingStatus: + def test_running_session_with_exhausted_trials_shows_verifying(self): + row = _session_row( + {"id": "ab12cd34", "status": "running", "max_trials": 2, "trials": {"count": 2}} + ) + + assert row["STATUS"] == "[bold deep_sky_blue1]verifying[/] [secondary](2/2)[/]" + + def test_running_session_with_remaining_trials_stays_clauding(self): + row = _session_row( + {"id": "ab12cd34", "status": "running", "max_trials": 2, "trials": {"count": 1}} + ) + + assert row["STATUS"].startswith("[bold sea_green3]clauding[/]") diff --git a/src/tests/_internal/cli/services/presets/test_prompt.py b/src/tests/_internal/cli/services/presets/test_prompt.py new file mode 100644 index 000000000..111576340 --- /dev/null +++ b/src/tests/_internal/cli/services/presets/test_prompt.py @@ -0,0 +1,44 @@ +import pytest + +from dstack._internal.cli.services.presets import prompt as prompt_module +from dstack._internal.cli.services.presets.prompt import get_preset_agent_system_prompt +from dstack._internal.core.errors import CLIError + +pytestmark = pytest.mark.windows + + +class TestSystemPrompt: + def test_stays_byte_identical_without_user_prompt(self): + text = get_preset_agent_system_prompt() + + assert text == get_preset_agent_system_prompt(None) == get_preset_agent_system_prompt("") + assert "## Additional instructions" not in text + assert " more.\n") + monkeypatch.setattr(prompt_module, "_SYSTEM_PROMPT_PATH", broken) + + with pytest.raises(CLIError, match="Unknown variable"): + get_preset_agent_system_prompt() diff --git a/src/tests/_internal/cli/services/presets/test_store.py b/src/tests/_internal/cli/services/presets/test_store.py new file mode 100644 index 000000000..a5b66cfe7 --- /dev/null +++ b/src/tests/_internal/cli/services/presets/test_store.py @@ -0,0 +1,182 @@ +from io import StringIO +from pathlib import Path +from unittest.mock import patch + +import pytest +import yaml + +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.services.presets import store as store_module +from dstack._internal.cli.services.presets.store import PresetStore +from dstack._internal.core.errors import ConfigurationError +from dstack._internal.core.models.envs import EnvSentinel +from tests._internal.cli.preset_factories import get_preset + +pytestmark = pytest.mark.windows + + +class TestPresetStore: + def test_saves_and_lists_self_contained_preset(self, tmp_path: Path): + store = PresetStore(tmp_path / "presets") + preset = get_preset() + + path = store.save(preset) + + assert path == (tmp_path / "presets" / "8f3a12c4" / "preset.yaml") + data = yaml.safe_load(path.read_text()) + assert data["base"] == preset.base + assert data["id"] == preset.id + assert data["model"] == preset.model + assert data["created_at"] == preset.created_at.isoformat() + assert "presets" not in data + assert store.list() == [preset] + assert store.get(preset.id) == preset + assert not list(path.parent.glob("*.tmp")) + + def test_saving_same_id_overwrites_existing_preset(self, tmp_path: Path): + store = PresetStore(tmp_path / "presets") + preset = get_preset() + store.save(preset) + + updated = preset.copy(update={"base": "Qwen/Another-Model", "context_length": 16384}) + store.save(updated) + + assert store.get(preset.id) == updated + + def test_migrates_legacy_layout_and_archives_on_delete(self, tmp_path: Path): + root = tmp_path / "presets" + store = PresetStore(root) + preset = get_preset() + legacy = root / "models--Qwen--Qwen3.5-27B" + legacy.mkdir(parents=True) + (legacy / f"{preset.id}.yaml").write_text( + yaml.safe_dump( + yaml.safe_load(PresetStore(tmp_path / "tmp").save(preset).read_text()), + sort_keys=False, + ) + ) + + assert store.list() == [preset] + assert (root / preset.id / "preset.yaml").is_file() + assert not legacy.exists() + + assert store.delete(preset.id) is True + assert store.get(preset.id) is None + assert (root / ".archive" / preset.id / "preset.yaml").is_file() + assert store.list() == [] + + def test_skips_invalid_preset_on_list_but_keeps_it_deletable(self, tmp_path: Path, capsys): + store = PresetStore(tmp_path / "presets") + valid = get_preset().copy(update={"id": "01234567"}) + store.save(valid) + path = store.save(get_preset()) + data = yaml.safe_load(path.read_text()) + data["validations"][0].pop("benchmark") + path.write_text(yaml.safe_dump(data, sort_keys=False)) + + # One corrupt file must not take down every read; the warning goes to + # stderr so --json stdout stays parseable... + assert [preset.id for preset in store.list()] == [valid.id] + assert "benchmark" in capsys.readouterr().err + # ...and the corrupt preset is still removable by ID. + assert store.delete(get_preset().id) is True + assert [preset.id for preset in store.list()] == [valid.id] + + def test_preserves_literal_env_values(self, tmp_path: Path): + store = PresetStore(tmp_path / "presets") + preset = get_preset() + preset.service.env.update( + { + "TOKENIZERS_PARALLELISM": "false", + "MODEL_LABEL": "monkey", + "HF_TOKEN": EnvSentinel(key="HF_TOKEN"), + } + ) + + path = store.save(preset) + env = yaml.safe_load(path.read_text())["service"]["env"] + + assert "TOKENIZERS_PARALLELISM=false" in env + assert "MODEL_LABEL=monkey" in env + assert "HF_TOKEN" in env + + +class TestParsePresetConfiguration: + @pytest.mark.parametrize("key", ["base", "repo"]) + def test_warns_on_nested_model_without_name(self, key: str): + stream = StringIO(f"type: preset\nmodel:\n {key}: Qwen/Qwen3.5-27B\n") + + with patch.object(store_module, "warn") as warn: + configuration = store_module._parse_preset_configuration(stream) + + warn.assert_called_once() + assert f"model.{key}" in warn.call_args.args[0] + assert f"`{key}:`" in warn.call_args.args[0] + assert configuration.model is not None + + @pytest.mark.parametrize( + "body", + [ + "base: Qwen/Qwen3.5-27B\n", + "repo: Qwen/Qwen3.5-27B\n", + "model: Qwen/Qwen3.5-27B\n", + "model:\n repo: community/Qwen3.5-27B-GPTQ-Int4\n name: Qwen/Qwen3.5-27B\n", + ], + ) + def test_does_not_warn_on_preferred_syntax(self, body: str): + stream = StringIO(f"type: preset\n{body}") + + with patch.object(store_module, "warn") as warn: + configuration = store_module._parse_preset_configuration(stream) + + warn.assert_not_called() + assert configuration.model is not None + + +class TestResolvePresetPrompt: + def test_resolves_inline_and_file_relative_to_configuration(self, tmp_path: Path): + (tmp_path / "notes.md").write_text("From a file.\n") + configuration_path = str(tmp_path / "preset.dstack.yml") + inline = PresetConfiguration(name="q", base="Q/M", prompt="Inline text.") + from_file = PresetConfiguration(name="q", base="Q/M", prompt={"path": "notes.md"}) + + assert store_module.resolve_preset_prompt(inline, configuration_path) == "Inline text." + assert store_module.resolve_preset_prompt(from_file, configuration_path) == "From a file." + assert ( + store_module.resolve_preset_prompt( + PresetConfiguration(name="q", base="Q/M"), configuration_path + ) + is None + ) + + def test_rejects_missing_and_empty_prompt_files(self, tmp_path: Path): + configuration_path = str(tmp_path / "preset.dstack.yml") + (tmp_path / "empty.md").write_text(" \n") + + with pytest.raises(ConfigurationError, match="Failed to read"): + store_module.resolve_preset_prompt( + PresetConfiguration(name="q", base="Q/M", prompt={"path": "missing.md"}), + configuration_path, + ) + with pytest.raises(ConfigurationError, match="is empty"): + store_module.resolve_preset_prompt( + PresetConfiguration(name="q", base="Q/M", prompt={"path": "empty.md"}), + configuration_path, + ) + + +class TestPresetNames: + def test_finds_and_detaches_names(self, tmp_path: Path): + store = PresetStore(tmp_path / "presets") + named = get_preset().copy(update={"name": "qwen"}) + store.save(named) + store.save(get_preset().copy(update={"id": "01234567"})) + + assert store.find_by_name("qwen").id == named.id + assert store.find_by_name("other") is None + + released = store.release_name("qwen") + + assert released.id == named.id + assert store.find_by_name("qwen") is None + assert store.get(named.id).name is None diff --git a/src/tests/_internal/cli/services/presets/test_verify.py b/src/tests/_internal/cli/services/presets/test_verify.py new file mode 100644 index 000000000..793bcda71 --- /dev/null +++ b/src/tests/_internal/cli/services/presets/test_verify.py @@ -0,0 +1,139 @@ +from datetime import datetime, timezone +from unittest.mock import patch + +import pytest +from pydantic import ValidationError + +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.models.preset_agent import AgentFinalReport +from dstack._internal.cli.services.presets.agent import ( + PresetAgentProcessOutput, +) +from dstack._internal.cli.services.presets.verify import ( + build_verified_preset, + load_preset_agent_report, +) +from dstack._internal.cli.services.presets.workspace import ( + PresetAgentWorkspace, +) +from dstack._internal.core.errors import CLIError +from dstack._internal.core.models.envs import EnvSentinel +from dstack._internal.core.models.profiles import ProfileParams +from tests._internal.cli.preset_factories import ( + get_running_service_run, + get_successful_preset_report, +) + +pytestmark = pytest.mark.windows + + +class TestBuildVerifiedPreset: + def test_successful_report_requires_benchmark(self): + run = get_running_service_run() + data = get_successful_preset_report(run).dict() + data.pop("benchmark") + + with pytest.raises(ValidationError, match="benchmark"): + AgentFinalReport.parse_obj(data) + + def test_builds_portable_self_contained_preset(self): + run = get_running_service_run() + created_at = datetime(2026, 1, 2, 3, 4, tzinfo=timezone.utc) + + with patch( + "dstack._internal.cli.services.presets.presets.get_current_datetime", + return_value=created_at, + ): + preset = build_verified_preset( + run=run, + preset_configuration=PresetConfiguration( + name="qwen-build", + model={"base": "Qwen/Qwen3.5-27B"}, + context_length=8192, + gateway="benchmark-gateway", + env=["LICENSE", "TOKENIZERS_PARALLELISM=false"], + ), + report=get_successful_preset_report(run), + ) + + assert preset.base == "Qwen/Qwen3.5-27B" + assert preset.model == "community/Qwen3.5-27B-GPTQ-Int4" + assert preset.context_length == 32768 + assert preset.created_at == created_at + assert preset.service.name is None + assert preset.service.gateway is None + assert all(getattr(preset.service, field) is None for field in ProfileParams.__fields__) + assert isinstance(preset.service.env["LICENSE"], EnvSentinel) + assert preset.service.env["TOKENIZERS_PARALLELISM"] == "false" + assert preset.service.resources.gpu.vendor.value == "nvidia" + validation = preset.validations[0] + assert validation.replicas[0].resources[0].gpu.name == ["A6000"] + assert validation.benchmark.target.type == "server-proxy" + assert validation.benchmark.client.type == "local" + + def test_rejects_variant_for_exact_model_request(self): + run = get_running_service_run() + report = get_successful_preset_report(run).copy(update={"model": "other/model"}) + + with pytest.raises(CLIError, match="changed an exact model request"): + build_verified_preset( + run=run, + preset_configuration=PresetConfiguration( + name="qwen-build", + model={ + "repo": "community/Qwen3.5-27B-GPTQ-Int4", + "name": "Qwen/Qwen3.5-27B", + }, + ), + report=report, + ) + + +class TestLoadPresetAgentReport: + def _load(self, tmp_path, report_data, redacted_values): + return load_preset_agent_report( + output=PresetAgentProcessOutput(report_data=report_data), + workspace=PresetAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home"), + redacted_values=redacted_values, + ) + + def test_redacts_known_secret_in_benchmark_command_instead_of_failing(self, tmp_path): + run = get_running_service_run() + data = get_successful_preset_report(run).dict() + data["run_id"] = str(data["run_id"]) + data["benchmark"]["command"] = ( + "python bench.py --header 'Authorization: Bearer sk-live-0123456789abcdef'" + ) + + report = self._load(tmp_path, data, redacted_values=("sk-live-0123456789abcdef",)) + + assert report.benchmark is not None + assert report.benchmark.command.endswith("Bearer [redacted]'") + assert "sk-live" not in report.benchmark.command + + def test_still_rejects_unknown_bearer_token(self, tmp_path): + run = get_running_service_run() + data = get_successful_preset_report(run).dict() + data["run_id"] = str(data["run_id"]) + data["benchmark"]["command"] = ( + "curl -H 'Authorization: Bearer sk-unknown-9876543210fedcba'" + ) + + with pytest.raises(CLIError, match="bearer token"): + self._load(tmp_path, data, redacted_values=("some-other-secret-value",)) + + def test_allows_bearer_prose_without_credential(self, tmp_path): + # Regression: "(auth via DSTACK_TOKEN bearer header from env)" failed + # two live sessions — the word after "bearer" is prose, not a token. + run = get_running_service_run() + data = get_successful_preset_report(run).dict() + data["run_id"] = str(data["run_id"]) + data["benchmark"]["command"] = ( + "./benchenv/bin/python bench_service.py --base $DSTACK_SERVER_URL/x" + " (auth via DSTACK_TOKEN bearer header from env)" + ) + + report = self._load(tmp_path, data, redacted_values=()) + + assert report.benchmark is not None + assert "bearer header" in report.benchmark.command