diff --git a/src/dstack/_internal/cli/commands/preset.py b/src/dstack/_internal/cli/commands/preset.py index 5f9952cc8..827452193 100644 --- a/src/dstack/_internal/cli/commands/preset.py +++ b/src/dstack/_internal/cli/commands/preset.py @@ -1,7 +1,9 @@ import argparse +import io import os +import sys import time -from contextlib import suppress +from contextlib import redirect_stderr, suppress from pathlib import Path from argcomplete import FilesCompleter # type: ignore[attr-defined] @@ -14,6 +16,7 @@ PresetListOutput, ) from dstack._internal.cli.services.completion import ProjectNameCompleter +from dstack._internal.cli.services.configurators import APPLY_STDIN_NAME from dstack._internal.cli.services.presets.apply import apply_preset from dstack._internal.cli.services.presets.create import ( CreationStopped, @@ -28,13 +31,14 @@ ) 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, + list_preset_sessions, + load_resumable_session, resolve_session_ref, ) from dstack._internal.cli.services.presets.store import ( PresetStore, load_preset_configuration, + parse_preset_configuration, resolve_preset_prompt, ) from dstack._internal.cli.services.profile import ( @@ -251,9 +255,13 @@ def _list(self, args: argparse.Namespace) -> None: limit=args.limit, ) return + # The store warns about unreadable presets on stderr once per read; + # inside Live that would tear the render on every refresh. The first + # read happens before Live starts so warnings print once, above the + # table; refreshes read with stderr suppressed. + presets, sessions = self._list_presets_and_sessions(base=base, repo=repo) 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, @@ -264,13 +272,15 @@ def _list(self, args: argparse.Namespace) -> None: ) ) time.sleep(LIVE_TABLE_PROVISION_INTERVAL_SECS) + with redirect_stderr(io.StringIO()): + presets, sessions = self._list_presets_and_sessions(base=base, repo=repo) 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() + sessions = list_preset_sessions() if base or repo: repo_to_base = {preset.model: preset.base for preset in presets} presets = _filter_presets(presets, base=base, repo=repo) @@ -282,13 +292,14 @@ def _list_presets_and_sessions( return presets, sessions def _create(self, args: argparse.Namespace) -> None: - configuration_path, configuration = load_preset_configuration(args.configuration_file) + _check_stdin_configuration_confirmable(args) + _, configuration = _read_configuration_arg(args.configuration_file) configuration = _get_effective_configuration(configuration, args, require_name=False) - user_prompt = resolve_preset_prompt(configuration, configuration_path) + user_prompt = resolve_preset_prompt(configuration, _prompt_base(args.configuration_file)) store = PresetStore() resume_session = None if getattr(args, "resume", None): - resume_session = load_resumable_agent_session(args.resume) + resume_session = load_resumable_session(args.resume) if getattr(args, "trials", None) is not None: console.print( "[warning]--trials is ignored when resuming: " @@ -368,7 +379,7 @@ def _get(self, args: argparse.Namespace) -> None: def _apply(self, args: argparse.Namespace) -> None: self._reconcile() - configuration_path, configuration = load_preset_configuration(args.configuration_file) + configuration_path, configuration = _read_configuration_arg(args.configuration_file) configuration = _get_effective_configuration(configuration, args) apply_preset( api=Client.from_config(project_name=args.project), @@ -532,6 +543,28 @@ def _confirm_preset_creation(store: PresetStore, name: str | None, *, assume_yes return True +def _read_configuration_arg(configuration_file: str) -> tuple[str, PresetConfiguration]: + """`-f `, or `-f -` for stdin — the same convention as `dstack apply`.""" + if configuration_file == APPLY_STDIN_NAME: + return APPLY_STDIN_NAME, parse_preset_configuration(sys.stdin) + path = Path(configuration_file) + return str(path.resolve()), load_preset_configuration(path) + + +def _prompt_base(configuration_file: str) -> Path: + """Prompt files resolve relative to the configuration file; cwd for stdin.""" + if configuration_file == APPLY_STDIN_NAME: + return Path.cwd() + return Path(configuration_file).resolve().parent + + +def _check_stdin_configuration_confirmable(args: argparse.Namespace) -> None: + # Same rule as `dstack apply`: the confirmation prompt cannot read from a + # stdin that is the configuration itself. + if not args.yes and args.configuration_file == APPLY_STDIN_NAME: + raise CLIError("Cannot read configuration from stdin if -y/--yes is not specified") + + def _get_effective_configuration( configuration: PresetConfiguration, args: argparse.Namespace, diff --git a/src/dstack/_internal/cli/models/configurations.py b/src/dstack/_internal/cli/models/configurations.py index 856b6c3e4..94b1fc578 100644 --- a/src/dstack/_internal/cli/models/configurations.py +++ b/src/dstack/_internal/cli/models/configurations.py @@ -135,14 +135,15 @@ class PresetConfiguration( ), ] = None min_context_length: Annotated[ - Optional[PositiveInt], Field(description="The minimum required context length") + Optional[PositiveInt], + Field(description="The minimum required context length. Required for creation"), ] = None max_ttft: Annotated[ Optional[PositiveInt], Field( description=( "The maximum p50 time to first token, in milliseconds, that any benchmark" - " may report" + " may report. Required for creation" ) ), ] = None @@ -151,7 +152,7 @@ class PresetConfiguration( Field( description=( "The number of benchmarked trials during preset creation" - " before the best one is promoted" + " before the best one is promoted. Required for creation" ) ), ] = None @@ -168,7 +169,8 @@ class PresetConfiguration( Optional[PositiveInt], Field( description=( - "The number of simultaneous requests used for benchmarks during preset creation" + "The number of simultaneous requests used for benchmarks during preset" + " creation. Required for creation" ) ), ] = None @@ -182,7 +184,7 @@ class PresetConfiguration( ), ] = None output_tokens: Annotated[ - Optional[PositiveInt], + Optional[Annotated[int, Field(ge=2)]], Field( description=( "The number of output tokens per request used for benchmarks during" @@ -191,7 +193,7 @@ class PresetConfiguration( ), ] = None shared_prefix_tokens: Annotated[ - Optional[PositiveInt], + Optional[Annotated[int, Field(ge=0)]], Field( description=( "How many of `input_tokens` are a prefix identical in every benchmark request," @@ -338,13 +340,23 @@ class PresetConstraints(CoreModel): max_ttft: PositiveInt trials_num: PositiveInt concurrency: PositiveInt - input_tokens: Optional[PositiveInt] = None - output_tokens: Optional[PositiveInt] = None - shared_prefix_tokens: Optional[int] = None - dataset: Optional[str] = None - baseline: bool = False - fleets: list[str] = Field(min_length=1) - env: list[str] = [] + baseline: bool + fleets: Annotated[list[str], Field(min_length=1)] + env: list[str] + + +class PresetRandomConstraints(PresetConstraints): + """Constraints for the synthetic `random` dataset, which the request shape defines.""" + + input_tokens: PositiveInt + output_tokens: Annotated[int, Field(ge=2)] + shared_prefix_tokens: Annotated[int, Field(ge=0)] + + +class PresetDatasetConstraints(PresetConstraints): + """Constraints for a named dataset, which defines its own request shape.""" + + dataset: str def _validate_model(value: Any, *, field: str) -> str: diff --git a/src/dstack/_internal/cli/models/preset_agent.py b/src/dstack/_internal/cli/models/preset_agent.py index bfa921c1d..abc06e099 100644 --- a/src/dstack/_internal/cli/models/preset_agent.py +++ b/src/dstack/_internal/cli/models/preset_agent.py @@ -1,152 +1,158 @@ import uuid -from typing import Any, Dict, Optional - -from pydantic import PositiveInt, model_validator -from typing_extensions import Self +from datetime import datetime +from typing import Annotated, Any, Dict, Literal, Optional, Union + +import yaml +from pydantic import ( + Field, + PositiveInt, + ValidationError, + ValidatorFunctionWrapHandler, + WrapValidator, +) from dstack._internal.cli.models.presets import PresetBenchmark from dstack._internal.core.models.common import CoreModel +from dstack._internal.core.models.configurations import ServiceConfiguration + + +class PresetAgentInvalidService(CoreModel): + """A reported service YAML that does not parse; the reason, not a failure.""" + + raw: str + error: str + + +def _service_from_yaml(value: Any, handler: ValidatorFunctionWrapHandler) -> Any: + if not isinstance(value, str): + return handler(value) + try: + return ServiceConfiguration.model_validate(yaml.safe_load(value)) + except ValidationError as e: + return PresetAgentInvalidService(raw=value, error=e.errors()[0]["msg"]) + except yaml.YAMLError as e: + return PresetAgentInvalidService(raw=value, error=str(e)) + + +# The agent writes YAML text; the parse yields the service or the typed reason +# it is not one. The schema keeps telling the agent `string`. +ReportedService = Annotated[ + Union[ServiceConfiguration, PresetAgentInvalidService], + WrapValidator(_service_from_yaml, json_schema_input_type=str), +] + + +class PresetAgentSuccess(CoreModel): + """A preset the agent created, cross-checked against the run in `verify`.""" + + success: Literal[True] + run_id: uuid.UUID + run_name: str + service_yaml: ReportedService + trial: PositiveInt + base: Annotated[str, Field(min_length=1)] + model: Annotated[str, Field(min_length=1)] + context_length: PositiveInt + benchmark: PresetBenchmark + + +class PresetAgentFailure(CoreModel): + """Why the agent created no preset.""" + + success: Literal[False] + failure_summary: str + + +AnyPresetAgentResult = Union[PresetAgentSuccess, PresetAgentFailure] + + +PresetSessionStatus = Literal["running", "interrupted", "success", "failed"] + +# The session models carry no defaults: `session.json` is this CLI's own state +# file, so every write states the whole state, and a `null` in the file is a +# statement (a detached session has `owner: null`), not an omission. + + +class PresetSessionProcess(CoreModel): + pid: int + # The start time guards against pid reuse; psutil may fail to provide it. + started_at: Optional[float] + + +class PresetSessionWorkspace(CoreModel): + path: str + # Aliased under a short stable path because SSH control sockets embedded in + # the workspace have a hard length limit. + alias: str + + +class PresetSessionFinalize(CoreModel): + """What a later detached reconcile needs to finalize the session.""" + + project: str + keep_service: bool + + +class PresetSessionRun(CoreModel): + """One agent run over the session, recorded whole when the run begins.""" + + workspace: PresetSessionWorkspace + finalize: PresetSessionFinalize + # Only known when this CLI launched the agent; a follower leaves it as is. + claude_model: Optional[str] + # None between claude process attempts and after a detach outlives them. + agent: Optional[PresetSessionProcess] + # None until the agent's stream reveals it. + claude_session_id: Optional[str] + + +class PresetSessionState(CoreModel): + """`session.json`. The record fields never change after creation; `status`, + `owner`, and `run` advance only through the `PresetSession` transitions.""" + + id: str + # Released when a newer preset claims the name — the one record mutation. + name: Optional[str] + model: str + # None when the configuration left the trial count to the agent. + trials_num: Optional[int] + previous: list[str] + created_at: datetime + debug: bool + status: PresetSessionStatus + # None is a detached session. + owner: Optional[PresetSessionProcess] + # None is a session that never began a run. + run: Optional[PresetSessionRun] + + +class ClaudeStreamEvent(CoreModel): + """One line of the claude CLI's `--output-format stream-json`. Not our format: + unknown fields are dropped and omitted fields default.""" + + # "system", "assistant", "user", "result", and whatever a newer claude CLI + # adds — which is why this is not a Literal. + type: str + # Identifies the claude conversation, so an interrupted creation can be + # resumed with `claude --resume`. Not every line carries it. + session_id: Optional[str] = None + + +class ClaudeResultEvent(ClaudeStreamEvent): + """The final line: the agent's structured output, or its error.""" + + type: Literal["result"] + is_error: bool = False + # An error's message, or the report itself when the agent printed it instead + # of submitting it through `StructuredOutput`; None when there was no text. + result: Optional[Any] = None + # The agent's final report, raw until `verify` parses it — that parse redacts + # known secret values as it validates, so nothing downstream sees them. + # None when the agent never submitted one. + structured_output: Optional[Dict[str, Any]] = None + -_LATENCY_JSON_SCHEMA = { - "type": "object", - "properties": { - "mean": {"type": "number", "minimum": 0}, - "p50": {"type": "number", "minimum": 0}, - "p99": {"type": "number", "minimum": 0}, - }, - "required": ["mean", "p50", "p99"], - "additionalProperties": False, -} - -_BENCHMARK_JSON_SCHEMA = { - "type": "object", - "properties": { - "tool": {"type": "string", "minLength": 1}, - "tool_version": {"type": "string", "minLength": 1}, - "command": {"type": "string", "minLength": 1}, - "workload": { - "type": "object", - "properties": { - "api": { - "type": "string", - "enum": ["chat_completions", "completions"], - }, - "num_requests": {"type": "integer", "minimum": 1}, - "input_tokens": {"type": "integer", "minimum": 1}, - "output_tokens": {"type": "integer", "minimum": 2}, - "concurrency": {"type": "integer", "minimum": 1}, - "shared_prefix_tokens": {"type": "integer", "minimum": 0}, - "dataset": {"type": "string", "minLength": 1}, - }, - # `shared_prefix_tokens` and `dataset` are not required: one schema - # serves both session modes, and each mode knows only its own field. - "required": [ - "api", - "num_requests", - "input_tokens", - "output_tokens", - "concurrency", - ], - "additionalProperties": False, - }, - "metrics": { - "type": "object", - "properties": { - "successful_requests": {"type": "integer", "minimum": 0}, - "failed_requests": {"type": "integer", "minimum": 0}, - "duration_seconds": {"type": "number", "exclusiveMinimum": 0}, - "total_input_tokens": {"type": "integer", "minimum": 0}, - "total_output_tokens": {"type": "integer", "minimum": 0}, - "output_tok_per_s": {"type": "number", "exclusiveMinimum": 0}, - "per_user_tok_per_s": {"type": "number", "exclusiveMinimum": 0}, - "ttft_ms": _LATENCY_JSON_SCHEMA, - "tpot_ms": _LATENCY_JSON_SCHEMA, - }, - "required": [ - "successful_requests", - "failed_requests", - "duration_seconds", - "total_input_tokens", - "total_output_tokens", - "output_tok_per_s", - "per_user_tok_per_s", - "ttft_ms", - "tpot_ms", - ], - "additionalProperties": False, - }, - }, - "required": ["tool", "tool_version", "command", "workload", "metrics"], - "additionalProperties": False, -} - -AGENT_FINAL_REPORT_JSON_SCHEMA = { - "type": "object", - "properties": { - "success": {"type": "boolean"}, - "run_id": {"type": "string"}, - "run_name": {"type": "string"}, - "service_yaml": {"type": "string"}, - "trial": {"type": "integer", "minimum": 1}, - "base": {"type": "string"}, - "model": {"type": "string"}, - "context_length": {"type": "integer", "minimum": 1}, - "benchmark": _BENCHMARK_JSON_SCHEMA, - "failure_summary": {"type": "string"}, - }, - "required": ["success"], - "additionalProperties": False, -} - - -class AgentFinalReport(CoreModel): - success: bool - run_id: Optional[uuid.UUID] = None - run_name: Optional[str] = None - service_yaml: Optional[str] = None - trial: Optional[PositiveInt] = None - base: Optional[str] = None - model: Optional[str] = None - context_length: Optional[PositiveInt] = None - benchmark: Optional[PresetBenchmark] = None - failure_summary: Optional[str] = None - - @model_validator(mode="after") - def validate_report(self) -> Self: - if self.success: - required = ( - "run_id", - "run_name", - "service_yaml", - "trial", - "base", - "model", - "context_length", - "benchmark", - ) - missing = [field for field in required if getattr(self, field) in (None, "")] - if missing: - raise ValueError("successful agent report must include " + ", ".join(missing)) - elif not self.failure_summary: - raise ValueError("failed agent report must include failure_summary") - return self - - -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] +# Left to right so a result line always reads as `ClaudeResultEvent`. +AnyClaudeStreamEvent = Annotated[ + Union[ClaudeResultEvent, ClaudeStreamEvent], Field(union_mode="left_to_right") +] diff --git a/src/dstack/_internal/cli/models/presets.py b/src/dstack/_internal/cli/models/presets.py index 87dde8f15..f7e1f2565 100644 --- a/src/dstack/_internal/cli/models/presets.py +++ b/src/dstack/_internal/cli/models/presets.py @@ -1,6 +1,6 @@ import re from datetime import datetime -from typing import Annotated, Literal, Optional +from typing import Annotated, Literal, Optional, Union from pydantic import ( Field, @@ -11,25 +11,29 @@ ) from typing_extensions import Self +from dstack._internal.cli.models.configurations import PresetConfiguration from dstack._internal.core.models.common import CoreModel from dstack._internal.core.models.configurations import ServiceConfiguration from dstack._internal.core.models.profiles import ProfileParams -from dstack._internal.core.models.resources import ResourcesSpec +from dstack._internal.core.models.resources import Range, ResourcesSpec +# The service name, the gateway, and the profile parameters are chosen by whoever +# runs `dstack apply` with the preset, so a preset never carries them. +PRESET_EXCLUDED_FIELDS = ("name", "gateway", *ProfileParams.model_fields) -class PresetBenchmarkWorkload(CoreModel): + +class PresetWorkload(CoreModel): api: Literal["chat_completions", "completions"] + dataset: str num_requests: PositiveInt - # With a dataset other than `random`, the measured means rather than the - # configured request shape. input_tokens: PositiveInt output_tokens: Annotated[int, Field(ge=2)] concurrency: PositiveInt - # Absent for presets saved before the field existed, and with a dataset - # other than `random`, where the dataset decides prefix sharing. - shared_prefix_tokens: Annotated[Optional[int], Field(ge=0)] = None - # Absent means the synthetic `random` dataset. - dataset: Optional[str] = None + + +class PresetRandomWorkload(PresetWorkload): + dataset: Literal["random"] = "random" + shared_prefix_tokens: Annotated[int, Field(ge=0)] = 0 class PresetBenchmarkLatency(CoreModel): @@ -44,30 +48,26 @@ class PresetBenchmarkMetrics(CoreModel): duration_seconds: PositiveFloat total_input_tokens: Annotated[int, Field(ge=0)] total_output_tokens: Annotated[int, Field(ge=0)] - # Defaulted rather than required: presets saved before these fields existed - # must still load. The `effective_*` properties derive them when absent. - output_tok_per_s: Optional[PositiveFloat] = None - per_user_tok_per_s: Optional[PositiveFloat] = None + # Stored as reported, but never read back: `effective_*` recomputes both + # from the totals rather than trusting self-reported rates. + output_tok_per_s: PositiveFloat + per_user_tok_per_s: PositiveFloat ttft_ms: PresetBenchmarkLatency tpot_ms: PresetBenchmarkLatency -class PresetBenchmarkTarget(CoreModel): - type: Literal["gateway", "server-proxy"] - - -class PresetBenchmarkClient(CoreModel): - type: Literal["local"] - - class PresetBenchmark(CoreModel): - tool: str - tool_version: str - command: str - workload: PresetBenchmarkWorkload + """The agent reports its benchmark in exactly this shape, and is forced to by + the schema generated from it. Changing a field here means also changing the + `## Benchmark` section of the system prompt, which tells it what to put there.""" + + tool: Annotated[str, Field(min_length=1)] + tool_version: Annotated[str, Field(min_length=1)] + command: Annotated[str, Field(min_length=1)] + # The subclass first: a report without `dataset` is a random workload, and a + # base-typed field would reject its `shared_prefix_tokens` as unknown. + workload: Union[PresetRandomWorkload, PresetWorkload] metrics: PresetBenchmarkMetrics - target: Optional[PresetBenchmarkTarget] = None - client: Optional[PresetBenchmarkClient] = None @property def effective_output_tok_per_s(self) -> float: @@ -77,13 +77,6 @@ def effective_output_tok_per_s(self) -> float: def effective_per_user_tok_per_s(self) -> float: return 1000 / self.metrics.tpot_ms.p50 - @field_validator("tool", "tool_version", "command") - @classmethod - def validate_non_empty(cls, value: str) -> str: - if not value.strip(): - raise ValueError("value must be non-empty") - return value - @field_validator("command") @classmethod def validate_command_has_no_bearer_token(cls, value: str) -> str: @@ -100,73 +93,61 @@ def validate_command_has_no_bearer_token(cls, value: str) -> str: @model_validator(mode="after") def validate_metrics(self) -> Self: - metrics = self.metrics - workload = self.workload - assert metrics is not None and workload is not None - if metrics.failed_requests != 0: + if self.metrics.failed_requests != 0: raise ValueError("benchmark must not include failed requests") - if metrics.successful_requests != workload.num_requests: + if self.metrics.successful_requests != self.workload.num_requests: raise ValueError("benchmark request count must match workload.num_requests") return self -class PresetValidationReplica(CoreModel): - resources: list[ResourcesSpec] - - -class PresetValidation(CoreModel): - replicas: list[PresetValidationReplica] - benchmark: PresetBenchmark +class PresetVerificationReplicaGroup(CoreModel): + # The service replica group this was measured for. + name: str + # One entry per replica that was running: its actual resources. + replicas: list[ResourcesSpec] class Preset(CoreModel): - base: str + """What was asked (`configuration`), what to deploy (`service`), and + the evidence it works (`verification_data`).""" + id: str name: Optional[str] = None - model: str + configuration: PresetConfiguration + base: Annotated[str, Field(min_length=1)] + model: Annotated[str, Field(min_length=1)] + # The largest context the service was verified to serve. context_length: PositiveInt - trial: Optional[PositiveInt] = None - min_context_length: Optional[PositiveInt] = None - max_ttft: Optional[PositiveInt] = None - created_at: datetime + # The session's `trials/` that won verification and became this preset. + best_trial: PositiveInt + submitted_at: datetime + # The service that passed verification, stripped of this machine's deployment + # choices; `apply` submits it with the user's own name, gateway, and profile. service: ServiceConfiguration - validations: list[PresetValidation] - - @field_validator("base", "id", "model") - @classmethod - def validate_non_empty(cls, value: str) -> str: - if not value.strip(): - raise ValueError("value must be non-empty") - return value + benchmark: PresetBenchmark + # The hardware it was verified on: the actual resources of every running + # replica, by service replica group. + verified_on: list[PresetVerificationReplicaGroup] @model_validator(mode="after") def validate_preset(self) -> Self: service = self.service - validations = self.validations - if service is None or validations is None: - return self if service.model is None: raise ValueError("preset service must specify model") if any(group.resources is None for group in service.replica_groups): raise ValueError("preset service must specify resources") - if service.name is not None or service.gateway is not None: - raise ValueError("preset service must not specify name or gateway") - if any(getattr(service, field) is not None for field in ProfileParams.model_fields): - raise ValueError("preset service must not specify placement constraints") - if not validations: - raise ValueError("preset must include validation evidence") - for validation in validations: - if len(validation.replicas) != len(service.replica_groups): - raise ValueError( - "preset validation replicas must match service replica group order" - ) - if validation.benchmark.target is None or validation.benchmark.client is None: - raise ValueError("preset benchmark must specify target and client") - for replica_group in validation.replicas: - if not replica_group.resources: - raise ValueError("preset validation replicas must specify resources") - for resources in replica_group.resources: - _validate_exact_resources(resources) + for field in PRESET_EXCLUDED_FIELDS: + if getattr(service, field) is not None: + raise ValueError(f"preset service must not specify {field}") + if [group.name for group in self.verified_on] != [ + group.name for group in service.replica_groups + ]: + raise ValueError("preset verification replica groups must match the service's") + for replica_group in self.verified_on: + if not replica_group.replicas: + raise ValueError("preset verification replica groups must not be empty") + for resources in replica_group.replicas: + _validate_exact_resources(resources) return self @@ -177,19 +158,19 @@ class PresetListOutput(CoreModel): def _validate_exact_resources(resources: ResourcesSpec) -> None: cpu = resources.cpu if not _is_exact(cpu.count) or not _is_exact(resources.memory): - raise ValueError("preset validation resources must be exact") + raise ValueError("preset verification resources must be exact") if resources.disk is None or not _is_exact(resources.disk.size): - raise ValueError("preset validation resources must be exact") + raise ValueError("preset verification resources must be exact") gpu = resources.gpu if gpu is None or not _is_exact(gpu.count): - raise ValueError("preset validation resources must be exact") + raise ValueError("preset verification resources must be exact") if gpu.count.min == 0: return if gpu.name is None or len(gpu.name) != 1 or not _is_exact(gpu.memory): - raise ValueError("preset validation resources must be exact") + raise ValueError("preset verification resources must be exact") -def _is_exact(value) -> bool: +def _is_exact(value: Optional[Range]) -> bool: return ( value is not None and value.min is not None diff --git a/src/dstack/_internal/cli/services/presets/agent.py b/src/dstack/_internal/cli/services/presets/agent.py index b28470891..f9e4a1934 100644 --- a/src/dstack/_internal/cli/services/presets/agent.py +++ b/src/dstack/_internal/cli/services/presets/agent.py @@ -9,38 +9,46 @@ from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path -from typing import Any, AsyncIterator, Callable, Optional, Sequence +from typing import Any, AsyncIterator, Callable, Literal, Optional, Sequence, get_args import psutil - -from dstack._internal.cli.models.preset_agent import AGENT_FINAL_REPORT_JSON_SCHEMA +from pydantic import ValidationError + +from dstack._internal.cli.models.preset_agent import ( + AnyClaudeStreamEvent, + ClaudeResultEvent, + PresetAgentFailure, + PresetAgentSuccess, + PresetSessionProcess, +) 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, + PresetSession, + pid_running, print_preset_progress, + process_alive, + process_started_at, ) from dstack._internal.cli.services.presets.tail import ( - _DirectoryMirror, - _FileLineReader, - _OffsetStore, - _ProgressTailer, - _RecordMirror, + DirectoryMirror, + FileLineReader, + OffsetStore, + ProgressTailer, + RecordMirror, open_session_offsets, ) from dstack._internal.cli.services.presets.workspace import ( - _PROGRESS_ENV, + PROGRESS_ENV, PresetAgentWorkspace, ) from dstack._internal.compat import IS_WINDOWS from dstack._internal.core.errors import CLIError +from dstack._internal.core.models.common import validate_json_extra_ignore 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") +ClaudeEffort = Literal["low", "medium", "high", "xhigh", "max"] _RESUME_DELAYS_SECONDS: tuple[int, ...] = (30, 60, 120) _TERMINATE_GRACE_SECONDS = 3 _RESUME_PROMPT = ( @@ -81,7 +89,8 @@ class ClaudeAuth: api_key: Optional[str] executable: str - effort: Optional[str] + # None uses the claude CLI's own default. + effort: Optional[ClaudeEffort] model: str @@ -106,9 +115,9 @@ def _get_claude_version(auth: "ClaudeAuth") -> Optional[str]: return None -def _get_claude_auth_status(auth: "ClaudeAuth") -> dict[str, Any]: +def _get_claude_auth_status(auth: "ClaudeAuth") -> str: if auth.api_key: - return {"authMethod": "api-key"} + return "api-key" try: result = subprocess.run( [auth.executable, "auth", "status", "--json"], @@ -116,12 +125,9 @@ def _get_claude_auth_status(auth: "ClaudeAuth") -> dict[str, Any]: 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"} + return result.stdout.strip() or "unknown" + except (OSError, subprocess.SubprocessError): + return "unknown" def get_claude_auth() -> ClaudeAuth: @@ -131,9 +137,9 @@ def get_claude_auth() -> ClaudeAuth: 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: + if effort is not None and effort not in get_args(ClaudeEffort): raise CLIError( - f"DSTACK_AGENT_CLAUDE_EFFORT must be one of: {', '.join(_CLAUDE_EFFORT_LEVELS)}" + f"DSTACK_AGENT_CLAUDE_EFFORT must be one of: {', '.join(get_args(ClaudeEffort))}" ) return ClaudeAuth( api_key=api_key, @@ -169,7 +175,7 @@ def build_preset_agent_env( env["DSTACK_SERVER_URL"] = api.client.base_url env["DSTACK_PROJECT"] = api.project env["DSTACK_TOKEN"] = token - env[_PROGRESS_ENV] = str(workspace.progress_path) + env[PROGRESS_ENV] = str(workspace.progress_path) for name in ["TMPDIR", "TEMP", "TMP"]: env[name] = str(workspace.temp_path) # Sandbox the agent's Claude config under the workspace home when we pass our @@ -194,13 +200,13 @@ async def run_preset_agent( workspace: PresetAgentWorkspace, auth: ClaudeAuth, redacted_values: Sequence[str], - agent_session: PresetAgentSession, + session: PresetSession, initial_resume_session_id: Optional[str] = None, ) -> PresetAgentProcessOutput: - offset_store = open_session_offsets(agent_session) + offset_store = open_session_offsets(session) async with _session_tailers( workspace=workspace, - agent_session=agent_session, + session=session, redacted_values=redacted_values, offset_store=offset_store, ): @@ -217,7 +223,7 @@ async def run_preset_agent( env=env, workspace=workspace, redacted_values=redacted_values, - agent_session=agent_session, + session=session, offset_store=offset_store, ) if output.report_data is None and returncode != 0: @@ -232,7 +238,8 @@ async def run_preset_agent( if output.made_progress: retry_delays = list(_RESUME_DELAYS_SECONDS) # Another process marked this session interrupted; don't restart it. - if agent_session.read_manifest().get("status") == "interrupted": + state = session.read_state() + if state is not None and state.status == "interrupted": return output if not retry_delays: return output @@ -246,7 +253,7 @@ async def run_preset_agent( action = "retrying" print_preset_progress( f"Agent process exited without a report; {action} in {delay}s.", - agent_session=agent_session, + session=session, ) await asyncio.sleep(delay) @@ -258,8 +265,8 @@ async def _run_claude_process( env: dict[str, str], workspace: PresetAgentWorkspace, redacted_values: Sequence[str], - agent_session: PresetAgentSession, - offset_store: _OffsetStore, + session: PresetSession, + offset_store: OffsetStore, ) -> tuple[PresetAgentProcessOutput, int]: proc: Optional[asyncio.subprocess.Process] = None try: @@ -283,8 +290,8 @@ async def _run_claude_process( # CreateProcess on Windows (WinError 87). close_fds=True, ) - agent_session.update_manifest( - agent_pid=proc.pid, agent_started_at=_process_started_at(proc.pid) + session.record_agent( + PresetSessionProcess(pid=proc.pid, started_at=process_started_at(proc.pid)) ) assert proc.stdin is not None proc.stdin.write(prompt.encode()) @@ -298,7 +305,7 @@ def agent_alive() -> bool: collect_task = asyncio.create_task( _collect_agent_output( workspace=workspace, - agent_session=agent_session, + session=session, redacted_values=redacted_values, is_alive=agent_alive, offset_store=offset_store, @@ -326,9 +333,7 @@ def agent_alive() -> bool: return output, returncode -def _build_claude_command( - *, auth: ClaudeAuth, resume_session_id: Optional[str] = None -) -> list[str]: +def _build_claude_command(*, auth: ClaudeAuth, resume_session_id: Optional[str]) -> list[str]: command = [ auth.executable, "-p", @@ -346,7 +351,7 @@ def _build_claude_command( "--model", auth.model, "--json-schema", - json.dumps(AGENT_FINAL_REPORT_JSON_SCHEMA), + json.dumps(_get_report_json_schema()), ] if auth.api_key is None: command[2:2] = ["--setting-sources", "project,local"] @@ -371,9 +376,9 @@ def _prepare_subprocess_command(command: list[str]) -> list[str]: def _write_debug_trace( - session: PresetAgentSession, + session: PresetSession, *, - stream_name: str, + stream_name: Literal["stdout", "stderr"], text: str, redacted_values: Sequence[str], ) -> None: @@ -399,36 +404,36 @@ def _write_debug_trace( async def _session_tailers( *, workspace: PresetAgentWorkspace, - agent_session: PresetAgentSession, + session: PresetSession, redacted_values: Sequence[str], - offset_store: _OffsetStore, + offset_store: OffsetStore, ) -> AsyncIterator[None]: - progress_tailer = _ProgressTailer( + progress_tailer = ProgressTailer( path=workspace.progress_path, redacted_values=redacted_values, - agent_session=agent_session, + session=session, offset_store=offset_store, ) record_mirrors = [ - _RecordMirror( + RecordMirror( source=workspace.runs_path, - target=agent_session.runs_path, + target=session.runs_path, redacted_values=redacted_values, offset_store=offset_store, offset_key="runs", - echo=agent_session.echo, + echo=session.echo, ), - _DirectoryMirror( + DirectoryMirror( source=workspace.trials_dir, - target=agent_session.trials_dir, + target=session.trials_dir, redacted_values=redacted_values, - echo=agent_session.echo, + echo=session.echo, ), - _DirectoryMirror( + DirectoryMirror( source=workspace.service_dir, - target=agent_session.service_dir, + target=session.service_dir, redacted_values=redacted_values, - echo=agent_session.echo, + echo=session.echo, ), ] tailer_tasks = [ @@ -450,40 +455,36 @@ async def _session_tailers( async def _collect_agent_output( *, workspace: PresetAgentWorkspace, - agent_session: PresetAgentSession, + session: PresetSession, redacted_values: Sequence[str], is_alive: Callable[[], bool], - offset_store: _OffsetStore, + offset_store: OffsetStore, ) -> PresetAgentProcessOutput: """Safe to run alongside a live agent or over the stream files a finished one left behind.""" stdout_output, _ = await asyncio.gather( _read_process_stream( - stream=_FileLineReader( + 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, + session=session, ), _read_process_stream( - stream=_FileLineReader( + 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, + session=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 @@ -491,24 +492,24 @@ async def attach_preset_agent( *, workspace: PresetAgentWorkspace, redacted_values: Sequence[str], - agent_session: PresetAgentSession, + session: PresetSession, ) -> PresetAgentProcessOutput: """Like `run_preset_agent`, but tails a detached agent it does not own.""" - offset_store = open_session_offsets(agent_session) + offset_store = open_session_offsets(session) async with _session_tailers( workspace=workspace, - agent_session=agent_session, + session=session, redacted_values=redacted_values, offset_store=offset_store, ): - manifest = agent_session.read_manifest() + state = session.read_state() def agent_alive() -> bool: - return _pid_alive(manifest.get("agent_pid"), manifest.get("agent_started_at")) + return state is not None and state.run is not None and process_alive(state.run.agent) return await _collect_agent_output( workspace=workspace, - agent_session=agent_session, + session=session, redacted_values=redacted_values, is_alive=agent_alive, offset_store=offset_store, @@ -517,21 +518,23 @@ def agent_alive() -> bool: async def _read_process_stream( *, - stream: "_FileLineReader", - stream_name: str, - parse_result: bool, + stream: "FileLineReader", + stream_name: Literal["stdout", "stderr"], redacted_values: Sequence[str], - agent_session: PresetAgentSession, + session: PresetSession, ) -> PresetAgentProcessOutput: + # stderr feeds the debug trace and advances the persisted offset, but only + # stdout can carry the report. + parse_result = stream_name == "stdout" output = PresetAgentProcessOutput() while True: line = await stream.readline() if not line: return output text = line.decode(errors="replace") - if agent_session.debug: + if session.debug: _write_debug_trace( - agent_session, + session, stream_name=stream_name, text=text, redacted_values=redacted_values, @@ -539,31 +542,26 @@ async def _read_process_stream( if not parse_result: continue try: - message = json.loads(text) - except json.JSONDecodeError: - continue - if not isinstance(message, dict): + event = validate_json_extra_ignore(AnyClaudeStreamEvent, text) + except ValidationError: 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": + if output.session_id is None and event.session_id: + output.session_id = event.session_id + session.record_claude_session_id(event.session_id) + if event.type == "assistant": output.made_progress = True - if message.get("type") != "result": + if not isinstance(event, ClaudeResultEvent): 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 + if event.is_error: + output.error = redact(str(event.result or "Claude failed"), redacted_values) + if event.structured_output is not None: + output.report_data = event.structured_output continue - result = message.get("result") - if isinstance(result, str): + # An agent may print the report as its final text instead of submitting + # it through `StructuredOutput`. + if isinstance(event.result, str): try: - parsed = json.loads(result) + parsed = json.loads(event.result) except json.JSONDecodeError: continue if isinstance(parsed, dict): @@ -596,20 +594,18 @@ async def _terminate_process(proc: asyncio.subprocess.Process) -> None: await proc.wait() -def terminate_agent_process(manifest: dict[str, Any]) -> None: +def terminate_agent_process(agent: Optional[PresetSessionProcess]) -> None: """Twin of `_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") - ): + if agent is None or not process_alive(agent): return + agent_pid = agent.pid 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): + if not pid_running(agent_pid): return time.sleep(0.1) with suppress(OSError): @@ -630,3 +626,22 @@ def _terminate_windows_process_tree(pid: int) -> None: with suppress(psutil.NoSuchProcess): process.kill() psutil.wait_procs(alive, timeout=3) + + +def _get_report_json_schema() -> dict[str, Any]: + """The one shape the API can enforce: a single object, no union, only + `success` required. `AnyPresetAgentResult` enforces the rest at parse.""" + success = PresetAgentSuccess.model_json_schema() + failure = PresetAgentFailure.model_json_schema() + return { + "type": "object", + "properties": { + **success["properties"], + **failure["properties"], + # Each outcome fixes its own value; only the merged shape offers both. + "success": {"type": "boolean"}, + }, + "required": ["success"], + "additionalProperties": False, + "$defs": {**success.get("$defs", {}), **failure.get("$defs", {})}, + } diff --git a/src/dstack/_internal/cli/services/presets/build.py b/src/dstack/_internal/cli/services/presets/build.py new file mode 100644 index 000000000..d4f1e855c --- /dev/null +++ b/src/dstack/_internal/cli/services/presets/build.py @@ -0,0 +1,200 @@ +from datetime import datetime +from typing import Any, Optional, TypeVar + +import gpuhunt + +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.models.presets import ( + PRESET_EXCLUDED_FIELDS, + Preset, + PresetBenchmark, + PresetVerificationReplicaGroup, +) +from dstack._internal.core.models.configurations import ServiceConfiguration +from dstack._internal.core.models.envs import Env, EnvSentinel +from dstack._internal.core.models.instances import Resources +from dstack._internal.core.models.resources import ( + CPUSpec, + DiskSpec, + GPUSpec, + Memory, + Range, + ResourcesSpec, +) +from dstack._internal.utils.gpu import detect_gpu_vendors_by_gpu_name + +ConfigurationT = TypeVar("ConfigurationT", ServiceConfiguration, PresetConfiguration) + + +def build_preset( + *, + service: ServiceConfiguration, + verification_replica_groups: list[PresetVerificationReplicaGroup], + base_model: str, + model: str, + context_length: int, + benchmark: PresetBenchmark, + configuration: PresetConfiguration, + best_trial: int, + preset_id: str, + name: Optional[str], + submitted_at: datetime, +) -> Preset: + service = _without_excluded_fields(service) + configuration = _without_excluded_fields(configuration) + configuration.env = Env() + set_service_gpu_vendor_from_verification(service, verification_replica_groups) + return Preset( + name=name, + base=base_model, + id=preset_id, + model=model, + context_length=context_length, + best_trial=best_trial, + configuration=configuration, + submitted_at=submitted_at, + service=service, + benchmark=benchmark, + verified_on=verification_replica_groups, + ) + + +def preset_to_yaml_dict(preset: Preset) -> dict[str, Any]: + """`Preset` in the plain types `yaml.safe_dump` accepts.""" + return { + **preset.model_dump(mode="json", exclude_none=True), + "service": service_configuration_to_yaml_dict(preset.service), + } + + +def service_configuration_to_yaml_dict( + configuration: ServiceConfiguration, +) -> dict[str, Any]: + """The service as a preset stores it. + + Env is rewritten as `key=value` because dumping it writes a passthrough + variable as `HF_TOKEN: {key: HF_TOKEN}`, and this file is meant to be read.""" + service = configuration.model_dump( + mode="json", + exclude={"type", *PRESET_EXCLUDED_FIELDS}, + exclude_none=True, + ) + if configuration.env: + service["env"] = [ + _env_item_to_yaml(key, value) for key, value in sorted(configuration.env.items()) + ] + return {field: value for field, value in service.items() if value not in ({}, [])} + + +def resources_spec_from_instance_resources(resources: Resources) -> ResourcesSpec: + gpu = GPUSpec(count=Range[int](min=0, max=0)) + if resources.gpus: + first = resources.gpus[0] + if any( + (g.name, g.memory_mib, g.vendor) != (first.name, first.memory_mib, first.vendor) + for g in resources.gpus + ): + raise ValueError("preset cannot be built from mixed-GPU instances") + gpu = GPUSpec( + vendor=first.vendor, + name=[first.name], + count=_exact(len(resources.gpus)), + memory=_exact_memory(first.memory_mib), + ) + return ResourcesSpec( + cpu=CPUSpec(arch=resources.cpu_arch, count=_exact(resources.cpus)), + memory=_exact_memory(resources.memory_mib), + disk=DiskSpec(size=_exact_memory(resources.disk.size_mib)), + gpu=gpu, + ) + + +def _exact(value: int) -> Range[int]: + return Range[int](min=value, max=value) + + +def _exact_memory(mib: int) -> Range[Memory]: + # Rounded to what a user would write in `resources`: 128887MiB is 125.9GB. + size = Memory(round(mib / 1024, 1)) + return Range[Memory](min=size, max=size) + + +def set_service_gpu_vendor_from_verification( + service: ServiceConfiguration, + verified_on: list[PresetVerificationReplicaGroup], +) -> None: + for group_num, group in enumerate(service.replica_groups): + resources = group.resources + if resources is None or not _requires_gpu(resources): + continue + verification_vendor = _get_verification_group_gpu_vendor(verified_on, group.name) + if verification_vendor is None or resources.gpu is None: + continue + if resources.gpu.vendor is not None and resources.gpu.vendor != verification_vendor: + raise ValueError("preset service GPU vendor does not match verification") + group_resources = _get_service_group_resources(service, group_num) + if group_resources.gpu is not None: + group_resources.gpu.vendor = verification_vendor + + +def _without_excluded_fields(configuration: ConfigurationT) -> ConfigurationT: + return configuration.model_copy(deep=True, update=dict.fromkeys(PRESET_EXCLUDED_FIELDS)) + + +def _env_item_to_yaml(key: str, value: str | EnvSentinel) -> str: + if isinstance(value, EnvSentinel): + return key + return f"{key}={value}" + + +def _get_verification_group_gpu_vendor( + verified_on: list[PresetVerificationReplicaGroup], + group_name: str, +) -> gpuhunt.AcceleratorVendor | None: + vendors = { + vendor + for group in verified_on + if group.name == group_name + for resources in group.replicas + if (vendor := _get_resources_gpu_vendor(resources)) is not None + } + if len(vendors) > 1: + raise ValueError("preset verification must not mix GPU vendors in a replica group") + return next(iter(vendors), None) + + +def _get_resources_gpu_vendor(resources: ResourcesSpec) -> gpuhunt.AcceleratorVendor | None: + gpu = resources.gpu + if gpu is None or gpu.count.min == 0: + return None + if gpu.vendor is not None: + return gpu.vendor + if not gpu.name: + return None + vendors: set[gpuhunt.AcceleratorVendor] = set() + for name in gpu.name: + vendors.update(detect_gpu_vendors_by_gpu_name(name)) + if len(vendors) > 1: + raise ValueError("preset verification must not mix GPU vendors in a replica group") + return next(iter(vendors), None) + + +def _get_service_group_resources( + service: ServiceConfiguration, + group_num: int, +) -> ResourcesSpec: + resources = ( + service.replicas[group_num].resources + if isinstance(service.replicas, list) + else service.resources + ) + if resources is None: + raise ValueError("preset service object must specify resources") + return resources + + +def _requires_gpu(resources: ResourcesSpec) -> bool: + gpu = resources.gpu + if gpu is None or gpu.count.max == 0: + return False + return gpu.count.min != 0 or gpu.count.max is not None diff --git a/src/dstack/_internal/cli/services/presets/create.py b/src/dstack/_internal/cli/services/presets/create.py index 84a0d0d16..34e455ad7 100644 --- a/src/dstack/_internal/cli/services/presets/create.py +++ b/src/dstack/_internal/cli/services/presets/create.py @@ -8,7 +8,7 @@ from contextlib import nullcontext, suppress from dataclasses import dataclass from pathlib import Path -from typing import Any, Optional, Sequence +from typing import Literal, Optional, Sequence import yaml from rich.table import Table @@ -18,8 +18,17 @@ DEFAULT_DATASET, PresetConfiguration, PresetConstraints, + PresetDatasetConstraints, + PresetRandomConstraints, +) +from dstack._internal.cli.models.preset_agent import ( + PresetAgentFailure, + PresetAgentSuccess, + PresetSessionFinalize, + PresetSessionState, + PresetSessionStatus, + PresetSessionWorkspace, ) -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, @@ -30,7 +39,7 @@ run_preset_agent, terminate_agent_process, ) -from dstack._internal.cli.services.presets.presets import preset_to_data +from dstack._internal.cli.services.presets.build import preset_to_yaml_dict from dstack._internal.cli.services.presets.prompt import get_preset_agent_system_prompt from dstack._internal.cli.services.presets.redaction import ( contains_redacted_value, @@ -39,17 +48,17 @@ redact, ) from dstack._internal.cli.services.presets.session import ( - PresetAgentSession, + PresetSession, SessionBusyError, claimed_session_name, - create_preset_agent_session, + create_preset_session, find_session_name_claims, - iter_agent_sessions, - load_agent_session, - load_attachable_agent_session, - mark_session_owner, + iter_preset_sessions, + load_attachable_session, + load_preset_session, print_preset_progress, print_session_log, + process_alive, release_session_claim, resolve_session_ref, session_process_alive, @@ -116,13 +125,13 @@ def follow_preset( Takes the exclusive finalize lock so a concurrent `logs -f` and reconcile can't both finalize the same session.""" - agent_session = load_attachable_agent_session(preset_id) - agent_session.echo = echo - lock = try_claim_session(agent_session) + session = load_attachable_session(preset_id) + session.echo = echo + lock = try_claim_session(session) if lock is None: raise SessionBusyError(f"Preset {preset_id} is being finalized by another process") try: - configuration = _load_session_configuration(agent_session) + configuration = _load_session_configuration(session) try: result = asyncio.run( _create_preset( @@ -131,39 +140,39 @@ def follow_preset( source_configuration=configuration, store=store, keep_service=keep_service, - agent_session=agent_session, - attach=True, + session=session, + mode="attach", 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) + _detach_agent_session(session) raise except AgentExitedWithoutReport as e: - _stop_active_session_runs(api, agent_session) - _suspend_agent_session(agent_session) + _stop_active_session_runs(api, session) + _suspend_agent_session(session) raise CLIError(str(e)) from e except CLIError: # A CLIError is definitive (bad report, unverifiable service, leaked # secret): the preset cannot be built, so fail it. - _close_agent_session(agent_session, "failed") + _close_agent_session(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") + _close_agent_session(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" +def _load_session_configuration(session: PresetSession) -> PresetConfiguration: + configuration_path = 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" + f"Preset {session.preset_id} has no saved configuration and cannot be" + f" followed; resume it with --resume {session.preset_id} instead" ) try: return PresetConfiguration.model_validate( @@ -181,9 +190,9 @@ def show_preset_session_logs( follow: bool, keep_service: bool, ) -> Optional[PresetCreateResult]: - session = load_agent_session(preset_id) - status = session.read_manifest().get("status") - if not follow or status in ("success", "failed", "interrupted"): + session = load_preset_session(preset_id) + state = session.read_state() + if not follow or state is None or state.status in ("success", "failed", "interrupted"): print_session_log(session) return None print_session_log(session) @@ -201,7 +210,7 @@ def show_preset_session_logs( return None -def _follow_session_log_readonly(session: PresetAgentSession) -> None: +def _follow_session_log_readonly(session: PresetSession) -> None: try: offset = session.log_path.stat().st_size except OSError: @@ -216,10 +225,10 @@ def _follow_session_log_readonly(session: PresetAgentSession) -> None: chunk = "" if chunk: console.print(Text(chunk.rstrip("\n")), soft_wrap=True) - manifest = session.read_manifest() - if manifest.get("status") in ("success", "failed", "interrupted"): + state = session.read_state() + if state is None or state.status in ("success", "failed", "interrupted"): return - if not chunk and not session_process_alive(manifest): + if not chunk and not session_process_alive(state): # The owner died without recording a terminal status; stop tailing # rather than poll forever. console.print( @@ -236,27 +245,30 @@ def reconcile_detached_sessions(store: PresetStore) -> None: 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()): + for session in iter_preset_sessions(): + if _is_reconcilable(session.read_state()): _reconcile_session(session, store) -def _is_reconcilable(manifest: dict[str, Any]) -> bool: +def _is_reconcilable(state: Optional[PresetSessionState]) -> bool: # An orphaned session (no live owner) whose agent left a completion report. # 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) + state is not None + and state.status == "running" + and state.run is not None + and session_report_exists(state) + and not session_process_alive(state) ) -def _reconcile_session(session: PresetAgentSession, store: PresetStore) -> None: - manifest = session.read_manifest() +def _reconcile_session(session: PresetSession, store: PresetStore) -> None: + state = session.read_state() + if state is None or state.run is None: + return try: - api = Client.from_config(project_name=str(manifest.get("project") or "")) + api = Client.from_config(project_name=state.run.finalize.project) except Exception: # noqa: BLE001 — offline/misconfigured must not break the read command return # follow_preset records the terminal status and is claim-safe; suppress every @@ -266,16 +278,18 @@ def _reconcile_session(session: PresetAgentSession, store: PresetStore) -> None: api=api, store=store, preset_id=session.preset_id, - keep_service=bool(manifest.get("keep_service")), + keep_service=state.run.finalize.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") + session = load_preset_session(preset_id) + state = session.read_state() + if state is None: + raise CLIError(f"Preset {preset_id} session state is unreadable") + status = state.status if status == "success": console.print(f"Preset [code]{preset_id}[/] is already created") return @@ -285,14 +299,14 @@ def stop_preset_session(api: Client, preset_id: str) -> None: 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): + if session_report_exists(state) and not session_process_alive(state): # 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")), + keep_service=state.run.finalize.keep_service if state.run else False, echo=False, ) console.print(f"Preset [code]{preset_id}[/] is already created") @@ -300,12 +314,12 @@ def stop_preset_session(api: Client, preset_id: str) -> None: # 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) + terminate_agent_process(state.run.agent if state.run else None) _stop_active_session_runs(api, session) _suspend_agent_session(session) -def _stop_active_session_runs(api: Client, session: PresetAgentSession) -> None: +def _stop_active_session_runs(api: Client, session: PresetSession) -> None: names = _load_submitted_run_names(session.runs_path) active = [] for name in names: @@ -342,24 +356,26 @@ def _resolve_preset_env( return configuration -def resolve_previous_sessions(refs: Sequence[str]) -> tuple[PresetAgentSession, ...]: - sessions: list[PresetAgentSession] = [] +def resolve_previous_sessions(refs: Sequence[str]) -> tuple[PresetSession, ...]: + sessions: list[PresetSession] = [] for ref in refs: try: - session = load_agent_session(resolve_session_ref(ref)) + session = load_preset_session(resolve_session_ref(ref)) except CLIError: raise CLIError(f"Previous session {ref!r} does not exist") if all(existing.preset_id != session.preset_id for existing in sessions): sessions.append(session) included = {session.preset_id for session in sessions} for session in sessions: - manifest = session.read_manifest() - if manifest.get("status") == "running" and session_process_alive(manifest): + state = session.read_state() + if state is None: + raise CLIError(f"Previous session {session.preset_id} state is unreadable") + if state.status == "running" and session_process_alive(state): raise CLIError( f"Previous session {session.preset_id} is still running;" " wait for it to finish or stop it" ) - for parent in manifest.get("previous") or []: + for parent in state.previous: if parent not in included: warn( f"{session.preset_id} was created with --previous {parent}," @@ -368,11 +384,11 @@ def resolve_previous_sessions(refs: Sequence[str]) -> tuple[PresetAgentSession, return tuple(sessions) -def _load_pinned_previous_sessions(ids: Sequence[str]) -> tuple[PresetAgentSession, ...]: +def _load_pinned_previous_sessions(ids: Sequence[str]) -> tuple[PresetSession, ...]: sessions = [] for preset_id in ids: try: - sessions.append(load_agent_session(preset_id)) + sessions.append(load_preset_session(preset_id)) except CLIError: warn(f"Previous session {preset_id} no longer exists; keeping its copied records") return tuple(sessions) @@ -386,12 +402,16 @@ def create_preset( keep_service: bool = False, build_name: Optional[str] = None, debug: bool = False, - resume_session: Optional[PresetAgentSession] = None, + resume_session: Optional[PresetSession] = None, user_prompt: Optional[str] = None, allowed_fleets: Optional[tuple[str, ...]] = None, - previous: Sequence[PresetAgentSession] = (), + previous: Sequence[PresetSession] = (), ) -> PresetCreateResult: - agent_session = resume_session or create_preset_agent_session(configuration, debug=debug) + session = resume_session or create_preset_session( + configuration, + previous=tuple(session.preset_id for session in previous), + debug=debug, + ) try: resolved_configuration = _resolve_preset_env(configuration) result = asyncio.run( @@ -402,23 +422,23 @@ def create_preset( store=store, keep_service=keep_service, build_name=build_name, - agent_session=agent_session, - resume=resume_session is not None, + session=session, + mode="resume" if resume_session is not None else "fresh", user_prompt=user_prompt, allowed_fleets=allowed_fleets, previous=previous, ) ) except KeyboardInterrupt: - _stop_or_detach_agent_session(agent_session, api) + _stop_or_detach_agent_session(session, api) raise except CreationStopped: # The stopping CLI already suspended the session. raise except BaseException: - _close_agent_session(agent_session, "failed") + _close_agent_session(session, "failed") raise - _close_agent_session(agent_session, "success") + _close_agent_session(session, "success") return result @@ -429,6 +449,7 @@ class _CreationSetup: auth: Optional[ClaudeAuth] workspace: PresetAgentWorkspace + workspace_record: PresetSessionWorkspace build_name: str allowed_fleets: tuple[str, ...] user_prompt: Optional[str] @@ -440,28 +461,28 @@ class _CreationSetup: def _fresh_setup( api: Client, configuration: PresetConfiguration, - agent_session: PresetAgentSession, + session: PresetSession, build_name: Optional[str], allowed_fleets: Optional[tuple[str, ...]], user_prompt: Optional[str], - previous: Sequence[PresetAgentSession] = (), + previous: Sequence[PresetSession], ) -> _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) + workspace, workspace_record = create_agent_workspace(session) previous_ids = tuple(session.preset_id for session in previous) if previous_ids: - agent_session.update_manifest(previous=list(previous_ids)) install_previous_records(workspace, previous) build_name = build_name or _get_build_name( - configuration.name, configuration.model.api_model_name, agent_session.preset_id + configuration.name, configuration.model.api_model_name, session.preset_id ) return _CreationSetup( auth=auth, workspace=workspace, + workspace_record=workspace_record, build_name=build_name, allowed_fleets=allowed_fleets, user_prompt=user_prompt, @@ -472,33 +493,32 @@ def _fresh_setup( def _resume_setup( - agent_session: PresetAgentSession, + session: PresetSession, 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() + pinned_prompt = 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 - previous_ids = tuple(manifest.get("previous") or []) + workspace, workspace_record = attach_agent_workspace(session) + state = session.read_state() + if state is None or state.run is None: + raise CLIError(f"Preset {session.preset_id} session state is unreadable") + if state.run.claude_model: + auth = dataclasses.replace(auth, model=state.run.claude_model) + initial_resume_session_id = state.run.claude_session_id + previous_ids = tuple(state.previous) if previous_ids: install_previous_records(workspace, _load_pinned_previous_sessions(previous_ids)) return _CreationSetup( auth=auth, workspace=workspace, + workspace_record=workspace_record, build_name=build_name or _load_build_name(workspace), allowed_fleets=(), user_prompt=user_prompt, @@ -509,19 +529,20 @@ def _resume_setup( def _attach_setup( - agent_session: PresetAgentSession, + session: PresetSession, build_name: Optional[str], ) -> _CreationSetup: - workspace = attach_agent_workspace(agent_session) + workspace, workspace_record = attach_agent_workspace(session) return _CreationSetup( auth=None, workspace=workspace, + workspace_record=workspace_record, build_name=build_name or _load_build_name(workspace), allowed_fleets=(), user_prompt=None, initial_resume_session_id=None, write_constraints=False, - previous=tuple(agent_session.read_manifest().get("previous") or []), + previous=_read_previous_ids(session), ) @@ -530,38 +551,35 @@ async def _create_preset( api: Client, configuration: PresetConfiguration, store: PresetStore, - source_configuration: Optional[PresetConfiguration] = None, + source_configuration: PresetConfiguration, keep_service: bool = False, build_name: Optional[str] = None, - agent_session: PresetAgentSession, - resume: bool = False, - attach: bool = False, + session: PresetSession, + mode: Literal["fresh", "resume", "attach"] = "fresh", wait_for_run_stop: bool = True, user_prompt: Optional[str] = None, allowed_fleets: Optional[tuple[str, ...]] = None, - previous: Sequence[PresetAgentSession] = (), + previous: Sequence[PresetSession] = (), ) -> 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) + if mode == "attach": + setup = _attach_setup(session, build_name) + elif mode == "resume": + setup = _resume_setup(session, build_name, user_prompt) else: setup = _fresh_setup( - api, configuration, agent_session, build_name, allowed_fleets, user_prompt, previous + api, configuration, session, build_name, allowed_fleets, user_prompt, previous ) - # 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, + # Take ownership and record the run whole: workspace, finalize context, and + # the model pin, so a later detached reconcile or resume works from disk alone. + session.begin_run( + workspace=setup.workspace_record, + finalize=PresetSessionFinalize(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: + token = api.client.token + if not token: raise CLIError("The configured dstack client has no authentication token") redacted_values = get_redacted_values( [ @@ -572,7 +590,7 @@ async def _create_preset( ] ) env: dict[str, str] = {} - report: Optional[AgentFinalReport] = None + report: Optional[PresetAgentSuccess] = None preset: Optional[Preset] = None preset_path: Optional[Path] = None creation_succeeded = False @@ -589,12 +607,12 @@ async def _create_preset( prompt = get_preset_agent_system_prompt( user_prompt=setup.user_prompt, baseline=configuration.effective_baseline, - previous=", ".join(setup.previous) if setup.previous else None, + previous=setup.previous, custom_dataset=configuration.effective_dataset != DEFAULT_DATASET, ) if setup.write_constraints: if setup.user_prompt: - agent_session.write_user_prompt(setup.user_prompt) + session.write_user_prompt(setup.user_prompt) constraints_text = _build_constraints( configuration=configuration, build_name=setup.build_name, @@ -603,19 +621,19 @@ async def _create_preset( setup.workspace.constraints_path.write_text(constraints_text, encoding="utf-8") # A second, persistent copy: the workspace above is deleted with the run, # while the listing and `--previous` read constraints from the session dir. - agent_session.write_constraints(constraints_text) - if agent_session.debug: - agent_session.write_prompt(prompt) + session.write_constraints(constraints_text) + if session.debug: + session.write_prompt(prompt) if setup.auth is not None: - agent_session.write_agent_info(setup.auth) + session.write_agent_info(setup.auth) try: - if attach: + if mode == "attach": process_output = await attach_preset_agent( workspace=setup.workspace, redacted_values=redacted_values, - agent_session=agent_session, + session=session, ) - _raise_if_externally_stopped(agent_session, process_output) + _raise_if_externally_stopped(session, process_output) if ( process_output.report_data is None and not setup.workspace.final_report_path.exists() @@ -629,26 +647,30 @@ async def _create_preset( workspace=setup.workspace, auth=setup.auth, redacted_values=redacted_values, - agent_session=agent_session, + session=session, initial_resume_session_id=setup.initial_resume_session_id, ) - _raise_if_externally_stopped(agent_session, process_output) - report = load_preset_agent_report( + _raise_if_externally_stopped(session, process_output) + result = load_preset_agent_report( output=process_output, workspace=setup.workspace, redacted_values=redacted_values, ) + if isinstance(result, PresetAgentFailure): + raise CLIError(result.failure_summary or "Claude did not create a preset") + report = result run = api.client.runs.get(api.project, report.run_name) preset = build_verified_preset( run=run, preset_configuration=source_configuration, report=report, workspace_path=setup.workspace.path, - session_path=agent_session.path, - preset_id=agent_session.preset_id or None, - name=claimed_session_name(agent_session.read_manifest()), + session_path=session.path, + preset_id=session.preset_id, + name=_read_claimed_name(session), + submitted_at=session.created_at, ) - if contains_redacted_value(preset_to_data(preset), redacted_values): + if contains_redacted_value(preset_to_yaml_dict(preset), redacted_values): raise CLIError("Generated preset contains a secret value") preset_path = store.save(preset) creation_succeeded = True @@ -656,10 +678,10 @@ async def _create_preset( interrupted = True raise finally: - if agent_session.debug: + if session.debug: _save_final_report_copy( workspace=setup.workspace, - agent_session=agent_session, + session=session, redacted_values=redacted_values, ) if not interrupted: @@ -671,7 +693,7 @@ async def _create_preset( 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, + session=session, wait_for_stop=wait_for_run_stop, ) except Exception as e: @@ -680,7 +702,7 @@ async def _create_preset( if cleanup_error is not None: # The preset is already saved; a failed cleanup only means trial runs may # still be running. Warn rather than fail — else a blip discards the work. - if agent_session.echo: + if session.echo: warn(f"Failed to stop preset creation runs: {cleanup_error}") assert preset is not None assert preset_path is not None @@ -695,16 +717,27 @@ async def _create_preset( ) +def _read_previous_ids(session: PresetSession) -> tuple[str, ...]: + state = session.read_state() + return tuple(state.previous) if state is not None else () + + +def _read_claimed_name(session: PresetSession) -> Optional[str]: + state = session.read_state() + return claimed_session_name(state) if state is not None else None + + def _raise_if_externally_stopped( - session: PresetAgentSession, output: "PresetAgentProcessOutput" + session: PresetSession, output: "PresetAgentProcessOutput" ) -> None: - if output.report_data is None and session.read_manifest().get("status") == "interrupted": + state = session.read_state() + if output.report_data is None and state is not None and state.status == "interrupted": raise CreationStopped def _finish_agent_session( - session: PresetAgentSession, - status: str, + session: PresetSession, + status: PresetSessionStatus, ) -> None: try: session.finish(status) @@ -713,24 +746,22 @@ def _finish_agent_session( warn(f"Could not finalize agent output. Files remain at {session.path}: {e}") -def _close_agent_session(session: PresetAgentSession, status: str) -> None: +def _close_agent_session(session: PresetSession, status: Literal["success", "failed"]) -> None: _finish_agent_session(session, status) remove_agent_workspace(session) -def _detach_agent_session(session: PresetAgentSession) -> None: +def _detach_agent_session(session: PresetSession) -> None: """Releases ownership but leaves the agent running (still reconcilable in `dstack preset`), and stays silent since `logs -f` calls this on Ctrl+C.""" - session.update_manifest(pid=None) + session.detach() -def _stop_or_detach_agent_session( - session: PresetAgentSession, api: Optional[Client] = None -) -> None: +def _stop_or_detach_agent_session(session: PresetSession, api: Client) -> None: """`create` interrupt: stop the session, or detach and leave the agent working as a running session in `dstack preset`.""" - manifest = session.read_manifest() - agent_alive = session_process_alive({**manifest, "pid": None}) + state = session.read_state() + agent_alive = state is not None and state.run is not None and process_alive(state.run.agent) stop = True if agent_alive: try: @@ -744,13 +775,13 @@ def _stop_or_detach_agent_session( 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) + if state is not None: + terminate_agent_process(state.run.agent if state.run else None) + _stop_active_session_runs(api, session) _suspend_agent_session(session) -def _suspend_agent_session(session: PresetAgentSession) -> None: +def _suspend_agent_session(session: PresetSession) -> None: try: session.finish("interrupted") except OSError as e: @@ -797,7 +828,7 @@ class PresetNameHolders: name: str preset: Optional[Preset] - sessions: list[PresetAgentSession] + sessions: list[PresetSession] @property def preset_ids(self) -> list[str]: @@ -820,7 +851,7 @@ def reassign_preset_name(store: PresetStore, holders: PresetNameHolders) -> None if holders.preset is not None: store.release_name(holders.name) for session in holders.sessions: - session.update_manifest(name=None) + session.release_name() def plan_preset(*, api: Client, configuration: PresetConfiguration) -> tuple[str, ...]: @@ -883,42 +914,48 @@ def _build_constraints( allowed_fleets: Sequence[str], ) -> str: dataset = configuration.effective_dataset - constraints = PresetConstraints.model_validate( - { - "run_name_prefix": build_name, - "model": json.loads(configuration.model.model_dump_json(exclude_none=True)), - "min_context_length": configuration.min_context_length, - "max_ttft": configuration.max_ttft, - "trials_num": configuration.trials, - "concurrency": configuration.concurrency, - **( - { - "input_tokens": configuration.effective_input_tokens, - "output_tokens": configuration.effective_output_tokens, - "shared_prefix_tokens": configuration.shared_prefix_tokens or 0, - } - if dataset == DEFAULT_DATASET - else {"dataset": dataset} - ), - "baseline": configuration.effective_baseline, - "fleets": list(allowed_fleets), - "env": list(configuration.env), - } - ) - return json.dumps(json.loads(constraints.model_dump_json(exclude_none=True)), indent=2) + "\n" + if dataset == DEFAULT_DATASET: + constraints: PresetConstraints = PresetRandomConstraints( + run_name_prefix=build_name, + model=configuration.model, + min_context_length=configuration.min_context_length, + max_ttft=configuration.max_ttft, + trials_num=configuration.trials, + concurrency=configuration.concurrency, + input_tokens=configuration.effective_input_tokens, + output_tokens=configuration.effective_output_tokens, + shared_prefix_tokens=configuration.shared_prefix_tokens or 0, + baseline=configuration.effective_baseline, + fleets=list(allowed_fleets), + env=list(configuration.env), + ) + else: + constraints = PresetDatasetConstraints( + run_name_prefix=build_name, + model=configuration.model, + min_context_length=configuration.min_context_length, + max_ttft=configuration.max_ttft, + trials_num=configuration.trials, + concurrency=configuration.concurrency, + dataset=dataset, + baseline=configuration.effective_baseline, + fleets=list(allowed_fleets), + env=list(configuration.env), + ) + return constraints.model_dump_json(exclude_none=True, indent=2) + "\n" def _save_final_report_copy( *, workspace: PresetAgentWorkspace, - agent_session: PresetAgentSession, + session: PresetSession, 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)) + session.write_final_report(redact(report_text, redacted_values)) except OSError as e: warn(f"Could not save a final report copy: {e}") @@ -929,7 +966,7 @@ async def _cleanup_runs( build_name: str, workspace: PresetAgentWorkspace, final_run_name: Optional[str], - agent_session: PresetAgentSession, + session: PresetSession, keep_final_service: bool = False, wait_for_stop: bool = True, ) -> None: @@ -956,7 +993,7 @@ async def _cleanup_runs( deadline = asyncio.get_running_loop().time() + _RUN_STOP_TIMEOUT_SECONDS pending = set(active_names) # Without a spinner the CLI looks hung while the runs terminate. - spinner = console.status("Stopping runs...") if agent_session.echo else nullcontext() + spinner = console.status("Stopping runs...") if session.echo else nullcontext() with spinner: while pending: if asyncio.get_running_loop().time() >= deadline: @@ -967,8 +1004,8 @@ async def _cleanup_runs( 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) + if session.debug: + print_preset_progress("All preset creation runs stopped.", session=session) def _load_submitted_run_names(path: Path) -> list[str]: diff --git a/src/dstack/_internal/cli/services/presets/output.py b/src/dstack/_internal/cli/services/presets/output.py index e316dc4b1..4f8a1023f 100644 --- a/src/dstack/_internal/cli/services/presets/output.py +++ b/src/dstack/_internal/cli/services/presets/output.py @@ -4,6 +4,7 @@ from rich.table import Table +from dstack._internal.cli.models.configurations import DEFAULT_DATASET from dstack._internal.cli.models.presets import ( Preset, ) @@ -11,7 +12,7 @@ from dstack._internal.utils.common import pretty_date, pretty_resources _STATUS_DISPLAY = { - "ready": ("verified", "grey"), + "ready": ("verified", "secondary"), "running": ("trialing", "bold sea_green3"), "verifying": ("verifying", "bold deep_sky_blue1"), "interrupted": ("interrupted", "secondary"), @@ -141,7 +142,7 @@ def get_presets_table( # so different models interleave); active-only by default, else the latest row. rows: list[tuple[str, Any, bool]] = [] for preset_list in presets_by_base.values(): - rows += [(preset.created_at.isoformat(), preset, True) for preset in preset_list] + rows += [(preset.submitted_at.isoformat(), preset, True) for preset in preset_list] for session_list in sessions_by_model.values(): rows += [ (str(session.get("created_at") or ""), session, False) for session in session_list @@ -271,15 +272,10 @@ def _add_preset( "": _format_trial_spark(creation), "CONSTRAINTS": format_preset_objective( preset, - # Fall back to the creation record for presets saved before the preset - # itself carried the requested values. - min_context_length=preset.min_context_length - or (creation or {}).get("constraints", {}).get("min_context_length"), - max_ttft=preset.max_ttft or (creation or {}).get("constraints", {}).get("max_ttft"), verbose=verbose, ), "BENCHMARK": format_preset_benchmark(preset, verbose=verbose), - "SUBMITTED": pretty_date(preset.created_at), + "SUBMITTED": pretty_date(preset.submitted_at), } if verbose and preset.model != preset.base: row["BASE"] = f"[secondary] repo={preset.model}[/]" @@ -299,41 +295,41 @@ def _add_preset( def format_preset_objective( preset: Preset, *, - min_context_length: Optional[int] = None, - max_ttft: Optional[float] = None, verbose: bool = False, ) -> str: - workload = preset.validations[0].benchmark.workload + configuration = preset.configuration + workload = preset.benchmark.workload parts = [] - if workload.dataset: - parts.append(f"data={workload.dataset}") + if configuration.effective_dataset != DEFAULT_DATASET: + parts.append(f"data={configuration.effective_dataset}") else: + input_tokens = configuration.effective_input_tokens parts.append( - f"io={_format_token_count(workload.input_tokens)}" - f"/{_format_token_count(workload.output_tokens)}" + f"io={_format_token_count(input_tokens)}" + f"/{_format_token_count(configuration.effective_output_tokens)}" ) - share = round(100 * (workload.shared_prefix_tokens or 0) / workload.input_tokens) + share = round(100 * (configuration.shared_prefix_tokens or 0) / input_tokens) parts.append(f"prefix={share}%") - parts.append(f"conc={workload.concurrency}") - # Absent for presets saved before the creation record was consulted. - if verbose and min_context_length is not None: - parts.append(f"ctx>={_format_token_count(min_context_length)}") - if verbose and max_ttft is not None: - parts.append(f"ttft<={_format_duration_ms(max_ttft)}") + parts.append(f"conc={configuration.concurrency or workload.concurrency}") + if verbose and configuration.min_context_length is not None: + parts.append(f"ctx>={_format_token_count(configuration.min_context_length)}") + if verbose and configuration.max_ttft is not None: + parts.append(f"ttft<={_format_duration_ms(configuration.max_ttft)}") return f"[secondary]{' '.join(parts)}[/]" def _breaches_constraints(preset: Preset) -> bool: - metrics = preset.validations[0].benchmark.metrics - if preset.max_ttft is not None and metrics.ttft_ms.p50 > preset.max_ttft: + configuration = preset.configuration + metrics = preset.benchmark.metrics + if configuration.max_ttft is not None and metrics.ttft_ms.p50 > configuration.max_ttft: return True - return preset.min_context_length is not None and ( - preset.context_length < preset.min_context_length + return configuration.min_context_length is not None and ( + preset.context_length < configuration.min_context_length ) def format_preset_benchmark(preset: Preset, *, verbose: bool = False) -> str: - benchmark = preset.validations[0].benchmark + benchmark = preset.benchmark metrics = benchmark.metrics parts = [ f"tok/s/user={_format_number(benchmark.effective_per_user_tok_per_s)}", @@ -372,12 +368,6 @@ 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_resources(resources, *, verbose: bool) -> str: if resources is None: return "-" diff --git a/src/dstack/_internal/cli/services/presets/presets.py b/src/dstack/_internal/cli/services/presets/presets.py deleted file mode 100644 index 2801e8e24..000000000 --- a/src/dstack/_internal/cli/services/presets/presets.py +++ /dev/null @@ -1,238 +0,0 @@ -import hashlib -import json -from typing import Any, Optional - -import gpuhunt - -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 -from dstack._internal.core.models.instances import Resources -from dstack._internal.core.models.profiles import ProfileParams -from dstack._internal.core.models.resources import ResourcesSpec -from dstack._internal.utils.common import format_mib_as_gb, get_current_datetime - - -def build_preset( - *, - service: ServiceConfiguration, - validation_replicas: list[PresetValidationReplica], - base_model: str, - model: str, - context_length: int, - benchmark: PresetBenchmark, - trial: Optional[int] = None, - min_context_length: Optional[int] = None, - max_ttft: Optional[int] = None, - preset_id: Optional[str] = None, - name: Optional[str] = None, -) -> Preset: - service = service.model_copy(deep=True) - service.name = None - service.gateway = None - for field in ProfileParams.model_fields: - setattr(service, field, None) - validation = PresetValidation( - replicas=validation_replicas, - benchmark=benchmark, - ) - set_service_gpu_vendors_from_validations(service, [validation]) - return Preset( - name=name, - base=base_model, - id=preset_id or make_preset_id(service, context_length=context_length), - model=model, - context_length=context_length, - trial=trial, - min_context_length=min_context_length, - max_ttft=max_ttft, - created_at=get_current_datetime(), - service=service, - validations=[validation], - ) - - -def make_preset_id( - service: ServiceConfiguration, - context_length: int, -) -> str: - payload = json.dumps( - { - "service": service_configuration_to_preset_data(service), - "context_length": context_length, - }, - sort_keys=True, - separators=(",", ":"), - ) - return hashlib.sha256(payload.encode()).hexdigest()[:8] - - -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, - **({"trial": preset.trial} if preset.trial is not None else {}), - **( - {"min_context_length": preset.min_context_length} - if preset.min_context_length is not None - else {} - ), - **({"max_ttft": preset.max_ttft} if preset.max_ttft is not None else {}), - "created_at": preset.created_at.isoformat(), - "service": service_configuration_to_preset_data(preset.service), - "validations": [ - json.loads(validation.model_dump_json(exclude_none=True)) - for validation in preset.validations - ], - } - - -def service_configuration_to_preset_data( - configuration: ServiceConfiguration, -) -> dict[str, Any]: - """The canonical service form used for preset identity and hashing: drops - type/name/gateway/profile fields, serializes env as sorted `key=value` - strings, and removes empty collections.""" - service_data = json.loads(configuration.model_dump_json(exclude_none=True)) - service_data.pop("type", None) - service_data.pop("name", None) - service_data.pop("gateway", None) - for field in ProfileParams.model_fields: - service_data.pop(field, None) - if configuration.env: - service_data["env"] = [ - _env_item_to_preset_data(key, value) - for key, value in sorted(configuration.env.items()) - ] - else: - service_data.pop("env", None) - for field, value in list(service_data.items()): - if value in ({}, []): - service_data.pop(field) - return service_data - - -def resources_spec_from_instance_resources(resources: Resources) -> ResourcesSpec: - data: dict[str, Any] = { - "cpu": str(resources.cpus), - "memory": format_mib_as_gb(resources.memory_mib), - "disk": format_mib_as_gb(resources.disk.size_mib), - } - if resources.cpu_arch is not None: - data["cpu"] = f"{resources.cpu_arch.value}:{resources.cpus}" - if resources.gpus: - first_gpu = resources.gpus[0] - if any( - gpu.name != first_gpu.name - or gpu.memory_mib != first_gpu.memory_mib - or gpu.vendor != first_gpu.vendor - for gpu in resources.gpus - ): - 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), - "count": len(resources.gpus), - } - if first_gpu.vendor is not None: - data["gpu"]["vendor"] = first_gpu.vendor.value - else: - data["gpu"] = 0 - return ResourcesSpec.model_validate(data) - - -def set_service_gpu_vendors_from_validations( - service: ServiceConfiguration, - validations: list[PresetValidation], -) -> None: - for group_num, group in enumerate(service.replica_groups): - resources = group.resources - if resources is None or not _requires_gpu(resources): - continue - validation_vendor = _get_validation_group_gpu_vendor(validations, group_num) - if validation_vendor is None or resources.gpu is None: - continue - if resources.gpu.vendor is not None and resources.gpu.vendor != validation_vendor: - raise ValueError("preset service GPU vendor does not match validation") - group_resources = _get_service_group_resources(service, group_num) - if group_resources.gpu is not None: - group_resources.gpu.vendor = validation_vendor - - -def _env_item_to_preset_data(key: str, value: str | EnvSentinel) -> str: - if isinstance(value, EnvSentinel): - return key - return f"{key}={value}" - - -def _get_validation_group_gpu_vendor( - validations: list[PresetValidation], - group_num: int, -) -> gpuhunt.AcceleratorVendor | None: - vendors = { - vendor - for validation in validations - for resources in validation.replicas[group_num].resources - if (vendor := _get_resources_gpu_vendor(resources)) is not None - } - if len(vendors) > 1: - raise ValueError("preset validations must not mix GPU vendors in a replica group") - return next(iter(vendors), None) - - -def _get_resources_gpu_vendor(resources: ResourcesSpec) -> gpuhunt.AcceleratorVendor | None: - gpu = resources.gpu - if gpu is None or gpu.count.min == 0: - return None - if gpu.vendor is not None: - return gpu.vendor - if not gpu.name: - return None - vendors = {_get_gpu_name_vendor(name) for name in gpu.name} - {None} - if len(vendors) > 1: - raise ValueError("preset validations must not mix GPU vendors in a replica group") - return next(iter(vendors), None) - - -def _get_gpu_name_vendor(name: str) -> gpuhunt.AcceleratorVendor | None: - known = ( - (gpuhunt.KNOWN_NVIDIA_GPUS, gpuhunt.AcceleratorVendor.NVIDIA), - (gpuhunt.KNOWN_AMD_GPUS, gpuhunt.AcceleratorVendor.AMD), - (gpuhunt.KNOWN_INTEL_ACCELERATORS, gpuhunt.AcceleratorVendor.INTEL), - (gpuhunt.KNOWN_TENSTORRENT_ACCELERATORS, gpuhunt.AcceleratorVendor.TENSTORRENT), - ) - for accelerators, vendor in known: - if any(accelerator.name.lower() == name.lower() for accelerator in accelerators): - return vendor - if name.startswith("tpu-"): - return gpuhunt.AcceleratorVendor.GOOGLE - return None - - -def _get_service_group_resources( - service: ServiceConfiguration, - group_num: int, -) -> ResourcesSpec: - resources = ( - service.replicas[group_num].resources - if isinstance(service.replicas, list) - else service.resources - ) - if resources is None: - raise ValueError("preset service object must specify resources") - return resources - - -def _requires_gpu(resources: ResourcesSpec) -> bool: - gpu = resources.gpu - if gpu is None or gpu.count.max == 0: - return False - return gpu.count.min != 0 or gpu.count.max is not None diff --git a/src/dstack/_internal/cli/services/presets/prompt.py b/src/dstack/_internal/cli/services/presets/prompt.py index b3b308079..ca19b03c1 100644 --- a/src/dstack/_internal/cli/services/presets/prompt.py +++ b/src/dstack/_internal/cli/services/presets/prompt.py @@ -1,7 +1,7 @@ import re from dataclasses import dataclass, field from pathlib import Path -from typing import Optional, Union +from typing import Optional, Sequence, Union from dstack._internal.core.errors import CLIError @@ -171,10 +171,10 @@ def _render_branch( # TODO: reintroduce a `# Resume` section in system_prompt.md once session resume # (seeded from `runs.jsonl` and the trial records) is designed. def get_preset_agent_system_prompt( - user_prompt: Optional[str] = None, - baseline: bool = False, - previous: Optional[str] = None, - custom_dataset: bool = False, + user_prompt: Optional[str], + baseline: bool, + previous: Sequence[str], + custom_dataset: bool, ) -> str: text = _SYSTEM_PROMPT_PATH.read_text(encoding="utf-8").strip() variables = { @@ -182,7 +182,7 @@ def get_preset_agent_system_prompt( # Rendered for its presence only; the directive body must not interpolate it. "baseline": "on" if baseline else None, # A comma-separated list of the previous session IDs. - "previous": previous.strip() if previous else None, + "previous": ", ".join(previous) if previous else None, # Rendered for its presence only; the dataset itself is in constraints.json. "dataset": "on" if custom_dataset else None, } diff --git a/src/dstack/_internal/cli/services/presets/session.py b/src/dstack/_internal/cli/services/presets/session.py index cdccd44ca..21fb0d482 100644 --- a/src/dstack/_internal/cli/services/presets/session.py +++ b/src/dstack/_internal/cli/services/presets/session.py @@ -10,17 +10,26 @@ from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path -from typing import TYPE_CHECKING, Any, Iterator, Optional +from typing import TYPE_CHECKING, Any, Iterator, Optional, Sequence import psutil import yaml +from pydantic import ValidationError 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.models.preset_agent import ( + PresetSessionFinalize, + PresetSessionProcess, + PresetSessionRun, + PresetSessionState, + PresetSessionStatus, + PresetSessionWorkspace, +) 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.models.common import validate_extra_ignore from dstack._internal.utils.common import get_dstack_dir if TYPE_CHECKING: @@ -29,10 +38,10 @@ _PROGRESS_FILENAME = "progress.jsonl" _RUNS_FILENAME = "runs.jsonl" -_TRIALS_DIRNAME = "trials" -_SERVICE_DIRNAME = "service" -_TRIAL_RESULT_FILENAME = "trial.json" -_VERIFICATION_RESULT_FILENAME = "verification.json" +TRIALS_DIRNAME = "trials" +SERVICE_DIRNAME = "service" +TRIAL_RESULT_FILENAME = "trial.json" +VERIFICATION_RESULT_FILENAME = "verification.json" _CONSTRAINTS_FILENAME = "constraints.json" _FINAL_REPORT_FILENAME = "final_report.json" _SESSION_FILENAME = "session.json" @@ -45,15 +54,22 @@ class SessionBusyError(CLIError): @dataclass -class PresetAgentSession: +class PresetSession: path: Path debug: bool - preset_id: str = "" + preset_id: str # Background reconcile sets this False so finalizing a detached session stays # silent on the read command; agent.log is written regardless. echo: bool = field(default=True, repr=False) _log_enabled: bool = field(default=True, init=False, repr=False) + @property + def created_at(self) -> datetime: + state = self.read_state() + if state is None: + raise CLIError(f"Unknown preset session {self.preset_id}") + return state.created_at + @property def log_path(self) -> Path: return self.path / "agent.log" @@ -68,11 +84,11 @@ def runs_path(self) -> Path: @property def trials_dir(self) -> Path: - return self.path / _TRIALS_DIRNAME + return self.path / TRIALS_DIRNAME @property def service_dir(self) -> Path: - return self.path / _SERVICE_DIRNAME + return self.path / SERVICE_DIRNAME def write_prompt(self, prompt: str) -> None: _write_private_text(self.path / "prompt.md", prompt + "\n") @@ -99,21 +115,15 @@ def write_agent_info(self, auth: "ClaudeAuth") -> None: _get_claude_version, ) - info = ClaudeAgentInfo.model_validate( - { - "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.model_dump_json()), indent=2) + "\n", - ) + # `agent.json`: a debug document written once and read by nothing, so it + # is a plain dump, not a model. + info = { + "executable": auth.executable, + "version": _get_claude_version(auth), + "model": {"name": auth.model, "effort": auth.effort or "default"}, + "auth_status": _get_claude_auth_status(auth), + } + _write_private_text(self.path / "agent.json", json.dumps(info, indent=2) + "\n") def append_log(self, line: str) -> None: if not self._log_enabled: @@ -127,35 +137,134 @@ def append_log(self, line: str) -> None: if self.echo: console.print(f"[warning]Could not write agent log {self.log_path}: {e}[/]") - def read_manifest(self) -> dict[str, Any]: + def read_state(self) -> Optional[PresetSessionState]: + """None marks a session with no readable state: never created, or corrupt.""" try: - manifest = json.loads((self.path / _SESSION_FILENAME).read_text(encoding="utf-8")) + data = json.loads((self.path / _SESSION_FILENAME).read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError): - return {} - return manifest if isinstance(manifest, dict) else {} + return None + if isinstance(data, dict) and "run" not in data and ("pid" in data or "workspace" in data): + data = _upgrade_pre_0_21_2_state(data) + try: + return validate_extra_ignore(PresetSessionState, data) + except ValidationError: + return None - 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 write_state(self, state: PresetSessionState) -> None: + _write_private_text( + self.path / _SESSION_FILENAME, + state.model_dump_json(indent=2) + "\n", + ) + + def begin_run( + self, + *, + workspace: PresetSessionWorkspace, + finalize: PresetSessionFinalize, + claude_model: Optional[str], + ) -> None: + """This CLI takes ownership and starts (or joins) the agent run. Everything + an earlier run established survives: the claude session id and model pin so + a resume finds them, and the agent process reference so following a live + detached agent keeps it alive instead of reading it as dead.""" + state = self.read_state() + if state is None: + raise CLIError(f"Preset {self.preset_id} session state is unreadable") + earlier = state.run + state.status = "running" + state.owner = _current_process() + state.run = PresetSessionRun( + workspace=workspace, + finalize=finalize, + claude_model=claude_model or (earlier.claude_model if earlier else None), + agent=earlier.agent if earlier else None, + claude_session_id=earlier.claude_session_id if earlier else None, + ) + self.write_state(state) + + def record_agent(self, agent: PresetSessionProcess) -> None: + state = self.read_state() + if state is None or state.run is None: + return + state.run.agent = agent + self.write_state(state) def record_claude_session_id(self, session_id: str) -> None: - self.update_manifest(claude_session_id=session_id) + # An unreadable state stays as it is: rewriting it would fabricate a + # session record out of one field. + state = self.read_state() + if state is None or state.run is None: + return + state.run.claude_session_id = session_id + self.write_state(state) + + def detach(self) -> None: + state = self.read_state() + if state is None: + return + state.owner = None + self.write_state(state) - def finish(self, status: str) -> Path: - self.update_manifest(status=status) + def release_name(self) -> None: + state = self.read_state() + if state is None: + return + state.name = None + self.write_state(state) + + def finish(self, status: PresetSessionStatus) -> Path: + state = self.read_state() + if state is not None: + state.status = status + self.write_state(state) return self.path +# TODO: Remove in 0.22 +def _upgrade_pre_0_21_2_state(data: dict[str, Any]) -> dict[str, Any]: + """A session file from before 0.21.2 held every field flat; the same facts now + live in `owner` and `run`. Pure regrouping for backward compatibility.""" + data = dict(data) + pid = data.pop("pid", None) + pid_started_at = data.pop("pid_started_at", None) + data["owner"] = {"pid": pid, "started_at": pid_started_at} if pid is not None else None + workspace = data.pop("workspace", None) + alias = data.pop("alias", None) + agent_pid = data.pop("agent_pid", None) + agent_started_at = data.pop("agent_started_at", None) + project = data.pop("project", None) + keep_service = data.pop("keep_service", None) + claude_model = data.pop("claude_model", None) + claude_session_id = data.pop("claude_session_id", None) + if workspace is None or project is None: + # Without the finalize context there is no run to reconcile or resume. + data["run"] = None + else: + data["run"] = { + "workspace": {"path": workspace, "alias": alias or workspace}, + "finalize": {"project": project, "keep_service": bool(keep_service)}, + "claude_model": claude_model, + "agent": ( + {"pid": agent_pid, "started_at": agent_started_at} + if agent_pid is not None + else None + ), + "claude_session_id": claude_session_id, + } + data.setdefault("previous", []) + return data + + def get_presets_dir() -> Path: return get_dstack_dir() / "presets" -def create_preset_agent_session( +def create_preset_session( configuration: PresetConfiguration, *, - debug: bool = False, -) -> PresetAgentSession: + previous: Sequence[str], + debug: bool, +) -> PresetSession: if configuration.name is None: raise CLIError("The service name is required to save agent output") parent = get_presets_dir() @@ -171,27 +280,30 @@ def create_preset_agent_session( 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), - "trials_num": configuration.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.model_dump_json(exclude_none=True)) - if configuration.env: - data["env"] = list(configuration.env) - else: - data.pop("env", None) + session = PresetSession(path=path, debug=debug, preset_id=preset_id) + session.write_state( + PresetSessionState( + id=preset_id, + name=configuration.name, + model=configuration.model.exact_repo or configuration.model.api_model_name, + trials_num=configuration.trials, + previous=list(previous), + created_at=datetime.now(timezone.utc), + debug=debug, + status="running", + owner=_current_process(), + run=None, + ) + ) + record = configuration.model_dump(mode="json", exclude_none=True) + # Env values may be secrets: the session records only the variable names, + # which read back as passthrough references resolved from the environment. + record["env"] = list(configuration.env) + if not configuration.env: + record.pop("env") _write_private_text( path / "preset.dstack.yml", - yaml.safe_dump(data, sort_keys=False), + yaml.safe_dump(record, sort_keys=False), ) if debug: _write_private_text(path / "trace.jsonl", "") @@ -199,96 +311,89 @@ def create_preset_agent_session( 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) + return session -def load_resumable_agent_session(preset_id: str) -> PresetAgentSession: +def load_resumable_session(preset_id: str) -> PresetSession: 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: + session = PresetSession(path=path, debug=False, preset_id=preset_id) + state = session.read_state() + if not path.is_dir() or state is None: raise CLIError(f"Unknown preset: {preset_id}") - status = manifest.get("status") - if status == "success": + if state.status == "success": raise CLIError(f"Preset {preset_id} is already created; nothing to resume") - if status == "failed": + if state.status == "failed": raise CLIError(f"Preset {preset_id} creation failed and cannot be resumed") - if status == "running" and session_process_alive(manifest): + if state.status == "running" and session_process_alive(state): 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"): + if state.run is None or state.run.claude_session_id is None: raise CLIError(f"Preset {preset_id} creation stopped before it started; create a new one") - session.debug = bool(manifest.get("debug")) + session.debug = state.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): +def _current_process() -> PresetSessionProcess: + return PresetSessionProcess(pid=os.getpid(), started_at=process_started_at(os.getpid())) + + +def process_alive(process: Optional[PresetSessionProcess]) -> bool: + if process is None or process.pid <= 0 or not psutil.pid_exists(process.pid): return False - if isinstance(started_at, (int, float)): - create_time = _process_started_at(pid) + if process.started_at is not None: + create_time = process_started_at(process.pid) # A recycled pid has a different start time. - if create_time is not None and abs(create_time - started_at) > 1.0: + if create_time is not None and abs(create_time - process.started_at) > 1.0: return False return True -def session_process_alive(manifest: dict[str, Any]) -> bool: +def session_process_alive(state: PresetSessionState) -> bool: """True if either a live agent (possibly detached) or a live CLI (possibly between agent retries) still owns the session.""" - if _pid_alive(manifest.get("agent_pid"), manifest.get("agent_started_at")): + if state.run is not None and process_alive(state.run.agent): return True - pid = manifest.get("pid") - if not isinstance(pid, int) or pid <= 0 or pid == os.getpid(): + if state.owner is None or state.owner.pid == os.getpid(): return False - # Guard the CLI pid with its start time: a recycled pid would otherwise read - # a dead session as still owned, falsely blocking reconcile / follow. - return _pid_alive(pid, manifest.get("pid_started_at")) + return process_alive(state.owner) -def load_attachable_agent_session(preset_id: str) -> PresetAgentSession: +def load_attachable_session(preset_id: str) -> PresetSession: 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: + session = PresetSession(path=path, debug=False, preset_id=preset_id) + state = session.read_state() + if not path.is_dir() or state is None: raise CLIError(f"Unknown preset: {preset_id}") - status = manifest.get("status") - if status == "success": + if state.status == "success": raise CLIError(f"Preset {preset_id} is already created") - if status == "failed": + if state.status == "failed": raise CLIError(f"Preset {preset_id} creation failed") - if status == "interrupted": + if state.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")) - ): + owner = state.owner + if owner is not None and owner.pid != os.getpid() and process_alive(owner): raise SessionBusyError( - f"Preset {preset_id} is already being followed by another CLI (pid {pid});" + f"Preset {preset_id} is already being followed by another CLI (pid {owner.pid});" f" stop or detach it there with Ctrl+C" ) - session.debug = bool(manifest.get("debug")) + session.debug = state.debug return session -def load_agent_session(preset_id: str) -> PresetAgentSession: +def load_preset_session(preset_id: str) -> PresetSession: 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(): + session = PresetSession(path=path, debug=False, preset_id=preset_id) + if not path.is_dir() or session.read_state() is None: raise CLIError(f"Unknown preset: {preset_id}") return session -def print_session_log(session: PresetAgentSession) -> None: +def print_session_log(session: PresetSession) -> None: try: content = session.log_path.read_text(encoding="utf-8") except OSError: @@ -299,39 +404,15 @@ def print_session_log(session: PresetAgentSession) -> None: 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: - """Beyond recording ownership, stores the finalize context a later detached - reconcile needs; `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: +def session_report_exists(state: PresetSessionState) -> bool: """True once the agent has written final_report.json, marking a detached session ready to finalize.""" - workspace = manifest.get("workspace") - if not isinstance(workspace, str) or not workspace: + if state.run is None: return False - return (Path(workspace) / "w" / _FINAL_REPORT_FILENAME).is_file() + return (Path(state.run.workspace.path) / "w" / _FINAL_REPORT_FILENAME).is_file() -def try_claim_session(session: PresetAgentSession) -> Optional[int]: +def try_claim_session(session: PresetSession) -> Optional[int]: """Takes an exclusive kernel lock so two readers can't both finalize the session; returns an fd to release via `release_session_claim`, or None if another process holds it. The kernel drops the lock if the holder dies, so @@ -377,12 +458,11 @@ def _try_lock_fd(fd: int) -> bool: return False -def claimed_session_name(manifest: dict[str, Any]) -> Optional[str]: - value = manifest.get("name") - return value if isinstance(value, str) and value else None +def claimed_session_name(state: PresetSessionState) -> Optional[str]: + return state.name or None -def iter_agent_sessions() -> Iterator[PresetAgentSession]: +def iter_preset_sessions() -> Iterator[PresetSession]: """Skips dotfiles and `models--*` HuggingFace cache dirs that share the presets directory but aren't sessions.""" root = get_presets_dir() @@ -390,15 +470,15 @@ def iter_agent_sessions() -> Iterator[PresetAgentSession]: 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) + yield PresetSession(path=path, debug=False, preset_id=path.name) -def find_session_name_claims(name: str) -> list[PresetAgentSession]: +def find_session_name_claims(name: str) -> list[PresetSession]: """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 + for session in iter_preset_sessions() + if (state := session.read_state()) is not None and claimed_session_name(state) == name ] @@ -412,22 +492,22 @@ def resolve_session_ref(ref: str) -> str: return ref -def list_agent_sessions() -> list[dict[str, Any]]: +def list_preset_sessions() -> list[dict[str, Any]]: entries = [] - for session in iter_agent_sessions(): + for session in iter_preset_sessions(): path = session.path - manifest = session.read_manifest() - status = manifest.get("status") - if status not in ("running", "interrupted", "success", "failed"): + state = session.read_state() + if state is None: continue - if status == "running" and not session_process_alive(manifest): + status = state.status + if status == "running" and not session_process_alive(state): status = "interrupted" - entry = dict(manifest) + entry = state.model_dump(mode="json") entry["id"] = path.name - entry["name"] = claimed_session_name(manifest) + entry["name"] = claimed_session_name(state) entry["status"] = status - entry["trials"] = _summarize_session_trials(path / _TRIALS_DIRNAME) - entry["verification"] = _read_last_session_verification(path / _SERVICE_DIRNAME) + entry["trials"] = _summarize_session_trials(path / TRIALS_DIRNAME) + entry["verification"] = _read_last_session_verification(path / SERVICE_DIRNAME) entry["constraints"] = _read_session_constraints(path) entries.append(entry) return entries @@ -475,7 +555,7 @@ def _read_last_session_verification(path: Path) -> Optional[dict[str, Any]]: if not attempts: return None last = attempts[-1] - record = _read_record(last / _VERIFICATION_RESULT_FILENAME) + record = _read_record(last / VERIFICATION_RESULT_FILENAME) if record is not None and isinstance(record.get("status"), str): return record return {"status": "verifying"} @@ -486,7 +566,7 @@ def _summarize_session_trials(path: Path) -> Optional[dict[str, Any]]: counted.""" records = [] for trial_dir in _numbered_subdirs(path): - record = _read_record(trial_dir / _TRIAL_RESULT_FILENAME) + record = _read_record(trial_dir / TRIAL_RESULT_FILENAME) if record is not None: records.append(record) count = 0 @@ -573,11 +653,11 @@ def _format_trial_gpu(record: dict[str, Any]) -> Optional[str]: return text -def print_preset_progress(message: str, *, agent_session: PresetAgentSession) -> None: +def print_preset_progress(message: str, *, session: PresetSession) -> 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: + session.append_log(f"[{timestamp}] {message}") + if not session.echo: return console.print( Text(f"[{timestamp}]", style="log.time"), @@ -586,14 +666,14 @@ def print_preset_progress(message: str, *, agent_session: PresetAgentSession) -> ) -def _pid_running(pid: int) -> bool: +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]: +def process_started_at(pid: int) -> Optional[float]: try: return psutil.Process(pid).create_time() except psutil.Error: @@ -602,7 +682,7 @@ def _process_started_at(pid: int) -> Optional[float]: def _write_private_bytes(path: Path, content: bytes) -> None: # Atomic tmp + fsync + replace (mkstemp already creates the file 0600), so - # a crash mid-write cannot leave a truncated manifest or offsets file. + # a crash mid-write cannot leave a truncated state or offsets file. # Binary mode: mirrored files are copies, and text mode rewrites newlines. fd, temporary = tempfile.mkstemp(dir=path.parent, prefix=f".{path.name}.", suffix=".tmp") try: @@ -615,7 +695,7 @@ def _write_private_bytes(path: Path, content: bytes) -> None: except PermissionError: if not IS_WINDOWS: raise - # A concurrent reader (a viewer polling the manifest) can hold the + # A concurrent reader (a viewer polling the state) 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): diff --git a/src/dstack/_internal/cli/services/presets/store.py b/src/dstack/_internal/cli/services/presets/store.py index b837f5096..9bb5a79d2 100644 --- a/src/dstack/_internal/cli/services/presets/store.py +++ b/src/dstack/_internal/cli/services/presets/store.py @@ -1,8 +1,6 @@ import os import shutil -import sys import tempfile -from contextlib import suppress from pathlib import Path from typing import TextIO @@ -14,15 +12,21 @@ 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.models.presets import PRESET_EXCLUDED_FIELDS, Preset +from dstack._internal.cli.services.presets.build import preset_to_yaml_dict from dstack._internal.cli.utils.common import warn from dstack._internal.core.errors import CLIError, ConfigurationError +from dstack._internal.core.models.configurations import ServiceConfiguration from dstack._internal.utils.common import get_dstack_dir +class EarlierVersionPresetError(CLIError): + """The file predates dstack 0.21 and cannot be read: that format did not + record the winning trial.""" + + class PresetStore: - """One `//preset.yaml` per preset.""" + """One `//preset.yml` per preset.""" def __init__(self, root: Path | None = None) -> None: self.root = root or get_dstack_dir() / "presets" @@ -30,23 +34,38 @@ def __init__(self, root: Path | None = None) -> None: def list(self) -> list[Preset]: if not self.root.exists(): return [] - self._migrate_legacy() presets = [] - for path in self.root.glob("*/preset.yaml"): + earlier_version_ids = [] + for path in self.root.glob("*/preset.yml"): try: presets.append(self._load(path)) + except EarlierVersionPresetError: + earlier_version_ids.append(path.parent.name) 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) + if len(earlier_version_ids) == 1: + warn( + f"Preset {earlier_version_ids[0]} was created before dstack 0.21 and cannot" + f" be read. Delete it with" + f" [code]dstack preset delete {earlier_version_ids[0]}[/], or recreate.", + stderr=True, + ) + elif earlier_version_ids: + warn( + f"{len(earlier_version_ids)} presets created before dstack 0.21 cannot be" + f" read: {', '.join(sorted(earlier_version_ids))}." + f" Delete them with [code]dstack preset delete [/], or recreate.", + 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" + path = self.root / preset_id / "preset.yml" if not path.is_file(): return None preset = self._load(path) @@ -56,25 +75,16 @@ def get(self, preset_id: str) -> Preset | None: 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" - data = preset_to_data(preset) + path = directory / "preset.yml" # Undo the load-time resolution: paths under the preset directory are saved # relative, or re-saving a loaded preset (e.g. `release_name`) would bake # this machine's absolute paths back in and break portability. - for mapping in data.get("service", {}).get("files", []): - local_path = Path(mapping["local_path"]) - if not local_path.is_absolute(): - continue - for base in (directory, directory.resolve()): - try: - mapping["local_path"] = local_path.relative_to(base).as_posix() - break - except ValueError: - continue - content = yaml.safe_dump(data, sort_keys=False) + preset = preset.model_copy(deep=True) + for mapping in preset.service.files: + mapping.local_path = _relative_to_preset_dir(mapping.local_path, directory) + content = yaml.safe_dump(preset_to_yaml_dict(preset), sort_keys=False) fd, temporary_path = tempfile.mkstemp( dir=directory, prefix=f".{preset.id}.", @@ -115,90 +125,163 @@ def delete(self, preset_id: str) -> bool: _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(): + if not (directory / "preset.yml").is_file(): return False shutil.rmtree(directory) return True - 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: + data = None try: with path.open(encoding="utf-8") as f: - preset = Preset.model_validate(yaml.safe_load(f)) + data = yaml.safe_load(f) + upgraded = data + # `validations` marks the pre-0.21.2 format; the current one + # stores `verified_on`. + if isinstance(data, dict) and "validations" in data: + upgraded = _upgrade_pre_0_21_2_preset(data, preset_id=path.parent.name) + preset = Preset.model_validate(upgraded) except (OSError, ValidationError, yaml.YAMLError) as e: + if isinstance(data, dict) and "validations" in data: + raise _earlier_version_preset_error(path.parent.name) from e raise CLIError(f"Invalid preset file {path}: {e}") from e - # `files` paths are saved relative to the preset directory so the directory - # is portable; resolve them so callers see absolute paths. Absolute values - # (presets saved before this) pass through. + # Paths under the preset directory are saved relative so the directory is + # portable; resolve them so callers see absolute paths. Files outside it are + # stored absolute already and pass through. for mapping in preset.service.files: if not Path(mapping.local_path).is_absolute(): mapping.local_path = str(path.parent / mapping.local_path) return preset +# TODO: Remove in 0.22 +def _upgrade_pre_0_21_2_preset(data: dict, *, preset_id: str) -> dict: + """A preset file from before 0.21.2 stored the service as `service`, the + verification as `validations`, and echoed the request as top-level fields. + Everything the old file recorded survives the regrouping; the old format + never stored a `prompt`, so the upgraded `configuration` cannot carry one.""" + try: + # The old writer always wrote exactly one validation; its `replicas` + # entries follow the service's replica group order. + (validation,) = data["validations"] + workload = validation["benchmark"]["workload"] + configuration = {"base": data["base"]} + for field in ("min_context_length", "max_ttft"): + if data.get(field) is not None: + configuration[field] = data[field] + # The old format never stored the request's workload constraints; the + # measured workload is its record of them, and without this echo the + # preset would display today's defaults as its objective. + if workload.get("dataset", "random") == "random": + for field in ("input_tokens", "output_tokens", "shared_prefix_tokens"): + if workload.get(field) is not None: + configuration[field] = workload[field] + else: + configuration["dataset"] = workload["dataset"] + if workload.get("concurrency") is not None: + configuration["concurrency"] = workload["concurrency"] + # Fields the preset excludes (the applier's deployment choices) would + # fail validation if an old file carried them. + service = { + field: value + for field, value in data["service"].items() + if field not in PRESET_EXCLUDED_FIELDS + } + # `validations` had no group names; they come from the service itself. + group_names = [ + group.name for group in ServiceConfiguration.model_validate(service).replica_groups + ] + verified_on = [ + {"name": group_names[num], "replicas": replica["resources"]} + for num, replica in enumerate(validation["replicas"]) + ] + benchmark = { + field: value + for field, value in validation["benchmark"].items() + # The old format recorded how the benchmark client was wired up; + # the preset no longer stores that. + if field not in ("target", "client") + } + return { + "id": data["id"], + "name": data.get("name"), + "configuration": configuration, + "base": data["base"], + "model": data["model"], + "context_length": data["context_length"], + "best_trial": data["trial"], + # The old format stamped this at save time; the closest fact it + # recorded for the submission moment. + "submitted_at": data["created_at"], + "service": service, + "benchmark": benchmark, + "verified_on": verified_on, + } + except (KeyError, IndexError, TypeError, AttributeError, ValueError) as e: + # Even older than the upgradable format (e.g. no `trial` yet), or a + # file that lost part of its structure. + raise _earlier_version_preset_error(preset_id) from e + + +def _earlier_version_preset_error(preset_id: str) -> EarlierVersionPresetError: + return EarlierVersionPresetError( + f"Preset {preset_id} was created before dstack 0.21 and cannot be read." + f" Delete it with `dstack preset delete {preset_id}`, or recreate it." + ) + + +def _relative_to_preset_dir(local_path: str, directory: Path) -> str: + path = Path(local_path) + for base in (directory, directory.resolve()): + try: + return path.relative_to(base).as_posix() + except ValueError: + continue + return local_path + + 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(): +def load_preset_configuration(path: Path) -> PresetConfiguration: + if not 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) + with path.open(encoding="utf-8") as f: + return 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: +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") + # Only checked here: a stored preset serializes `model` as an object, so + # the model itself has to keep accepting the form a file may not use. + model = data.get("model") + if isinstance(model, dict) and model.get("name") is None: + key = "base" if "base" in model else "repo" + raise ConfigurationError(f"Use top-level `{key}:` instead of nested `model.{key}`") configuration = PresetConfiguration.model_validate(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: - """Prompt-file paths resolve relative to the configuration file's directory (cwd for stdin).""" +def resolve_preset_prompt(configuration: PresetConfiguration, base: Path) -> str | None: + """The prompt text; a prompt file is read relative to `base`.""" 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() diff --git a/src/dstack/_internal/cli/services/presets/tail.py b/src/dstack/_internal/cli/services/presets/tail.py index 682e4caa2..f77f74b72 100644 --- a/src/dstack/_internal/cli/services/presets/tail.py +++ b/src/dstack/_internal/cli/services/presets/tail.py @@ -9,7 +9,7 @@ from dstack._internal.cli.services.presets.redaction import redact, redact_bytes from dstack._internal.cli.services.presets.session import ( - PresetAgentSession, + PresetSession, _write_private_bytes, _write_private_text, print_preset_progress, @@ -17,7 +17,7 @@ from dstack._internal.cli.utils.common import console -class _FileLineReader: +class FileLineReader: """`readline()` over a growing file whose persisted offset lets a later attach resume exactly where the previous reader stopped.""" _POLL_SECONDS = 0.2 @@ -27,7 +27,7 @@ def __init__( self, path: Path, *, - offset_store: "_OffsetStore", + offset_store: "OffsetStore", offset_key: str, is_alive: Callable[[], bool], ) -> None: @@ -71,7 +71,7 @@ async def readline(self) -> bytes: await asyncio.sleep(self._POLL_SECONDS) -class _OffsetStore: +class OffsetStore: """Persists per-stream read offsets under disjoint keys; a thread lock suffices because the session claim guarantees no other process writes this file.""" @@ -97,30 +97,30 @@ def set(self, key: str, value: int) -> None: _write_private_text(self._path, json.dumps(self._data) + "\n") -def open_session_offsets(session: PresetAgentSession) -> _OffsetStore: - return _OffsetStore(session.path / ".offsets.json") +def open_session_offsets(session: PresetSession) -> OffsetStore: + return OffsetStore(session.path / ".offsets.json") -class _ProgressTailer: +class ProgressTailer: def __init__( self, *, path: Path, redacted_values: Sequence[str], - agent_session: PresetAgentSession, - offset_store: Optional[_OffsetStore] = None, + session: PresetSession, + offset_store: OffsetStore, offset_key: str = "progress", ) -> None: self._path = path self._redacted_values = redacted_values - self._agent_session = agent_session + self._agent_session = session self._offset_store = offset_store self._offset_key = offset_key - self._offset = offset_store.get(offset_key) if offset_store else 0 + self._offset = offset_store.get(offset_key) async def run(self) -> None: while True: - # File IO runs off the event loop; see _FileLineReader.readline. + # File IO runs off the event loop; see FileLineReader.readline. await asyncio.to_thread(self.flush) await asyncio.sleep(1) @@ -131,40 +131,40 @@ def flush(self) -> None: f.seek(self._offset) lines = f.readlines() self._offset = f.tell() - if lines and self._offset_store is not None: + if lines: 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, + session=self._agent_session, ) -class _RecordMirror: +class RecordMirror: def __init__( self, *, source: Path, target: Path, redacted_values: Sequence[str], - offset_store: Optional[_OffsetStore] = None, - offset_key: str = "", - echo: bool = True, + offset_store: OffsetStore, + offset_key: str, + echo: bool, ) -> 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._offset = offset_store.get(offset_key) self._enabled = True self._echo = echo async def run(self) -> None: while True: - # File IO runs off the event loop; see _FileLineReader.readline. + # File IO runs off the event loop; see FileLineReader.readline. await asyncio.to_thread(self.flush) await asyncio.sleep(1) @@ -180,8 +180,7 @@ def flush(self) -> None: 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) + self._offset_store.set(self._offset_key, self._offset) try: if not self._target.exists(): _write_private_text(self._target, "") @@ -197,7 +196,7 @@ def flush(self) -> None: console.print(f"[warning]Could not mirror {self._target.name}: {e}[/]") -class _DirectoryMirror: +class DirectoryMirror: """Mirrors a directory by re-copying each whole file whose size or mtime changed; a half-written file is simply recopied complete on a later flush.""" @@ -209,7 +208,7 @@ def __init__( source: Path, target: Path, redacted_values: Sequence[str], - echo: bool = True, + echo: bool, ) -> None: self._source = source self._target = target @@ -220,7 +219,7 @@ def __init__( async def run(self) -> None: while True: - # File IO runs off the event loop; see _FileLineReader.readline. + # File IO runs off the event loop; see FileLineReader.readline. await asyncio.to_thread(self.flush) await asyncio.sleep(1) diff --git a/src/dstack/_internal/cli/services/presets/verify.py b/src/dstack/_internal/cli/services/presets/verify.py index e7d56519a..1263e45fc 100644 --- a/src/dstack/_internal/cli/services/presets/verify.py +++ b/src/dstack/_internal/cli/services/presets/verify.py @@ -1,22 +1,23 @@ import json +from datetime import datetime from pathlib import Path -from typing import Any, Optional, Sequence -from urllib.parse import urlparse +from typing import Annotated, Any, Optional, Sequence -from pydantic import ValidationError +from pydantic import BeforeValidator, ConfigDict, TypeAdapter, ValidationError, ValidationInfo -from dstack._internal.cli.models.configurations import DEFAULT_DATASET, PresetConfiguration -from dstack._internal.cli.models.preset_agent import AgentFinalReport +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.models.preset_agent import ( + AnyPresetAgentResult, + PresetAgentSuccess, +) from dstack._internal.cli.models.presets import ( Preset, - PresetBenchmarkClient, - PresetBenchmarkTarget, - PresetValidationReplica, + PresetVerificationReplicaGroup, ) from dstack._internal.cli.services.presets.agent import ( PresetAgentProcessOutput, ) -from dstack._internal.cli.services.presets.presets import ( +from dstack._internal.cli.services.presets.build import ( build_preset, resources_spec_from_instance_resources, ) @@ -33,12 +34,34 @@ from dstack._internal.core.models.runs import JobStatus, Run, RunStatus +def _prepare_report(value: Any, info: ValidationInfo) -> Any: + """The parse itself redacts, so an unredacted report cannot be parsed at all — + the redaction cannot be forgotten at a call site. Explicit nulls are dropped + because the wire schema cannot forbid them on either outcome's fields.""" + if info.context is None or "redacted_values" not in info.context: + raise ValueError("an agent report must be parsed with redacted_values in context") + value = redact_structure(value, info.context["redacted_values"]) + if isinstance(value, dict): + value = {key: item for key, item in value.items() if item is not None} + return value + + +# Input hiding is defense in depth: it keeps credentials that redaction did not +# know about out of the error text shown to the user. +_AGENT_RESULT_ADAPTER = TypeAdapter( + Annotated[AnyPresetAgentResult, BeforeValidator(_prepare_report)], + config=ConfigDict(hide_input_in_errors=True), +) + + def load_preset_agent_report( *, output: PresetAgentProcessOutput, workspace: PresetAgentWorkspace, redacted_values: Sequence[str], -) -> AgentFinalReport: +) -> AnyPresetAgentResult: + """The agent's outcome, success or failure, with its secrets redacted. + Raises only when there is no valid report at all.""" report_data = output.report_data or _load_json_object(workspace.final_report_path) if report_data is None: raise CLIError( @@ -47,124 +70,126 @@ def load_preset_agent_report( redacted_values, ) ) - # Redact known secrets; an unknown leaked token is still caught downstream by - # the command bearer-token check. - report_data = redact_structure(report_data, redacted_values) try: - report = AgentFinalReport.model_validate(report_data) + return _AGENT_RESULT_ADAPTER.validate_python( + report_data, context={"redacted_values": tuple(redacted_values)} + ) except ValidationError as e: raise CLIError(f"Claude returned an invalid final report: {e}") from e - if not report.success: - raise CLIError( - redact( - report.failure_summary or "Claude did not create a preset", - redacted_values, - ) - ) - return report - - -def _rewrite_workspace_file_paths( - service: ServiceConfiguration, *, workspace_path: Path, session_path: Path -) -> None: - """Re-roots `files` onto the session's mirrored copies because the submission - workspace is deleted when the session ends; only `trials/` and `service/` are - mirrored. Paths are written relative to the preset directory so the saved - preset stays portable; the store resolves them at load.""" - workspace_root = workspace_path.resolve() - for mapping in service.files: - try: - relative = Path(mapping.local_path).resolve().relative_to(workspace_root) - except ValueError: - raise CLIError( - f"Claude final service file '{mapping.local_path}' is outside the agent workspace" - ) - target = session_path / relative - if relative.parts[:1] not in (("trials",), ("service",)) or not target.exists(): - raise CLIError( - f"Claude final service file '{mapping.local_path}' has no mirrored copy" - f" at '{target}'" - ) - mapping.local_path = relative.as_posix() def build_verified_preset( *, run: Run, preset_configuration: PresetConfiguration, - report: AgentFinalReport, - workspace_path: Optional[Path] = None, - session_path: Optional[Path] = None, - preset_id: Optional[str] = None, - name: Optional[str] = None, + report: PresetAgentSuccess, + workspace_path: Path, + session_path: Path, + preset_id: str, + name: Optional[str], + submitted_at: datetime, ) -> Preset: """Cross-checks the agent's self-reported final report against the actual run - and service state before trusting it to build a preset.""" + and service state before trusting it to build a preset. The preset's service is + taken from the run the server verified, not from anything the agent wrote, and + is then stripped of this machine's deployment choices; the session keeps the + agent's own files verbatim as the record of what ran.""" + service = _verified_run_service(run, report) + _check_report_answers_request(report, preset_configuration) + if service.model is None or service.model.name != preset_configuration.model.api_model_name: + raise CLIError("Claude final service model name does not match the requested model") + return build_preset( + name=name, + service=_portable_service( + service, + preset_configuration, + workspace_path=workspace_path, + session_path=session_path, + ), + verification_replica_groups=_get_verification_replica_groups(run, service), + base_model=report.base, + model=report.model, + context_length=report.context_length, + benchmark=report.benchmark, + best_trial=report.trial, + configuration=preset_configuration, + preset_id=preset_id, + submitted_at=submitted_at, + ) + + +def _verified_run_service(run: Run, report: PresetAgentSuccess) -> ServiceConfiguration: + """The service the server actually runs, after proving the report talks about + this run and the run is a live model service.""" 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: raise CLIError("Claude final service is not running") service = run.run_spec.configuration - if not isinstance(service, ServiceConfiguration) or service.model is None: + if not isinstance(service, ServiceConfiguration): raise CLIError("Claude final run is not a model service") - 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 - # Only when a dataset was requested: a `random` session is never told the - # field exists, so whatever it reports there means nothing. - requested_dataset = preset_configuration.effective_dataset - if requested_dataset != DEFAULT_DATASET: - if report.benchmark.workload.dataset != requested_dataset: - raise CLIError("Claude final benchmark dataset does not match the requested dataset") - if preset_configuration.model.allows_variant_selection: - if report.base != preset_configuration.model.api_model_name: + return service + + +def _check_report_answers_request( + report: PresetAgentSuccess, configuration: PresetConfiguration +) -> None: + """The report must answer what the configuration asked: the same dataset, and + the requested model — exactly when it was exact, any variant of the base + otherwise.""" + if report.benchmark.workload.dataset != configuration.effective_dataset: + raise CLIError("Claude final benchmark dataset does not match the requested dataset") + if configuration.model.allows_variant_selection: + if report.base != 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: + elif report.model != configuration.model.exact_repo: raise CLIError("Claude changed an exact model request") - target_type = ( - "gateway" if urlparse(run.service.url).scheme in {"http", "https"} else "server-proxy" - ) - benchmark = report.benchmark.model_copy( - update={ - "target": PresetBenchmarkTarget(type=target_type), - "client": PresetBenchmarkClient(type="local"), - } - ) - portable_service = service.model_copy(deep=True) - # 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 - if portable_service.files: - if workspace_path is None or session_path is None: - raise CLIError("Claude final service uses files but no workspace is attached") - _rewrite_workspace_file_paths( - portable_service, workspace_path=workspace_path, session_path=session_path + +def _portable_service( + service: ServiceConfiguration, + configuration: PresetConfiguration, + *, + workspace_path: Path, + session_path: Path, +) -> ServiceConfiguration: + """The service as the preset carries it: env values become the references the + user wrote, and workspace file paths are re-rooted onto the session's mirrored + copies, because the submission workspace is deleted when the session ends.""" + portable = service.model_copy(deep=True) + for key, value in configuration.env.items(): + if isinstance(value, EnvSentinel) and key in portable.env: + portable.env[key] = value + for mapping in portable.files: + mapping.local_path = _mirrored_file_path( + mapping.local_path, workspace_path=workspace_path, session_path=session_path ) - 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, - trial=report.trial, - min_context_length=preset_configuration.min_context_length, - max_ttft=preset_configuration.max_ttft, - preset_id=preset_id, - ) + return portable + + +def _mirrored_file_path(local_path: str, *, workspace_path: Path, session_path: Path) -> str: + """The path relative to the preset directory (the store resolves it at load), + proven to have a mirrored copy: only `trials/` and `service/` are mirrored.""" + try: + relative = Path(local_path).resolve().relative_to(workspace_path.resolve()) + except ValueError: + raise CLIError(f"Claude final service file '{local_path}' is outside the agent workspace") + if ( + relative.parts[:1] not in (("trials",), ("service",)) + or not (session_path / relative).exists() + ): + raise CLIError( + f"Claude final service file '{local_path}' has no mirrored copy" + f" at '{session_path / relative}'" + ) + return relative.as_posix() -def _get_validation_replicas( +def _get_verification_replica_groups( run: Run, service: ServiceConfiguration, -) -> list[PresetValidationReplica]: - replicas: list[PresetValidationReplica] = [] +) -> list[PresetVerificationReplicaGroup]: + groups: list[PresetVerificationReplicaGroup] = [] for group in service.replica_groups: resources = [] for job in sorted(run.jobs, key=lambda job: job.job_spec.replica_num): @@ -186,8 +211,8 @@ def _get_validation_replicas( ) if not resources: raise CLIError(f"Final service replica group {group.name!r} has no running replicas") - replicas.append(PresetValidationReplica(resources=resources)) - return replicas + groups.append(PresetVerificationReplicaGroup(name=group.name, replicas=resources)) + return groups def _load_json_object(path: Path) -> Optional[dict[str, Any]]: diff --git a/src/dstack/_internal/cli/services/presets/workspace.py b/src/dstack/_internal/cli/services/presets/workspace.py index ae1a0bfee..acb336731 100644 --- a/src/dstack/_internal/cli/services/presets/workspace.py +++ b/src/dstack/_internal/cli/services/presets/workspace.py @@ -12,21 +12,22 @@ from pathlib import Path from typing import Optional, Sequence +from dstack._internal.cli.models.preset_agent import PresetSessionWorkspace from dstack._internal.cli.services.presets.session import ( _CONSTRAINTS_FILENAME, _FINAL_REPORT_FILENAME, _PROGRESS_FILENAME, _RUNS_FILENAME, - _SERVICE_DIRNAME, - _TRIALS_DIRNAME, - PresetAgentSession, + SERVICE_DIRNAME, + TRIALS_DIRNAME, + PresetSession, ) from dstack._internal.cli.utils.common import warn 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" +PROGRESS_ENV = "DSTACK_PRESET_PROGRESS_LOG" _MAX_RUN_NAME_LENGTH = 41 _MAX_UNIX_SOCKET_PATH_BYTES = 103 @@ -54,11 +55,11 @@ def runs_path(self) -> Path: @property def trials_dir(self) -> Path: - return self.path / _TRIALS_DIRNAME + return self.path / TRIALS_DIRNAME @property def service_dir(self) -> Path: - return self.path / _SERVICE_DIRNAME + return self.path / SERVICE_DIRNAME @property def constraints_path(self) -> Path: @@ -77,7 +78,9 @@ def agent_stderr_path(self) -> Path: return self.path / ".agent-stderr.log" -def create_agent_workspace(session: PresetAgentSession) -> PresetAgentWorkspace: +def create_agent_workspace( + session: PresetSession, +) -> tuple[PresetAgentWorkspace, PresetSessionWorkspace]: real = session.path / "workspace" try: real.mkdir(mode=0o700) @@ -92,15 +95,16 @@ def create_agent_workspace(session: PresetAgentSession) -> PresetAgentWorkspace: _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 + return workspace, PresetSessionWorkspace(path=str(real), alias=str(alias)) -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: +def attach_agent_workspace( + session: PresetSession, +) -> tuple[PresetAgentWorkspace, PresetSessionWorkspace]: + state = session.read_state() + if state is None or state.run is None: raise CLIError("The preset creation has no workspace to resume") + real_value, alias_value = state.run.workspace.path, state.run.workspace.alias real, alias = Path(real_value), Path(alias_value) if not real.is_dir(): raise CLIError( @@ -109,13 +113,16 @@ def attach_agent_workspace(session: PresetAgentSession) -> PresetAgentWorkspace: ) if alias != real: _ensure_workspace_alias(alias, real) - return PresetAgentWorkspace(path=alias / "w", dstack_home=alias / "h") + workspace = PresetAgentWorkspace(path=alias / "w", dstack_home=alias / "h") + return workspace, state.run.workspace -def remove_agent_workspace(session: PresetAgentSession) -> None: - manifest = session.read_manifest() - alias = manifest.get("alias") - workspace = manifest.get("workspace") +def remove_agent_workspace(session: PresetSession) -> None: + state = session.read_state() + if state is None or state.run is None: + return + alias = state.run.workspace.alias + workspace = state.run.workspace.path if alias and alias != workspace and Path(alias).is_symlink(): with suppress(OSError): os.unlink(alias) @@ -123,11 +130,12 @@ def remove_agent_workspace(session: PresetAgentSession) -> None: shutil.rmtree(workspace, ignore_errors=True) -def scrub_workspace_token(session: PresetAgentSession) -> None: +def scrub_workspace_token(session: PresetSession) -> None: """The dstack config holds a live token; scrubbing it leaves an interrupted session with no on-disk credential, and resume re-mints it via `build_preset_agent_env`.""" - workspace = session.read_manifest().get("workspace") + state = session.read_state() + workspace = state.run.workspace.path if state is not None and state.run else None if not workspace: return with suppress(OSError): @@ -265,14 +273,14 @@ def _get_progress_script() -> str: if not message: print("Usage: progress ", file=sys.stderr) raise SystemExit(2) -path = Path(os.environ.get("{_PROGRESS_ENV}", "{_PROGRESS_FILENAME}")) +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_previous_records( - workspace: PresetAgentWorkspace, previous_sessions: Sequence[PresetAgentSession] + workspace: PresetAgentWorkspace, previous_sessions: Sequence[PresetSession] ) -> None: """Remove-then-recopy, so a crashed partial copy heals on the next run.""" for session in previous_sessions: @@ -290,8 +298,8 @@ def _copy_session_records(source_root: Path, target_root: Path) -> bool: shutil.copyfile(source_root / name, target_root / name) copied = True for group, filenames in ( - (_TRIALS_DIRNAME, ("trial.json", "task.dstack.yml")), - (_SERVICE_DIRNAME, ("service.dstack.yml", "verification.json")), + (TRIALS_DIRNAME, ("trial.json", "task.dstack.yml")), + (SERVICE_DIRNAME, ("service.dstack.yml", "verification.json")), ): source_group = source_root / group if not source_group.is_dir(): @@ -306,7 +314,7 @@ def _copy_session_records(source_root: Path, target_root: Path) -> bool: shutil.copyfile(record_dir / name, target_dir / name) copied = True patches = record_dir / "patches" - if group == _TRIALS_DIRNAME and patches.is_dir(): + if group == TRIALS_DIRNAME and patches.is_dir(): shutil.copytree(patches, target_dir / "patches", dirs_exist_ok=True) copied = True return copied diff --git a/src/dstack/api/server/__init__.py b/src/dstack/api/server/__init__.py index 8a2f2a512..4b344ed27 100644 --- a/src/dstack/api/server/__init__.py +++ b/src/dstack/api/server/__init__.py @@ -148,6 +148,10 @@ def files(self) -> FilesAPIClient: def events(self) -> EventsAPIClient: return EventsAPIClient(self._request, self._logger) + @property + def token(self) -> Optional[str]: + return self._token + def get_token_hash(self) -> str: if self._token is None: raise ValueError("Token not set") diff --git a/src/tests/_internal/cli/commands/test_preset.py b/src/tests/_internal/cli/commands/test_preset.py index 84d8a6e28..982f64a4b 100644 --- a/src/tests/_internal/cli/commands/test_preset.py +++ b/src/tests/_internal/cli/commands/test_preset.py @@ -1,3 +1,4 @@ +import argparse import json from contextlib import contextmanager from io import StringIO @@ -6,11 +7,12 @@ import pytest +from dstack._internal.cli.commands.preset import _check_stdin_configuration_confirmable from dstack._internal.cli.services.presets import output as presets_utils from dstack._internal.cli.services.presets.store import PresetStore +from dstack._internal.core.errors import CLIError from dstack._internal.utils.common import render_datetime_as_api -from tests._internal.cli.common import plain_console, run_dstack_cli -from tests._internal.cli.preset_factories import get_preset +from tests._internal.cli.common import get_preset, plain_console, run_dstack_cli pytestmark = pytest.mark.windows @@ -36,11 +38,24 @@ def mock_ssh_client_info(): yield +class TestStdinConfiguration: + def test_create_from_stdin_requires_yes(self): + # Same rule as `dstack apply`: the prompt cannot read from a stdin that + # is the configuration itself. + args = argparse.Namespace(yes=False, configuration_file="-") + with pytest.raises(CLIError, match="stdin"): + _check_stdin_configuration_confirmable(args) + + _check_stdin_configuration_confirmable( + argparse.Namespace(yes=True, configuration_file="-") + ) + + 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\ntrials: 1\nconcurrency: 8\nmax_ttft: 5000\nmin_context_length: 8192\n" + "type: preset\nname: qwen\nbase: Qwen/Qwen3.5-27B\ntrials: 1\nconcurrency: 8\nmax_ttft: 5000\nmin_context_length: 8192\n" ) with _patched_create_preset(side_effect=KeyboardInterrupt): @@ -59,7 +74,7 @@ def test_create_ends_quietly_when_stopped_from_another_cli(self, tmp_path, capsy configuration_path = tmp_path / "preset.dstack.yml" configuration_path.write_text( - "type: preset\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\ntrials: 1\nconcurrency: 8\nmax_ttft: 5000\nmin_context_length: 8192\n" + "type: preset\nname: qwen\nbase: Qwen/Qwen3.5-27B\ntrials: 1\nconcurrency: 8\nmax_ttft: 5000\nmin_context_length: 8192\n" ) with _patched_create_preset(side_effect=CreationStopped): @@ -104,7 +119,7 @@ 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) + output = self._list_output(tmp_path, ["preset", "list"], created_at=preset.submitted_at) assert "Qwen/Qwen3.5-27B" in output assert "8f3a12c4" in output @@ -133,7 +148,7 @@ def test_verbose_list_adds_repo(self, tmp_path): joined_verbose = "".join( self._list_output( - tmp_path, ["preset", "list", "-v"], created_at=preset.created_at + tmp_path, ["preset", "list", "-v"], created_at=preset.submitted_at ).split() ) @@ -169,9 +184,9 @@ def test_gets_complete_preset_as_json_without_api_client(self, tmp_path, capsys) data = json.loads(capsys.readouterr().out) assert data["id"] == preset.id - assert data["created_at"] == render_datetime_as_api(preset.created_at) + assert data["submitted_at"] == render_datetime_as_api(preset.submitted_at) assert data["context_length"] == 32768 - assert data["validations"][0]["benchmark"]["metrics"]["total_output_tokens"] == 2048 + assert data["benchmark"]["metrics"]["total_output_tokens"] == 2048 @pytest.mark.parametrize( "args", @@ -190,9 +205,9 @@ def test_lists_complete_presets_as_json(self, tmp_path, capsys, args): assert len(output["presets"]) == 1 data = output["presets"][0] assert data["id"] == preset.id - assert data["created_at"] == render_datetime_as_api(preset.created_at) + assert data["submitted_at"] == render_datetime_as_api(preset.submitted_at) assert data["context_length"] == 32768 - assert data["validations"][0]["benchmark"]["metrics"]["total_output_tokens"] == 2048 + assert data["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( @@ -246,7 +261,7 @@ def test_corrupt_preset_stays_deletable_and_keeps_json_parseable(self, tmp_path, store.save(preset) corrupt_dir = tmp_path / ".dstack" / "presets" / "deadbeef" corrupt_dir.mkdir(parents=True) - (corrupt_dir / "preset.yaml").write_text("{not valid yaml") + (corrupt_dir / "preset.yml").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 @@ -276,8 +291,7 @@ def test_merges_profile_configuration_and_cli_args(self, tmp_path): configuration_path.write_text( """type: preset name: file-name -model: - base: Qwen/Qwen3.5-27B +base: Qwen/Qwen3.5-27B regions: [file-region] max_price: 0.5 trials: 1 @@ -291,7 +305,7 @@ def test_merges_profile_configuration_and_cli_args(self, tmp_path): preset = get_preset() result = SimpleNamespace( preset=preset, - path=tmp_path / "preset.yaml", + path=tmp_path / "preset.yml", final_run_name="qwen-build-2", ) @@ -334,9 +348,7 @@ def test_apply_passes_selected_profile_and_preset_id(self, tmp_path): "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" - ) + configuration_path.write_text("type: preset\nname: qwen\nbase: Qwen/Qwen3.5-27B\n") with ( patch("dstack.api.Client.from_config"), @@ -364,9 +376,7 @@ def test_apply_passes_selected_profile_and_preset_id(self, tmp_path): 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" - ) + configuration_path.write_text("type: preset\nname: qwen\nbase: Qwen/Qwen3.5-27B\n") with patch("dstack._internal.cli.commands.preset.apply_preset") as apply: exit_code = run_dstack_cli( @@ -390,7 +400,7 @@ def test_create_detaches_the_name_from_the_old_preset(self, tmp_path): "type: preset\nname: qwen\nbase: Qwen/Qwen3.5-27B\ntrials: 1\nconcurrency: 8\nmax_ttft: 5000\nmin_context_length: 8192\n" ) result = SimpleNamespace( - preset=preset, path=tmp_path / "preset.yaml", final_run_name="qwen-1" + preset=preset, path=tmp_path / "preset.yml", final_run_name="qwen-1" ) with _patched_create_preset(return_value=result) as create: diff --git a/src/tests/_internal/cli/common.py b/src/tests/_internal/cli/common.py index 4e4442ba3..e98048d5f 100644 --- a/src/tests/_internal/cli/common.py +++ b/src/tests/_internal/cli/common.py @@ -1,13 +1,36 @@ import os +from datetime import datetime, timezone from pathlib import Path -from typing import IO, List, Optional +from types import SimpleNamespace +from typing import IO, Any, List, Optional from unittest.mock import patch +from uuid import uuid4 from rich.console import Console from rich.theme import Theme from dstack._internal.cli.main import main +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.models.preset_agent import ( + PresetAgentSuccess, + PresetSessionFinalize, + PresetSessionRun, + PresetSessionState, + PresetSessionWorkspace, +) +from dstack._internal.cli.models.presets import ( + Preset, + PresetBenchmark, + PresetVerificationReplicaGroup, +) from dstack._internal.compat import IS_WINDOWS +from dstack._internal.core.models.configurations import ( + DEFAULT_REPLICA_GROUP_NAME, + ServiceConfiguration, +) +from dstack._internal.core.models.instances import Disk, Gpu, Resources +from dstack._internal.core.models.resources import ResourcesSpec +from dstack._internal.core.models.runs import JobStatus, Run, RunStatus, ServiceSpec def plain_console(file: IO[str], *, width: int = 250) -> Console: @@ -54,3 +77,169 @@ def run_dstack_cli( if repo_dir is not None: os.chdir(cwd) return exit_code + + +def get_preset_benchmark() -> PresetBenchmark: + benchmark = PresetBenchmark( + tool="vllm bench serve", + tool_version="0.11.0", + command="vllm bench serve --base-url $SERVICE_URL", + workload={ + "api": "chat_completions", + "num_requests": 16, + "input_tokens": 1024, + "output_tokens": 128, + "concurrency": 1, + }, + metrics={ + "successful_requests": 16, + "failed_requests": 0, + "duration_seconds": 48.64, + "total_input_tokens": 16384, + "total_output_tokens": 2048, + "output_tok_per_s": 42.1, + "per_user_tok_per_s": 42.1, + "ttft_ms": {"mean": 110.9, "p50": 108.2, "p99": 121.6}, + "tpot_ms": {"mean": 7.5, "p50": 7.4, "p99": 8.1}, + }, + ) + return benchmark + + +def get_preset( + *, + preset_id: str = "8f3a12c4", + context_length: int = 32768, +) -> Preset: + resources = ResourcesSpec.model_validate( + { + "cpu": "16", + "memory": "64GB", + "disk": "200GB", + "gpu": {"name": "A6000", "memory": "48GB", "count": 1}, + } + ) + return Preset( + configuration=PresetConfiguration.model_validate( + { + "type": "preset", + "base": "Qwen/Qwen3.5-27B", + "trials": 3, + "concurrency": 1, + "input_tokens": 1024, + "output_tokens": 128, + } + ), + base="Qwen/Qwen3.5-27B", + id=preset_id, + model="community/Qwen3.5-27B-GPTQ-Int4", + submitted_at=datetime(2026, 1, 2, 3, 4, tzinfo=timezone.utc), + service=ServiceConfiguration.model_validate( + { + "image": "vllm/vllm-openai:v0.11.0", + "commands": ["vllm serve community/Qwen3.5-27B-GPTQ-Int4"], + "port": 8000, + "model": "Qwen/Qwen3.5-27B", + "resources": {"gpu": "nvidia:40GB..48GB:1"}, + "env": ["HF_TOKEN"], + } + ), + context_length=context_length, + best_trial=1, + benchmark=get_preset_benchmark(), + verified_on=[ + PresetVerificationReplicaGroup(name=DEFAULT_REPLICA_GROUP_NAME, replicas=[resources]) + ], + ) + + +def get_running_service_run() -> Run: + service = ServiceConfiguration.model_validate( + { + "name": "qwen-build-2", + "image": "vllm/vllm-openai:v0.11.0", + "commands": [ + "vllm serve community/Qwen3.5-27B-GPTQ-Int4 --served-model-name Qwen/Qwen3.5-27B" + ], + "port": 8000, + "model": "Qwen/Qwen3.5-27B", + "gateway": "benchmark-gateway", + "fleets": ["gpu-fleet"], + "backends": ["verda"], + "spot_policy": "auto", + "max_price": 0.5, + "env": {"LICENSE": "license-secret", "TOKENIZERS_PARALLELISM": "false"}, + "resources": {"gpu": "40GB..48GB:1"}, + } + ) + resources = Resources( + cpus=16, + memory_mib=64 * 1024, + gpus=[Gpu(name="A6000", memory_mib=48 * 1024)], + spot=False, + disk=Disk(size_mib=200 * 1024), + ) + job = SimpleNamespace( + job_spec=SimpleNamespace(job_num=0, replica_num=0, replica_group="0"), + job_submissions=[ + SimpleNamespace( + deployment_num=0, + status=JobStatus.RUNNING, + job_runtime_data=SimpleNamespace( + offer=SimpleNamespace(instance=SimpleNamespace(resources=resources)) + ), + ) + ], + ) + return Run.model_construct( + id=uuid4(), + project_name="main", + status=RunStatus.RUNNING, + run_spec=SimpleNamespace(run_name="qwen-build-2", configuration=service), + jobs=[job], + service=ServiceSpec(url="/proxy/services/main/qwen-build-2/"), + deployment_num=0, + ) + + +def get_session_state(**overrides: Any) -> PresetSessionState: + fields: dict[str, Any] = { + "id": "ab12cd34", + "name": None, + "model": "Qwen/Qwen3.5-27B", + "trials_num": None, + "previous": [], + "created_at": datetime(2026, 1, 2, 3, 4, tzinfo=timezone.utc), + "debug": False, + "status": "running", + "owner": None, + "run": None, + } + fields.update(overrides) + return PresetSessionState(**fields) + + +def get_session_run(**overrides: Any) -> PresetSessionRun: + fields: dict[str, Any] = { + "workspace": PresetSessionWorkspace(path="/tmp/preset-ws", alias="/tmp/preset-ws"), + "finalize": PresetSessionFinalize(project="main", keep_service=False), + "claude_model": None, + "agent": None, + "claude_session_id": None, + } + fields.update(overrides) + return PresetSessionRun(**fields) + + +def get_successful_preset_report(run: Run) -> PresetAgentSuccess: + return PresetAgentSuccess( + success=True, + run_id=run.id, + run_name=run.run_spec.run_name, + service_yaml="type: service", + trial=1, + base="Qwen/Qwen3.5-27B", + model="community/Qwen3.5-27B-GPTQ-Int4", + context_length=32768, + benchmark=get_preset_benchmark(), + ) diff --git a/src/tests/_internal/cli/models/test_configurations.py b/src/tests/_internal/cli/models/test_configurations.py index 4bec88ed1..c54045163 100644 --- a/src/tests/_internal/cli/models/test_configurations.py +++ b/src/tests/_internal/cli/models/test_configurations.py @@ -28,7 +28,7 @@ def test_parses_string_as_exact_repo(self): assert not configuration.model.allows_variant_selection def test_parses_base_model(self): - configuration = PresetConfiguration(model={"base": "Qwen/Qwen3.5-27B"}) + configuration = PresetConfiguration(base="Qwen/Qwen3.5-27B") assert isinstance(configuration.model, PresetModelBase) assert configuration.model.exact_repo is None @@ -76,7 +76,7 @@ def test_rejects_combined_base_and_repo_shorthand(self): PresetConfiguration(base="Qwen/base", repo="Qwen/repo") def test_rejects_shorthand_combined_with_model(self): - with pytest.raises(ValidationError): + with pytest.raises(ValidationError, match="cannot be combined"): PresetConfiguration(base="Qwen/base", model={"repo": "Qwen/repo"}) def test_requires_model(self): diff --git a/src/tests/_internal/cli/models/test_presets.py b/src/tests/_internal/cli/models/test_presets.py index bc4a77136..5564aabe8 100644 --- a/src/tests/_internal/cli/models/test_presets.py +++ b/src/tests/_internal/cli/models/test_presets.py @@ -1,42 +1,15 @@ 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 +from tests._internal.cli.common 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.model_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.model_fields) - # Mode-dependent fields are optional: a `random` session records - # `shared_prefix_tokens` and never hears of `dataset`; a custom-dataset - # session records `dataset` and omits `shared_prefix_tokens`. - assert set(workload_schema["required"]) == set(workload_schema["properties"]) - { - "dataset", - "shared_prefix_tokens", - } - assert set(metrics_schema["properties"]) == set(PresetBenchmarkMetrics.model_fields) - assert set(metrics_schema["required"]) == set(metrics_schema["properties"]) - assert set(metrics_schema["properties"]["ttft_ms"]["properties"]) == set( - PresetBenchmarkLatency.model_fields - ) - @pytest.mark.parametrize( ("field", "value", "error"), [ diff --git a/src/tests/_internal/cli/preset_factories.py b/src/tests/_internal/cli/preset_factories.py deleted file mode 100644 index 8a799bc8d..000000000 --- a/src/tests/_internal/cli/preset_factories.py +++ /dev/null @@ -1,150 +0,0 @@ -from datetime import datetime, timezone -from types import SimpleNamespace -from uuid import uuid4 - -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 -from dstack._internal.core.models.resources import ResourcesSpec -from dstack._internal.core.models.runs import JobStatus, Run, RunStatus, ServiceSpec - - -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", - workload={ - "api": "chat_completions", - "num_requests": 16, - "input_tokens": 1024, - "output_tokens": 128, - "concurrency": 1, - }, - metrics={ - "successful_requests": 16, - "failed_requests": 0, - "duration_seconds": 48.64, - "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}, - }, - ) - if not verified: - return benchmark - return benchmark.model_copy( - update={ - "target": PresetBenchmarkTarget(type="server-proxy"), - "client": PresetBenchmarkClient(type="local"), - } - ) - - -def get_preset( - *, - preset_id: str = "8f3a12c4", - context_length: int = 32768, -) -> Preset: - resources = ResourcesSpec.model_validate( - { - "cpu": "16", - "memory": "64GB", - "disk": "200GB", - "gpu": {"name": "A6000", "memory": "48GB", "count": 1}, - } - ) - return Preset( - base="Qwen/Qwen3.5-27B", - id=preset_id, - model="community/Qwen3.5-27B-GPTQ-Int4", - context_length=context_length, - created_at=datetime(2026, 1, 2, 3, 4, tzinfo=timezone.utc), - service=ServiceConfiguration.model_validate( - { - "image": "vllm/vllm-openai:v0.11.0", - "commands": ["vllm serve community/Qwen3.5-27B-GPTQ-Int4"], - "port": 8000, - "model": "Qwen/Qwen3.5-27B", - "resources": {"gpu": "nvidia:40GB..48GB:1"}, - "env": ["HF_TOKEN"], - } - ), - validations=[ - PresetValidation( - replicas=[PresetValidationReplica(resources=[resources])], - benchmark=get_preset_benchmark(), - ) - ], - ) - - -def get_running_service_run() -> Run: - service = ServiceConfiguration.model_validate( - { - "name": "qwen-build-2", - "image": "vllm/vllm-openai:v0.11.0", - "commands": [ - "vllm serve community/Qwen3.5-27B-GPTQ-Int4 --served-model-name Qwen/Qwen3.5-27B" - ], - "port": 8000, - "model": "Qwen/Qwen3.5-27B", - "gateway": "benchmark-gateway", - "fleets": ["gpu-fleet"], - "backends": ["verda"], - "spot_policy": "auto", - "max_price": 0.5, - "env": {"LICENSE": "license-secret", "TOKENIZERS_PARALLELISM": "false"}, - "resources": {"gpu": "40GB..48GB:1"}, - } - ) - resources = Resources( - cpus=16, - memory_mib=64 * 1024, - gpus=[Gpu(name="A6000", memory_mib=48 * 1024)], - spot=False, - disk=Disk(size_mib=200 * 1024), - ) - job = SimpleNamespace( - job_spec=SimpleNamespace(job_num=0, replica_num=0, replica_group="0"), - job_submissions=[ - SimpleNamespace( - deployment_num=0, - status=JobStatus.RUNNING, - job_runtime_data=SimpleNamespace( - offer=SimpleNamespace(instance=SimpleNamespace(resources=resources)) - ), - ) - ], - ) - return Run.model_construct( - id=uuid4(), - project_name="main", - status=RunStatus.RUNNING, - run_spec=SimpleNamespace(run_name="qwen-build-2", configuration=service), - jobs=[job], - service=ServiceSpec(url="/proxy/services/main/qwen-build-2/"), - deployment_num=0, - ) - - -def get_successful_preset_report(run: Run) -> AgentFinalReport: - return AgentFinalReport( - success=True, - run_id=run.id, - run_name=run.run_spec.run_name, - service_yaml="type: service", - trial=1, - base="Qwen/Qwen3.5-27B", - model="community/Qwen3.5-27B-GPTQ-Int4", - context_length=32768, - benchmark=get_preset_benchmark(verified=False), - ) diff --git a/src/tests/_internal/cli/services/presets/conftest.py b/src/tests/_internal/cli/services/presets/conftest.py index 843d796f0..e2dc42c28 100644 --- a/src/tests/_internal/cli/services/presets/conftest.py +++ b/src/tests/_internal/cli/services/presets/conftest.py @@ -1,6 +1,6 @@ import pytest -from dstack._internal.cli.services.presets.tail import _FileLineReader +from dstack._internal.cli.services.presets.tail import FileLineReader @pytest.fixture(autouse=True) @@ -12,4 +12,4 @@ def no_tail_poll_wait(monkeypatch: pytest.MonkeyPatch): and `_POLL_SECONDS` makes every run wait once after the agent has already exited. Tests write the whole output up front, so there is nothing to wait for. """ - monkeypatch.setattr(_FileLineReader, "_POLL_SECONDS", 0) + monkeypatch.setattr(FileLineReader, "_POLL_SECONDS", 0) diff --git a/src/tests/_internal/cli/services/presets/test_agent.py b/src/tests/_internal/cli/services/presets/test_agent.py index 22fa2bfb0..b353f1bba 100644 --- a/src/tests/_internal/cli/services/presets/test_agent.py +++ b/src/tests/_internal/cli/services/presets/test_agent.py @@ -14,6 +14,7 @@ import yaml from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.models.preset_agent import PresetSessionProcess from dstack._internal.cli.services.presets.agent import ( ClaudeAuth, _build_claude_command, @@ -28,16 +29,17 @@ redact, ) from dstack._internal.cli.services.presets.session import ( - PresetAgentSession, + PresetSession, _read_last_session_verification, _summarize_session_trials, - create_preset_agent_session, - load_resumable_agent_session, + create_preset_session, + load_resumable_session, print_preset_progress, ) from dstack._internal.cli.services.presets.tail import ( - _ProgressTailer, - _RecordMirror, + ProgressTailer, + RecordMirror, + open_session_offsets, ) from dstack._internal.cli.services.presets.workspace import ( PresetAgentWorkspace, @@ -48,6 +50,15 @@ from dstack._internal.compat import IS_WINDOWS from dstack._internal.core.errors import CLIError from dstack._internal.core.services.configs import ConfigManager +from tests._internal.cli.common import get_session_run, get_session_state + + +def _record_run(session, workspace_record): + state = session.read_state() + assert state is not None + state.run = get_session_run(workspace=workspace_record) + session.write_state(state) + pytestmark = pytest.mark.windows @@ -76,7 +87,9 @@ def test_uses_api_key_only_when_env_is_set(self, monkeypatch, 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")) + command = _build_claude_command( + auth=_claude_auth(api_key=api_key, effort="high"), resume_session_id=None + ) assert ("--bare" in command) is (api_key is not None) assert ("--setting-sources" in command) is (api_key is None) @@ -164,15 +177,17 @@ def test_detects_known_secret_in_generated_artifact(self): 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) + session = PresetSession(path=session_dir, debug=False, preset_id="abcd1234") + session.write_state(get_session_state(id="abcd1234")) + workspace, _ = create_agent_workspace(session) + return workspace class TestAgentSession: def _configuration(self) -> PresetConfiguration: return PresetConfiguration( name="qwen", - model={"base": "Qwen/Qwen3.5-27B"}, + base="Qwen/Qwen3.5-27B", max_price=0.5, env=["HF_TOKEN", "TOKENIZERS_PARALLELISM=false"], ) @@ -184,7 +199,7 @@ def _home(self, tmp_path, monkeypatch) -> None: 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()) + session = create_preset_session(self._configuration(), previous=(), debug=False) assert session.path.parent == tmp_path / ".dstack" / "presets" assert session.path.name == session.preset_id @@ -194,13 +209,13 @@ def test_creates_private_session_with_log_and_manifest(self, tmp_path, monkeypat "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) + state = json.loads((session.path / "session.json").read_text()) + assert state["id"] == session.preset_id + assert state["status"] == "running" + assert state["name"] == "qwen" + assert state["model"] == "Qwen/Qwen3.5-27B" + assert state["owner"]["pid"] == os.getpid() + print_preset_progress("creating preset", session=session) assert "creating preset" in session.log_path.read_text() assert "creating preset" in capsys.readouterr().out if not IS_WINDOWS: @@ -210,7 +225,7 @@ def test_creates_private_session_with_log_and_manifest(self, tmp_path, monkeypat 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) + debug_session = create_preset_session(self._configuration(), previous=(), debug=True) data = yaml.safe_load((debug_session.path / "preset.dstack.yml").read_text()) assert {path.name for path in debug_session.path.iterdir()} == { @@ -226,7 +241,7 @@ def test_debug_session_saves_scrubbed_configuration_and_trace(self, tmp_path, mo @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()) + session = create_preset_session(self._configuration(), previous=(), debug=False) finished_path = session.finish(status) @@ -237,15 +252,19 @@ 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( + session = PresetSession( path=session_dir, debug=False, + preset_id="ab12cd34", ) + session.write_state(get_session_state()) path = session.finish("failed") assert path == session_dir - assert json.loads((session_dir / "session.json").read_text()) == {"status": "failed"} + state = session.read_state() + assert state is not None + assert state.status == "failed" def test_reports_invalid_existing_parent(self, tmp_path, monkeypatch): monkeypatch.setenv("HOME", str(tmp_path)) @@ -254,17 +273,18 @@ def test_reports_invalid_existing_parent(self, tmp_path, monkeypatch): (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"}) + create_preset_session( + PresetConfiguration(name="qwen", base="Qwen/Qwen3.5-27B"), previous=(), debug=False ) def test_log_write_failure_warns_once(self, tmp_path, capsys): path = tmp_path / "agent-running" path.mkdir() (path / "agent.log").touch() - session = PresetAgentSession( + session = PresetSession( path=path, debug=False, + preset_id="ab12cd34", ) shutil.rmtree(path) @@ -310,9 +330,10 @@ async def test_sends_prompt_and_redacts_raw_output(self, tmp_path, monkeypatch, session_path.mkdir() (session_path / "agent.log").touch() (session_path / "trace.jsonl").touch() - agent_session = PresetAgentSession( + session = PresetSession( path=session_path, debug=True, + preset_id="ab12cd34", ) output = await run_preset_agent( prompt="full preset prompt", @@ -320,12 +341,12 @@ async def test_sends_prompt_and_redacts_raw_output(self, tmp_path, monkeypatch, workspace=workspace, auth=_claude_auth(), redacted_values=("secret-token",), - agent_session=agent_session, + session=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()] + trace = [json.loads(line) for line in session.trace_path.read_text().splitlines()] assert len(trace) == 1 assert trace[0]["timestamp"].endswith("Z") assert trace[0]["stream"] == "stdout" @@ -364,7 +385,7 @@ async def test_mirrors_trial_and_service_records_into_the_session(self, tmp_path session_path = tmp_path / "session" session_path.mkdir() (session_path / "agent.log").touch() - agent_session = PresetAgentSession(path=session_path, debug=False) + session = PresetSession(path=session_path, debug=False, preset_id="ab12cd34") output = await run_preset_agent( prompt="p", @@ -372,7 +393,7 @@ async def test_mirrors_trial_and_service_records_into_the_session(self, tmp_path workspace=workspace, auth=_claude_auth(), redacted_values=("secret-token",), - agent_session=agent_session, + session=session, ) assert output.report_data == {"ok": True} @@ -411,9 +432,10 @@ async def test_accepts_stream_event_larger_than_64_kib(self, tmp_path, monkeypat workspace=PresetAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home"), auth=_claude_auth(), redacted_values=(), - agent_session=PresetAgentSession( + session=PresetSession( path=session_path, debug=False, + preset_id="ab12cd34", ), ) @@ -426,21 +448,23 @@ def test_progress_stream_prints_only_redacted_messages(self, tmp_path, capsys): session_path = tmp_path / "agent-running" session_path.mkdir() (session_path / "agent.log").touch() - agent_session = PresetAgentSession( + session = PresetSession( path=session_path, debug=False, + preset_id="ab12cd34", ) - _ProgressTailer( + ProgressTailer( path=progress_path, redacted_values=("secret-token",), - agent_session=agent_session, + session=session, + offset_store=open_session_offsets(session), ).flush() output = capsys.readouterr().out assert "using [redacted]" in output assert "secret-token" not in output - log = agent_session.log_path.read_text() + log = session.log_path.read_text() assert "using [redacted]" in log assert "secret-token" not in log @@ -473,7 +497,14 @@ 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"]) + mirror = RecordMirror( + source=source, + target=target, + redacted_values=["dstack-secret"], + offset_store=_offsets(tmp_path), + offset_key="runs", + echo=False, + ) source.write_text( '{"name":"run-1","note":"dstack-secret"}\n{"name":"run-2"', encoding="utf-8" @@ -492,10 +523,13 @@ def test_mirrors_complete_lines_redacted(self, tmp_path): ] def test_missing_source_is_no_op(self, tmp_path): - mirror = _RecordMirror( + mirror = RecordMirror( source=tmp_path / "absent.jsonl", target=tmp_path / "target.jsonl", redacted_values=[], + offset_store=_offsets(tmp_path), + offset_key="runs", + echo=False, ) mirror.flush() @@ -505,12 +539,12 @@ def test_missing_source_is_no_op(self, tmp_path): class TestDirectoryMirror: def _mirror(self, tmp_path, **kwargs): - from dstack._internal.cli.services.presets.tail import _DirectoryMirror + from dstack._internal.cli.services.presets.tail import DirectoryMirror - return _DirectoryMirror( + return DirectoryMirror( source=tmp_path / "w" / "trials", target=tmp_path / "session" / "trials", - **{"redacted_values": ["dstack-secret"], **kwargs}, + **{"redacted_values": ["dstack-secret"], "echo": False, **kwargs}, ) def test_copies_the_tree_redacted(self, tmp_path): @@ -578,9 +612,9 @@ def test_missing_source_is_no_op(self, tmp_path): assert not (tmp_path / "session").exists() def test_skips_files_above_the_size_limit(self, tmp_path, monkeypatch): - from dstack._internal.cli.services.presets.tail import _DirectoryMirror + from dstack._internal.cli.services.presets.tail import DirectoryMirror - monkeypatch.setattr(_DirectoryMirror, "_MAX_FILE_BYTES", 8) + monkeypatch.setattr(DirectoryMirror, "_MAX_FILE_BYTES", 8) source = tmp_path / "w" / "trials" / "1" source.mkdir(parents=True) (source / "trial.json").write_text('{"learned": "far larger than eight bytes"}') @@ -599,11 +633,11 @@ def test_writes_model_params_and_auth(self, tmp_path, monkeypatch): ) monkeypatch.setattr( "dstack._internal.cli.services.presets.agent._get_claude_auth_status", - lambda auth: {"authMethod": "claude.ai", "loggedIn": True}, + lambda auth: '{"authMethod": "claude.ai", "loggedIn": true}', ) session_dir = tmp_path / "session" session_dir.mkdir() - session = PresetAgentSession(path=session_dir, debug=True) + session = PresetSession(path=session_dir, debug=True, preset_id="ab12cd34") session.write_agent_info( ClaudeAuth(api_key=None, executable="claude", effort=None, model="claude-opus-4-8") @@ -613,10 +647,16 @@ def test_writes_model_params_and_auth(self, tmp_path, monkeypatch): "executable": "claude", "version": "2.1.0 (Claude Code)", "model": {"name": "claude-opus-4-8", "effort": "default"}, - "auth": {"authMethod": "claude.ai", "loggedIn": True}, + "auth_status": '{"authMethod": "claude.ai", "loggedIn": true}', } +def _offsets(tmp_path): + session_dir = tmp_path / "offsets-session" + session_dir.mkdir(exist_ok=True) + return open_session_offsets(PresetSession(path=session_dir, debug=False, preset_id="offsets0")) + + 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 @@ -637,11 +677,13 @@ def _agent_setup(tmp_path): session_path = tmp_path / "session" session_path.mkdir() (session_path / "agent.log").touch() - agent_session = PresetAgentSession( + session = PresetSession( path=session_path, debug=False, + preset_id="ab12cd34", ) - return workspace, agent_session + session.write_state(get_session_state()) + return workspace, session def _patch_claude_command(monkeypatch, script): @@ -684,7 +726,7 @@ async def test_resumes_after_connection_error(self, tmp_path, monkeypatch, capsy "dstack._internal.cli.services.presets.agent._RESUME_DELAYS_SECONDS", (0,) ) _patch_claude_command(monkeypatch, script) - workspace, agent_session = _agent_setup(tmp_path) + workspace, session = _agent_setup(tmp_path) output = await run_preset_agent( prompt="system prompt", @@ -692,7 +734,7 @@ async def test_resumes_after_connection_error(self, tmp_path, monkeypatch, capsy workspace=workspace, auth=ClaudeAuth(api_key=None, executable="claude", effort=None, model="m"), redacted_values=(), - agent_session=agent_session, + session=session, ) assert output.report_data == {"resumed": True} @@ -726,7 +768,7 @@ async def test_resumes_on_any_unreported_death(self, tmp_path, monkeypatch): "dstack._internal.cli.services.presets.agent._RESUME_DELAYS_SECONDS", (0,) ) _patch_claude_command(monkeypatch, script) - workspace, agent_session = _agent_setup(tmp_path) + workspace, session = _agent_setup(tmp_path) output = await run_preset_agent( prompt="system prompt", @@ -734,7 +776,7 @@ async def test_resumes_on_any_unreported_death(self, tmp_path, monkeypatch): workspace=workspace, auth=ClaudeAuth(api_key=None, executable="claude", effort=None, model="m"), redacted_values=(), - agent_session=agent_session, + session=session, ) assert output.report_data == {"recovered": True} @@ -765,7 +807,7 @@ async def test_gives_up_after_repeated_no_progress_failures(self, tmp_path, monk "dstack._internal.cli.services.presets.agent._RESUME_DELAYS_SECONDS", (0, 0) ) _patch_claude_command(monkeypatch, script) - workspace, agent_session = _agent_setup(tmp_path) + workspace, session = _agent_setup(tmp_path) output = await run_preset_agent( prompt="system prompt", @@ -773,7 +815,7 @@ async def test_gives_up_after_repeated_no_progress_failures(self, tmp_path, monk workspace=workspace, auth=ClaudeAuth(api_key=None, executable="claude", effort=None, model="m"), redacted_values=(), - agent_session=agent_session, + session=session, ) assert output.report_data is None @@ -804,10 +846,12 @@ async def test_does_not_resurrect_an_externally_stopped_agent(self, tmp_path, mo "dstack._internal.cli.services.presets.agent._RESUME_DELAYS_SECONDS", (0, 0) ) _patch_claude_command(monkeypatch, script) - workspace, agent_session = _agent_setup(tmp_path) + workspace, 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") + state = session.read_state() + state.status = "interrupted" + session.write_state(state) output = await run_preset_agent( prompt="system prompt", @@ -815,7 +859,7 @@ async def test_does_not_resurrect_an_externally_stopped_agent(self, tmp_path, mo workspace=workspace, auth=ClaudeAuth(api_key=None, executable="claude", effort=None, model="m"), redacted_values=(), - agent_session=agent_session, + session=session, ) assert output.report_data is None @@ -826,21 +870,23 @@ 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") + session = PresetSession(path=session_dir, debug=False, preset_id="ab12cd34") + session.write_state(get_session_state()) + return session @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" + workspace, workspace_record = create_agent_workspace(session) + _record_run(session, workspace_record) + alias = Path(workspace_record.alias) + assert Path(workspace_record.path) == 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) + attached, _ = attach_agent_workspace(session) assert alias.is_symlink() assert (attached.path / "note.txt").read_text() == "kept" @@ -851,9 +897,9 @@ def test_create_attach_and_remove(self, tmp_path): @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"]) + _, workspace_record = create_agent_workspace(session) + _record_run(session, workspace_record) + alias = Path(workspace_record.alias) os.unlink(alias) alias.mkdir() try: @@ -864,7 +910,8 @@ def test_attach_refuses_occupied_alias(self, tmp_path): def test_attach_fails_when_workspace_is_gone(self, tmp_path): session = self._session(tmp_path) - create_agent_workspace(session) + _, workspace_record = create_agent_workspace(session) + _record_run(session, workspace_record) remove_agent_workspace(session) with pytest.raises(CLIError, match="no longer exists"): attach_agent_workspace(session) @@ -872,43 +919,89 @@ def test_attach_fails_when_workspace_is_gone(self, tmp_path): class TestOffsetPersistence: def test_mirror_does_not_duplicate_after_restart(self, tmp_path): - from dstack._internal.cli.services.presets.tail import _OffsetStore + 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( + mirror = RecordMirror( source=source, target=target, redacted_values=(), - offset_store=_OffsetStore(state), + offset_store=OffsetStore(state), offset_key="runs", + echo=False, ) mirror.flush() with source.open("a", encoding="utf-8") as f: f.write('{"name":"two"}\n') - restarted = _RecordMirror( + restarted = RecordMirror( source=source, target=target, redacted_values=(), - offset_store=_OffsetStore(state), + offset_store=OffsetStore(state), offset_key="runs", + echo=False, ) restarted.flush() assert target.read_text().splitlines() == ['{"name":"one"}', '{"name":"two"}'] +class TestOldFlatSessionState: + def test_reads_a_pre_0_22_flat_session_file(self, tmp_path): + # Verbatim shape of a session written by the pre-0.21.2 CLI: every field flat. + flat = { + "id": "30a012bf", + "status": "success", + "pid": 70516, + "pid_started_at": 1755116373.0, + "name": "dsv4-flash-mi300x-kernel2", + "model": "deepseek-ai/DeepSeek-V4-Flash", + "trials_num": 7, + "created_at": "2026-08-13T20:19:33+00:00", + "debug": False, + "agent_pid": 70521, + "agent_started_at": 1755116374.0, + "claude_model": "claude-opus-5", + "claude_session_id": "71b025f9-fba0-42b9-8734-e357deca5281", + "workspace": "/tmp/w", + "alias": "/tmp/dpe-1", + "project": "main", + "keep_service": True, + } + session_dir = tmp_path / "30a012bf" + session_dir.mkdir() + (session_dir / "session.json").write_text(json.dumps(flat)) + session = PresetSession(path=session_dir, debug=False, preset_id="30a012bf") + + state = session.read_state() + + assert state is not None + assert state.status == "success" + assert state.owner == PresetSessionProcess(pid=70516, started_at=1755116373.0) + assert state.run is not None + assert state.run.workspace.path == "/tmp/w" + assert state.run.workspace.alias == "/tmp/dpe-1" + assert state.run.finalize.project == "main" + assert state.run.finalize.keep_service is True + assert state.run.agent == PresetSessionProcess(pid=70521, started_at=1755116374.0) + assert state.run.claude_session_id == "71b025f9-fba0-42b9-8734-e357deca5281" + assert state.previous == [] + + class TestLoadResumableSession: - def _write_session(self, tmp_path, monkeypatch, manifest): + def _write_session(self, tmp_path, monkeypatch, state): monkeypatch.setenv("HOME", str(tmp_path)) monkeypatch.setenv("USERPROFILE", str(tmp_path)) - path = tmp_path / ".dstack" / "presets" / manifest["id"] + path = tmp_path / ".dstack" / "presets" / state["id"] path.mkdir(parents=True) - (path / "session.json").write_text(json.dumps(manifest), encoding="utf-8") + (path / "session.json").write_text( + get_session_state(**state).model_dump_json(), encoding="utf-8" + ) return path def test_loads_interrupted_session(self, tmp_path, monkeypatch): @@ -918,13 +1011,13 @@ def test_loads_interrupted_session(self, tmp_path, monkeypatch): { "id": "ab12cd34", "status": "interrupted", - "claude_session_id": "sid-1", + "run": get_session_run(claude_session_id="sid-1"), "debug": True, "created_at": "2026-07-20T10:00:00Z", }, ) - session = load_resumable_agent_session("ab12cd34") + session = load_resumable_session("ab12cd34") assert session.preset_id == "ab12cd34" assert session.debug is True @@ -936,8 +1029,8 @@ def test_treats_dead_running_session_as_resumable(self, tmp_path, monkeypatch): { "id": "ab12cd34", "status": "running", - "pid": 4242, - "claude_session_id": "sid-1", + "owner": {"pid": 4242, "started_at": None}, + "run": get_session_run(claude_session_id="sid-1"), }, ) monkeypatch.setattr( @@ -945,10 +1038,10 @@ def test_treats_dead_running_session_as_resumable(self, tmp_path, monkeypatch): lambda pid: False, ) - assert load_resumable_agent_session("ab12cd34").preset_id == "ab12cd34" + assert load_resumable_session("ab12cd34").preset_id == "ab12cd34" @pytest.mark.parametrize( - ("manifest", "match"), + ("state", "match"), [ pytest.param(None, "Unknown preset", id="unknown-preset"), pytest.param( @@ -959,8 +1052,8 @@ def test_treats_dead_running_session_as_resumable(self, tmp_path, monkeypatch): { "id": "aa000003", "status": "running", - "pid": 4242, - "claude_session_id": "sid-1", + "owner": {"pid": 4242, "started_at": None}, + "run": get_session_run(claude_session_id="sid-1"), }, "still being created", id="still-running", @@ -972,19 +1065,19 @@ def test_treats_dead_running_session_as_resumable(self, tmp_path, monkeypatch): ), ], ) - def test_refusals(self, tmp_path, monkeypatch, manifest, match): + def test_refusals(self, tmp_path, monkeypatch, state, 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) + if state is not None: + self._write_session(tmp_path, monkeypatch, state) 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" + preset_id = state["id"] if state is not None else "00000000" with pytest.raises(CLIError, match=match): - load_resumable_agent_session(preset_id) + load_resumable_session(preset_id) def _write_trials(tmp_path, records): @@ -1113,16 +1206,16 @@ def test_a_torn_result_copy_reads_as_verifying(self, tmp_path): 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 + 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( + reader = FileLineReader( stream, - offset_store=_OffsetStore(state), + offset_store=OffsetStore(state), offset_key="agent_stdout", is_alive=lambda: alive, ) @@ -1133,9 +1226,9 @@ async def test_reads_lines_and_continues_from_persisted_offset(self, tmp_path): with stream.open("ab") as f: f.write(b"third\npartial") alive = False - attached = _FileLineReader( + attached = FileLineReader( stream, - offset_store=_OffsetStore(state), + offset_store=OffsetStore(state), offset_key="agent_stdout", is_alive=lambda: alive, ) @@ -1153,23 +1246,28 @@ def test_detach_keeps_the_agent_and_stop_terminates_it(self, tmp_path, monkeypat session_dir = tmp_path / "ab12cd34" session_dir.mkdir() (session_dir / "agent.log").touch() - session = PresetAgentSession(path=session_dir, debug=False, preset_id="ab12cd34") + session = PresetSession(path=session_dir, debug=False, preset_id="ab12cd34") + session.write_state(get_session_state()) 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) + state = session.read_state() + assert state is not None + state.status = "running" + state.run = get_session_run(agent=PresetSessionProcess(pid=agent.pid, started_at=None)) + session.write_state(state) monkeypatch.setattr(create_module, "confirm_ask", lambda *_: False) - _stop_or_detach_agent_session(session) + _stop_or_detach_agent_session(session, None) assert agent.poll() is None - assert session.read_manifest()["status"] == "running" + assert session.read_state().status == "running" assert "Detached" in capsys.readouterr().out monkeypatch.setattr(create_module, "confirm_ask", lambda *_: True) - _stop_or_detach_agent_session(session) + _stop_or_detach_agent_session(session, None) psutil.Process(agent.pid).wait(timeout=10) - assert session.read_manifest()["status"] == "interrupted" + assert session.read_state().status == "interrupted" finally: with suppress(OSError): os.killpg(agent.pid, signal.SIGKILL) @@ -1179,10 +1277,10 @@ class TestOffsetStoreSharing: def test_shared_store_keeps_every_writers_keys(self, tmp_path): import threading - from dstack._internal.cli.services.presets.tail import _OffsetStore + from dstack._internal.cli.services.presets.tail import OffsetStore state = tmp_path / ".offsets.json" - store = _OffsetStore(state) + store = OffsetStore(state) # One store serves the whole session; readers and mirrors write # disjoint keys from worker threads. @@ -1199,6 +1297,6 @@ def advance(key: str) -> None: for thread in threads: thread.join() - reloaded = _OffsetStore(state) + 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 index d6c628042..968ed2005 100644 --- a/src/tests/_internal/cli/services/presets/test_apply.py +++ b/src/tests/_internal/cli/services/presets/test_apply.py @@ -11,7 +11,7 @@ ) from dstack._internal.core.errors import CLIError from dstack._internal.core.models.instances import InstanceAvailability -from tests._internal.cli.preset_factories import get_preset +from tests._internal.cli.common import get_preset pytestmark = pytest.mark.windows @@ -21,7 +21,7 @@ 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"}, + base="Qwen/Qwen3.5-27B", min_context_length=8192, ) @@ -33,7 +33,7 @@ def test_warns_on_insufficient_context_instead_of_failing(self, capsys): preset = get_preset(preset_id="small", context_length=4096) configuration = PresetConfiguration( name="qwen", - model={"base": "Qwen/Qwen3.5-27B"}, + base="Qwen/Qwen3.5-27B", min_context_length=8192, ) @@ -65,7 +65,7 @@ class TestBuildService: def test_applies_preset_name_env_gateway_and_constraints(self): configuration = PresetConfiguration( name="qwen-production", - model={"base": "Qwen/Qwen3.5-27B"}, + base="Qwen/Qwen3.5-27B", gateway="inference", env={"HF_TOKEN": "token"}, fleets=["gpu-fleet"], @@ -86,7 +86,7 @@ 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=PresetConfiguration(name="qwen", base="Qwen/Qwen3.5-27B"), configuration_path="preset.dstack.yml", preset_id="ee55ff66", profile_name=None, @@ -112,7 +112,7 @@ def test_applies_the_referenced_preset(self, monkeypatch): api=Mock(), configuration=PresetConfiguration( name="qwen", - model={"base": "Qwen/Qwen3.5-27B"}, + base="Qwen/Qwen3.5-27B", ), configuration_path="preset.dstack.yml", preset_id="8f3a12c4", diff --git a/src/tests/_internal/cli/services/presets/test_create.py b/src/tests/_internal/cli/services/presets/test_create.py index 92394d7a4..684ac2ece 100644 --- a/src/tests/_internal/cli/services/presets/test_create.py +++ b/src/tests/_internal/cli/services/presets/test_create.py @@ -8,6 +8,12 @@ from pydantic import ValidationError from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.models.preset_agent import ( + PresetSessionFinalize, + PresetSessionProcess, + PresetSessionState, + PresetSessionWorkspace, +) from dstack._internal.cli.services.presets.agent import ( ClaudeAuth, PresetAgentProcessOutput, @@ -30,9 +36,8 @@ stop_preset_session, ) from dstack._internal.cli.services.presets.session import ( - PresetAgentSession, - load_agent_session, - mark_session_owner, + PresetSession, + load_preset_session, print_preset_progress, print_session_log, release_session_claim, @@ -48,9 +53,11 @@ 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 ( +from tests._internal.cli.common import ( get_preset, get_running_service_run, + get_session_run, + get_session_state, get_successful_preset_report, ) @@ -82,14 +89,14 @@ def creation_context(tmp_path, monkeypatch): project="main", runs=run_apis, client=SimpleNamespace( - _token="dstack-secret", + 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"}, + base="Qwen/Qwen3.5-27B", min_context_length=8192, max_ttft=5000, concurrency=8, @@ -99,7 +106,7 @@ def creation_context(tmp_path, monkeypatch): ) source_configuration = PresetConfiguration( name="qwen-build", - model={"base": "Qwen/Qwen3.5-27B"}, + base="Qwen/Qwen3.5-27B", min_context_length=8192, max_ttft=5000, concurrency=8, @@ -132,10 +139,10 @@ def test_saves_agent_log_without_debug(self, tmp_path, monkeypatch): preset = get_preset() async def create(**kwargs): - print_preset_progress("testing preset", agent_session=kwargs["agent_session"]) + print_preset_progress("testing preset", session=kwargs["session"]) return PresetCreateResult( preset=preset, - path=tmp_path / "preset.yaml", + path=tmp_path / "preset.yml", final_run_id=uuid.uuid4(), final_run_name="qwen-build-2", ) @@ -149,7 +156,7 @@ async def create(**kwargs): api=SimpleNamespace(), configuration=PresetConfiguration( name="qwen", - model={"base": "Qwen/Qwen3.5-27B"}, + base="Qwen/Qwen3.5-27B", ), store=PresetStore(tmp_path / "presets"), ) @@ -161,9 +168,9 @@ async def create(**kwargs): "session.json", "preset.dstack.yml", } - manifest = json.loads((paths[0] / "session.json").read_text()) - assert manifest["status"] == "success" - assert manifest["id"] == paths[0].name + state = json.loads((paths[0] / "session.json").read_text()) + assert state["status"] == "success" + assert state["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): @@ -179,10 +186,10 @@ async def create(**kwargs): } assert isinstance(kwargs["source_configuration"].env["HF_TOKEN"], EnvSentinel) assert kwargs["source_configuration"].env["TOKENIZERS_PARALLELISM"] == "false" - kwargs["agent_session"].write_prompt("test prompt") + kwargs["session"].write_prompt("test prompt") return PresetCreateResult( preset=preset, - path=tmp_path / "preset.yaml", + path=tmp_path / "preset.yml", final_run_id=uuid.uuid4(), final_run_name="qwen-build-2", ) @@ -194,13 +201,13 @@ def fail_finish(self, preset_id=None): "dstack._internal.cli.services.presets.create._create_preset", create, ) - monkeypatch.setattr(PresetAgentSession, "finish", fail_finish) + monkeypatch.setattr(PresetSession, "finish", fail_finish) result = create_preset( api=SimpleNamespace(), configuration=PresetConfiguration( name="qwen", - model={"base": "Qwen/Qwen3.5-27B"}, + base="Qwen/Qwen3.5-27B", env=["HF_TOKEN", "TOKENIZERS_PARALLELISM=false"], ), store=PresetStore(tmp_path / "presets"), @@ -235,14 +242,14 @@ def fail_finish(self, preset_id=None): "dstack._internal.cli.services.presets.create._create_preset", create, ) - monkeypatch.setattr(PresetAgentSession, "finish", fail_finish) + monkeypatch.setattr(PresetSession, "finish", fail_finish) with pytest.raises(RuntimeError, match="creation failed"): create_preset( api=SimpleNamespace(), configuration=PresetConfiguration( name="qwen", - model={"base": "Qwen/Qwen3.5-27B"}, + base="Qwen/Qwen3.5-27B", ), store=PresetStore(tmp_path / "presets"), debug=True, @@ -260,14 +267,16 @@ async def test_checks_active_fleets_before_claude_auth(self, tmp_path, monkeypat ) with pytest.raises(CLIError, match="no fleets"): + configuration = PresetConfiguration( + name="qwen-build", + base="Qwen/Qwen3.5-27B", + ) await _create_preset( api=api, - configuration=PresetConfiguration( - name="qwen-build", - model={"base": "Qwen/Qwen3.5-27B"}, - ), + configuration=configuration, + source_configuration=configuration, store=PresetStore(tmp_path / "presets"), - agent_session=_agent_session(tmp_path), + session=_agent_session(tmp_path), ) @pytest.mark.asyncio @@ -295,7 +304,7 @@ async def cleanup_runs(**kwargs): configuration=creation_context.configuration, source_configuration=creation_context.source_configuration, store=creation_context.store, - agent_session=_agent_session(tmp_path), + session=_agent_session(tmp_path), ) assert cleanup_calls == [] @@ -312,13 +321,15 @@ async def test_saves_preset_and_cleans_up_runs( session_path.mkdir() (session_path / "agent.log").touch() (session_path / "trace.jsonl").touch() - agent_session = PresetAgentSession( + session = PresetSession( path=session_path, debug=True, + preset_id="ab12cd34", ) + session.write_state(get_session_state()) async def run_agent(**kwargs): - assert kwargs["agent_session"] is agent_session + assert kwargs["session"] is session assert (session_path / "prompt.md").is_file() return PresetAgentProcessOutput( report_data=json.loads( @@ -337,7 +348,7 @@ async def run_agent(**kwargs): store=creation_context.store, keep_service=keep_service, build_name="qwen-build", - agent_session=agent_session, + session=session, ) assert result.preset.base == "Qwen/Qwen3.5-27B" @@ -355,7 +366,9 @@ def _store(self, tmp_path, monkeypatch, *ids): root = store / preset_id (root / "trials" / "1").mkdir(parents=True) (root / "trials" / "1" / "trial.json").write_text("{}") - (root / "session.json").write_text(json.dumps({"status": "failed"})) + (root / "session.json").write_text( + get_session_state(status="failed").model_dump_json() + ) monkeypatch.setattr( "dstack._internal.cli.services.presets.session.get_presets_dir", lambda: store, @@ -378,7 +391,9 @@ def test_rejects_an_unknown_reference(self, tmp_path, monkeypatch): def test_warns_when_a_chained_session_is_not_included(self, tmp_path, monkeypatch, capsys): store = self._store(tmp_path, monkeypatch, "a1b2c3d4", "e5f6a7b8") (store / "e5f6a7b8" / "session.json").write_text( - json.dumps({"status": "failed", "previous": ["a1b2c3d4", "00000000"]}) + get_session_state( + id="e5f6a7b8", status="failed", previous=["a1b2c3d4", "00000000"] + ).model_dump_json() ) resolve_previous_sessions(["e5f6a7b8", "a1b2c3d4"]) @@ -390,10 +405,12 @@ def test_warns_when_a_chained_session_is_not_included(self, tmp_path, monkeypatc def test_rejects_a_previous_session_that_is_still_running(self, tmp_path, monkeypatch): store = self._store(tmp_path, monkeypatch, "a1b2c3d4") - (store / "a1b2c3d4" / "session.json").write_text(json.dumps({"status": "running"})) + (store / "a1b2c3d4" / "session.json").write_text( + get_session_state(id="a1b2c3d4").model_dump_json() + ) monkeypatch.setattr( "dstack._internal.cli.services.presets.create.session_process_alive", - lambda manifest: True, + lambda state: True, ) with pytest.raises(CLIError, match="still running"): @@ -401,10 +418,12 @@ def test_rejects_a_previous_session_that_is_still_running(self, tmp_path, monkey def test_accepts_a_stale_running_session_whose_process_died(self, tmp_path, monkeypatch): store = self._store(tmp_path, monkeypatch, "a1b2c3d4") - (store / "a1b2c3d4" / "session.json").write_text(json.dumps({"status": "running"})) + (store / "a1b2c3d4" / "session.json").write_text( + get_session_state(id="a1b2c3d4").model_dump_json() + ) monkeypatch.setattr( "dstack._internal.cli.services.presets.create.session_process_alive", - lambda manifest: False, + lambda state: False, ) sessions = resolve_previous_sessions(["a1b2c3d4"]) @@ -435,7 +454,7 @@ def test_flag_overrides_and_property_stands_without_it(self): def configuration(): # A fresh object per call: the merger mutates its input. return PresetConfiguration( - name="qwen", model={"base": "Qwen/Qwen3.5-27B"}, previous=["from-config"] + name="qwen", base="Qwen/Qwen3.5-27B", previous=["from-config"] ) overridden = _get_effective_configuration( @@ -456,14 +475,15 @@ async def test_installs_records_pins_manifest_and_extends_the_prompt( root = store / "8d3b01aa" (root / "trials" / "1").mkdir(parents=True) (root / "trials" / "1" / "trial.json").write_text('{"learned": "x"}') - (root / "session.json").write_text(json.dumps({"status": "failed"})) + (root / "session.json").write_text(get_session_state(status="failed").model_dump_json()) monkeypatch.setattr( "dstack._internal.cli.services.presets.session.get_presets_dir", lambda: store, ) session_path = tmp_path / "fresh" session_path.mkdir() - agent_session = PresetAgentSession(path=session_path, debug=False) + session = PresetSession(path=session_path, debug=False, preset_id="ab12cd34") + session.write_state(get_session_state(previous=["8d3b01aa"])) seen = {} async def run_agent(**kwargs): @@ -487,14 +507,14 @@ async def run_agent(**kwargs): source_configuration=creation_context.source_configuration, store=creation_context.store, build_name="qwen-build", - agent_session=agent_session, + session=session, previous=resolve_previous_sessions(["8d3b01aa"]), ) assert seen["record"] is True assert "## Previous Sessions" in seen["prompt"] assert "8d3b01aa" in seen["prompt"] - assert agent_session.read_manifest()["previous"] == ["8d3b01aa"] + assert session.read_state().previous == ["8d3b01aa"] # constraints.json is a session record even without --debug. assert (session_path / "constraints.json").is_file() @@ -537,23 +557,25 @@ async def no_sleep(_): ), final_run_name="qwen-build-2", keep_final_service=True, - agent_session=_agent_session(tmp_path), + session=_agent_session(tmp_path), ) assert runs.stopped_names == ["qwen-build-1"] -def _agent_session(tmp_path, *, debug: bool = False) -> PresetAgentSession: +def _agent_session(tmp_path, *, debug: bool = False) -> PresetSession: path = tmp_path / "agent-running" path.mkdir() (path / "agent.log").touch() if debug: (path / "trace.jsonl").touch() - return PresetAgentSession( + session = PresetSession( path=path, debug=debug, preset_id="ab12cd34", ) + session.write_state(get_session_state()) + return session class _FakeRuns: @@ -592,7 +614,7 @@ def test_a_live_session_prints_the_log_alone(self, tmp_path, monkeypatch, capsys monkeypatch.setenv("HOME", str(tmp_path)) monkeypatch.setenv("USERPROFILE", str(tmp_path)) session = _agent_session(tmp_path) - print_preset_progress("provisioning", agent_session=session) + print_preset_progress("provisioning", session=session) print_session_log(session) @@ -605,7 +627,7 @@ class TestFindings: def test_passes_through_to_constraints(self): configuration = PresetConfiguration( name="qwen", - model={"base": "Qwen/Qwen3-32B"}, + base="Qwen/Qwen3-32B", max_ttft=5000, min_context_length=8192, concurrency=8, @@ -628,7 +650,7 @@ def test_rejects_a_prefix_that_leaves_nothing_unique(self): with pytest.raises(ValidationError, match="less than input_tokens"): PresetConfiguration( name="qwen", - model={"base": "Qwen/Qwen3-32B"}, + base="Qwen/Qwen3-32B", max_ttft=5000, min_context_length=8192, concurrency=8, @@ -642,7 +664,7 @@ def test_checks_against_the_default_input_tokens(self): with pytest.raises(ValidationError, match="less than input_tokens"): PresetConfiguration( name="qwen", - model={"base": "Qwen/Qwen3-32B"}, + base="Qwen/Qwen3-32B", max_ttft=5000, min_context_length=8192, concurrency=8, @@ -655,7 +677,7 @@ class TestPerformanceConstraints: def test_max_ttft_reaches_the_constraints(self): configuration = PresetConfiguration( name="qwen", - model={"base": "Qwen/Qwen3-32B"}, + base="Qwen/Qwen3-32B", min_context_length=8192, concurrency=8, trials=1, @@ -673,7 +695,7 @@ def test_max_ttft_reaches_the_constraints(self): def test_throughput_is_derived_not_read(self): # A miscomputed field must not become the number we rank on. preset = get_preset() - benchmark = preset.validations[0].benchmark + benchmark = preset.benchmark benchmark.metrics.output_tok_per_s = 999999.0 benchmark.metrics.per_user_tok_per_s = 999999.0 @@ -687,7 +709,7 @@ class TestBuildConstraints: def test_renders_defaults_for_the_optional_fields(self): configuration = PresetConfiguration( name="qwen", - model={"base": "Qwen/Qwen3-32B"}, + base="Qwen/Qwen3-32B", concurrency=8, trials=3, max_ttft=5000, @@ -720,7 +742,7 @@ def test_renders_defaults_for_the_optional_fields(self): def test_renders_custom_dataset_without_request_shape(self): configuration = PresetConfiguration( name="qwen", - model={"base": "Qwen/Qwen3-32B"}, + base="Qwen/Qwen3-32B", min_context_length=32768, max_ttft=5000, trials=3, @@ -783,7 +805,7 @@ def test_copies_report_redacted(self, tmp_path): _save_final_report_copy( workspace=workspace, - agent_session=session, + session=session, redacted_values=["dstack-secret"], ) @@ -798,7 +820,7 @@ def test_missing_report_is_no_op(self, tmp_path): _save_final_report_copy( workspace=workspace, - agent_session=session, + session=session, redacted_values=["dstack-secret"], ) @@ -821,14 +843,14 @@ async def create(**kwargs): with pytest.raises(KeyboardInterrupt): create_preset( api=SimpleNamespace(), - configuration=PresetConfiguration(name="qwen", model={"base": "Qwen/Qwen3.5-27B"}), + configuration=PresetConfiguration(name="qwen", 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" + state = json.loads((sessions[0] / "session.json").read_text()) + assert state["status"] == "interrupted" output = capsys.readouterr().out assert "--resume" in output assert sessions[0].name in output @@ -837,33 +859,47 @@ 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") + session = PresetSession(path=session_dir, debug=False, preset_id="ab12cd34") + session.write_state(get_session_state()) 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)) + state = session.read_state() + assert state is not None + state.run = get_session_run( + workspace=PresetSessionWorkspace(path=str(workspace_root), alias=str(workspace_root)) + ) + session.write_state(state) _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" + assert session.read_state().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) + session = PresetSession(path=session_dir, debug=False, preset_id="fe98dc76") + session.write_state(get_session_state(id="fe98dc76")) + workspace, workspace_record = create_agent_workspace(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") + state = session.read_state() + assert state is not None + state.run = get_session_run( + workspace=workspace_record, + claude_model="claude-pinned", + claude_session_id="sid-xyz", + ) + session.write_state(state) captured = {} async def run_agent(**kwargs): @@ -884,22 +920,23 @@ async def run_agent(**kwargs): configuration=creation_context.configuration, source_configuration=creation_context.source_configuration, store=creation_context.store, - agent_session=agent_session, - resume=True, + session=session, + mode="resume", ) 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) + remove_agent_workspace(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") + session = PresetSession(path=session_dir, debug=False, preset_id="ab34ef12") + session.write_state(get_session_state(id="ab34ef12")) captured = {} async def run_agent(**kwargs): @@ -921,11 +958,11 @@ async def run_agent(**kwargs): source_configuration=creation_context.source_configuration, store=creation_context.store, build_name="qwen-build", - agent_session=agent_session, + session=session, user_prompt="Optimize for RAG traffic.", ) - assert agent_session.read_user_prompt() == "Optimize for RAG traffic." + assert session.read_user_prompt() == "Optimize for RAG traffic." assert "## Additional instructions" in captured["prompt"] assert "Optimize for RAG traffic." in captured["prompt"] @@ -936,13 +973,17 @@ async def test_resume_keeps_the_pinned_user_prompt( 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) + session = PresetSession(path=session_dir, debug=False, preset_id="ab34ef12") + session.write_state(get_session_state(id="ab34ef12")) + workspace, workspace_record = create_agent_workspace(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.") + state = session.read_state() + assert state is not None + state.run = get_session_run(workspace=workspace_record, claude_session_id="sid-abc") + session.write_state(state) + session.write_user_prompt("Optimize for RAG traffic.") captured = {} async def run_agent(**kwargs): @@ -963,8 +1004,8 @@ async def run_agent(**kwargs): configuration=creation_context.configuration, source_configuration=creation_context.source_configuration, store=creation_context.store, - agent_session=agent_session, - resume=True, + session=session, + mode="resume", user_prompt="A different prompt.", ) @@ -972,7 +1013,7 @@ async def run_agent(**kwargs): 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) + remove_agent_workspace(session) class TestFleetOffersPreview: @@ -995,12 +1036,12 @@ def test_no_offers_shows_the_shared_warning_without_failing(self, capsys): class TestSessionLog: - def _session(self, tmp_path, preset_id: str, status: str, log: str) -> PresetAgentSession: + def _session(self, tmp_path, preset_id: str, status: str, log: str) -> PresetSession: 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) + session = PresetSession(path=session_dir, debug=False, preset_id=preset_id) + session.write_state(get_session_state(id=preset_id, status=status)) return session def test_load_agent_session_reads_any_status(self, tmp_path, monkeypatch): @@ -1010,13 +1051,13 @@ def test_load_agent_session_reads_any_status(self, tmp_path, monkeypatch): lambda: tmp_path, ) # A failed session is off-limits to follow/resume, but its log is readable. - session = load_agent_session("dead0000") + session = load_preset_session("dead0000") assert session.preset_id == "dead0000" with pytest.raises(CLIError, match="Unknown preset"): - load_agent_session("nope0000") + load_preset_session("nope0000") def test_lists_a_failed_session(self, tmp_path, monkeypatch): - from dstack._internal.cli.services.presets.session import list_agent_sessions + from dstack._internal.cli.services.presets.session import list_preset_sessions self._session(tmp_path, "dead0000", "failed", "[t] boom\n") self._session(tmp_path, "beef0000", "success", "[t] saved preset\n") @@ -1025,7 +1066,7 @@ def test_lists_a_failed_session(self, tmp_path, monkeypatch): lambda: tmp_path, ) - listed = {entry["id"]: entry["status"] for entry in list_agent_sessions()} + listed = {entry["id"]: entry["status"] for entry in list_preset_sessions()} assert listed == {"dead0000": "failed", "beef0000": "success"} @@ -1045,15 +1086,22 @@ def test_print_session_log_notes_empty_log(self, tmp_path, capsys): class TestFollowPreset: - def _detached_session(self, tmp_path, configuration_yaml: str) -> PresetAgentSession: + def _detached_session(self, tmp_path, configuration_yaml: str) -> PresetSession: 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) + session = PresetSession(path=session_dir, debug=False, preset_id="ab12cd34") + workspace, workspace_record = create_agent_workspace(session) workspace.constraints_path.write_text('{"run_name_prefix": "qwen-build"}') - session.update_manifest(status="running", agent_pid=987654321) + session.write_state( + get_session_state( + run=get_session_run( + workspace=workspace_record, + agent=PresetSessionProcess(pid=987654321, started_at=None), + ) + ) + ) return session def test_finalizes_a_detached_session(self, creation_context, monkeypatch, tmp_path): @@ -1065,7 +1113,7 @@ def test_finalizes_a_detached_session(self, creation_context, monkeypatch, tmp_p lambda: tmp_path, ) monkeypatch.setattr( - "dstack._internal.cli.services.presets.create.load_attachable_agent_session", + "dstack._internal.cli.services.presets.create.load_attachable_session", lambda preset_id: session, ) @@ -1089,7 +1137,7 @@ async def fake_attach(**kwargs): assert result.preset.id == "ab12cd34" assert creation_context.store.get("ab12cd34") is not None - assert session.read_manifest()["status"] == "success" + assert session.read_state().status == "success" def test_agent_death_without_report_suspends_instead_of_failing( self, creation_context, monkeypatch, tmp_path @@ -1098,7 +1146,7 @@ def test_agent_death_without_report_suspends_instead_of_failing( 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", + "dstack._internal.cli.services.presets.create.load_attachable_session", lambda preset_id: session, ) @@ -1117,7 +1165,7 @@ async def fake_attach(**kwargs): preset_id="ab12cd34", ) - assert session.read_manifest()["status"] == "interrupted" + assert session.read_state().status == "interrupted" def test_backs_off_when_claim_is_held(self, creation_context, monkeypatch, tmp_path): session = self._detached_session( @@ -1127,7 +1175,7 @@ def test_backs_off_when_claim_is_held(self, creation_context, monkeypatch, tmp_p held = try_claim_session(session) assert held is not None monkeypatch.setattr( - "dstack._internal.cli.services.presets.create.load_attachable_agent_session", + "dstack._internal.cli.services.presets.create.load_attachable_session", lambda preset_id: session, ) # follow must refuse rather than double-finalize; the session is untouched. @@ -1137,16 +1185,16 @@ def test_backs_off_when_claim_is_held(self, creation_context, monkeypatch, tmp_p store=creation_context.store, preset_id="ab12cd34", ) - assert session.read_manifest()["status"] == "running" + assert session.read_state().status == "running" release_session_claim(held) class TestStopPresetSession: - def _session_dir(self, tmp_path, manifest: dict): - session_dir = tmp_path / ".dstack" / "presets" / manifest["id"] + def _session_dir(self, tmp_path, state: PresetSessionState): + session_dir = tmp_path / ".dstack" / "presets" / state.id session_dir.mkdir(parents=True) (session_dir / "agent.log").touch() - (session_dir / "session.json").write_text(json.dumps(manifest)) + (session_dir / "session.json").write_text(state.model_dump_json()) return session_dir def _patch_root(self, monkeypatch, tmp_path): @@ -1167,7 +1215,7 @@ 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}) + self._session_dir(tmp_path, get_session_state(**{"id": "ab12cd34", "status": status})) stop_preset_session(SimpleNamespace(), "ab12cd34") @@ -1180,12 +1228,12 @@ def test_finalizes_completed_session_and_reports_created(self, tmp_path, monkeyp (workspace / "w" / "final_report.json").write_text("{}") self._session_dir( tmp_path, - { - "id": "ab12cd34", - "status": "running", - "workspace": str(workspace), - "keep_service": True, - }, + get_session_state( + run=get_session_run( + workspace=PresetSessionWorkspace(path=str(workspace), alias=str(workspace)), + finalize=PresetSessionFinalize(project="main", keep_service=True), + ), + ), ) calls = [] monkeypatch.setattr( @@ -1206,22 +1254,19 @@ 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, - }, + # A live owner: this very process. + get_session_state( + run=get_session_run(agent=PresetSessionProcess(pid=os.getpid(), started_at=None)) + ), ) order = [] monkeypatch.setattr( - "dstack._internal.cli.services.presets.session._pid_alive", + "dstack._internal.cli.services.presets.session.process_alive", lambda pid, started_at=None: True, ) monkeypatch.setattr( "dstack._internal.cli.services.presets.create.terminate_agent_process", - lambda manifest: order.append("terminate"), + lambda state: order.append("terminate"), ) monkeypatch.setattr( "dstack._internal.cli.services.presets.create._stop_active_session_runs", @@ -1243,20 +1288,20 @@ def record_finish(session, status): # 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" + state = json.loads((session_dir / "session.json").read_text()) + assert state["status"] == "interrupted" out = capsys.readouterr().out assert "creation interrupted" in out class TestStopActiveSessionRuns: - def _session(self, tmp_path) -> PresetAgentSession: + def _session(self, tmp_path) -> PresetSession: 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") + return PresetSession(path=session_dir, debug=False, preset_id="ab12cd34") def _api(self, statuses: dict, stopped: list) -> SimpleNamespace: def get(project, name): @@ -1308,27 +1353,33 @@ def _session_dir( ): session_dir = tmp_path / preset_id (session_dir / "workspace" / "w").mkdir(parents=True) - manifest = { - "id": preset_id, - "status": status, - "keep_service": keep_service, + workspace_path = str(session_dir / "workspace") + state = get_session_state( + id=preset_id, + status=status, # 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: + owner=PresetSessionProcess(pid=987654321, started_at=0.0), + run=( + get_session_run( + workspace=PresetSessionWorkspace(path=workspace_path, alias=workspace_path), + finalize=PresetSessionFinalize(project=project, keep_service=keep_service), + ) + if project is not None + # A session from before the finalize context was persisted has + # no reconcilable run. + else None + ), + ) + if owner_alive and state.run is not None: # 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)) + state.run.agent = PresetSessionProcess(pid=os.getpid(), started_at=None) + (session_dir / "session.json").write_text(state.model_dump_json()) 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. + # reconcile iterates via session.iter_preset_sessions -> session.get_presets_dir. monkeypatch.setattr( "dstack._internal.cli.services.presets.session.get_presets_dir", lambda: tmp_path, @@ -1398,7 +1449,7 @@ def boom(**kwargs): class TestSessionClaim: def _session(self, tmp_path): (tmp_path / "sess").mkdir() - return PresetAgentSession(path=tmp_path / "sess", debug=False, preset_id="sess") + return PresetSession(path=tmp_path / "sess", debug=False, preset_id="sess") def test_claim_is_exclusive_and_releasable(self, tmp_path): session = self._session(tmp_path) @@ -1423,21 +1474,61 @@ def test_claim_acquires_when_lock_file_is_unheld(self, tmp_path): 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 + assert ( + session_process_alive( + get_session_state( + run=get_session_run( + agent=PresetSessionProcess(pid=os.getpid(), 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 + assert ( + session_process_alive( + get_session_state(owner=PresetSessionProcess(pid=987654321, started_at=0.0)) + ) + is False + ) + assert session_process_alive(get_session_state()) is False -class TestMarkSessionOwner: - def test_persists_finalize_context(self, tmp_path): +class TestBeginRun: + def test_records_the_run_whole_and_keeps_claude_state(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 + session = PresetSession(path=tmp_path / "s", debug=False, preset_id="s") + workspace = PresetSessionWorkspace(path=str(tmp_path / "w"), alias=str(tmp_path / "w")) + session.write_state( + get_session_state( + id="s", + run=get_session_run( + workspace=workspace, + finalize=PresetSessionFinalize(project="old", keep_service=False), + claude_model="claude-pinned", + agent=PresetSessionProcess(pid=1, started_at=None), + claude_session_id="sid-1", + ), + ) + ) + + session.begin_run( + workspace=workspace, + finalize=PresetSessionFinalize(project="main", keep_service=True), + claude_model=None, + ) + + state = session.read_state() + assert state is not None + assert state.owner is not None + assert state.owner.pid == os.getpid() + assert state.owner.started_at is not None + assert state.run is not None + assert state.run.finalize == PresetSessionFinalize(project="main", keep_service=True) + # Everything the earlier run established survives: the claude state so a + # resume finds it, and the agent reference so following a live detached + # agent does not read it as dead (and kill it). + assert state.run.claude_model == "claude-pinned" + assert state.run.claude_session_id == "sid-1" + assert state.run.agent == PresetSessionProcess(pid=1, started_at=None) diff --git a/src/tests/_internal/cli/services/presets/test_output.py b/src/tests/_internal/cli/services/presets/test_output.py index 8c1328054..8d1032923 100644 --- a/src/tests/_internal/cli/services/presets/test_output.py +++ b/src/tests/_internal/cli/services/presets/test_output.py @@ -5,10 +5,10 @@ import pytest from rich.table import Table +from dstack._internal.cli.models.presets import PresetWorkload 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 +from tests._internal.cli.common import get_preset, plain_console pytestmark = pytest.mark.windows @@ -27,7 +27,7 @@ def test_small_values_keep_three_significant_digits(self): class TestFormatPresetBenchmark: def test_formats_second_scale_ttft_without_scientific_notation(self): preset = get_preset() - ttft = preset.validations[0].benchmark.metrics.ttft_ms + ttft = preset.benchmark.metrics.ttft_ms ttft.mean = 8148.3 ttft.p50 = 8151.4 ttft.p99 = 8334.2 @@ -43,28 +43,44 @@ def test_formats_second_scale_ttft_without_scientific_notation(self): class TestFormatPresetObjective: - def test_shows_the_shared_prefix_the_benchmark_actually_used(self): + def test_shows_the_requested_shared_prefix(self): preset = get_preset() - preset.validations[0].benchmark.workload.shared_prefix_tokens = 768 + preset.configuration.shared_prefix_tokens = 768 assert output_module.format_preset_objective(preset) == ( "[secondary]io=1K/128 prefix=75% conc=1[/]" ) - def test_a_preset_saved_before_the_field_existed_still_loads(self): - # `shared_prefix_tokens` is absent from every preset saved so far. + def test_treats_an_unset_shared_prefix_as_none_shared(self): preset = get_preset() - assert preset.validations[0].benchmark.workload.shared_prefix_tokens is None + assert preset.configuration.shared_prefix_tokens is None assert output_module.format_preset_objective(preset) == ( "[secondary]io=1K/128 prefix=0% conc=1[/]" ) def test_shows_the_dataset_instead_of_the_request_shape(self): - # With a custom dataset the io shape is measured, not configured, so the - # contract cell names the dataset instead. + # A dataset defines its own request shape, so the cell names it instead. preset = get_preset() - preset.validations[0].benchmark.workload.dataset = "spec_bench" + preset.configuration.dataset = "spec_bench" + + assert output_module.format_preset_objective(preset) == ( + "[secondary]data=spec_bench conc=1[/]" + ) + + def test_renders_the_requested_workload_not_the_measured_one(self): + # They diverge with a dataset, where the benchmark records measured means. + preset = get_preset() + preset.configuration.dataset = "spec_bench" + benchmark = preset.benchmark + benchmark.workload = PresetWorkload( + api=benchmark.workload.api, + num_requests=benchmark.workload.num_requests, + input_tokens=347, + output_tokens=2451, + concurrency=8, + dataset="spec_bench", + ) assert output_module.format_preset_objective(preset) == ( "[secondary]data=spec_bench conc=1[/]" @@ -213,7 +229,7 @@ def test_sorts_all_rows_newest_first_without_grouping(self, monkeypatch): monkeypatch.setattr(output_module, "console", plain_console(buffer, width=200)) old = get_preset() new = old.model_copy( - update={"id": "11aa22bb", "created_at": old.created_at + timedelta(days=2)} + update={"id": "11aa22bb", "submitted_at": old.submitted_at + timedelta(days=2)} ) sessions = [ { @@ -456,3 +472,18 @@ def test_the_best_failed_trial_is_gold_while_none_passes(self): assert spark.count("gold1") == 1 assert spark.count("indian_red1") == 2 assert "sea_green3" not in spark + + +def test_warn_respects_stderr_redirect(): + """`preset -w` suppresses store warnings during refreshes via + redirect_stderr; that only works while `error_console` resolves + `sys.stderr` at write time.""" + import io + from contextlib import redirect_stderr + + from dstack._internal.cli.utils.common import warn + + buffer = io.StringIO() + with redirect_stderr(buffer): + warn("torn render", stderr=True) + assert "torn render" in buffer.getvalue() diff --git a/src/tests/_internal/cli/services/presets/test_prompt.py b/src/tests/_internal/cli/services/presets/test_prompt.py index e429a5ac5..1c163e627 100644 --- a/src/tests/_internal/cli/services/presets/test_prompt.py +++ b/src/tests/_internal/cli/services/presets/test_prompt.py @@ -9,16 +9,31 @@ class TestSystemPrompt: def test_stays_byte_identical_without_user_prompt(self): - text = get_preset_agent_system_prompt() + text = get_preset_agent_system_prompt( + user_prompt=None, baseline=False, previous=(), custom_dataset=False + ) - assert text == get_preset_agent_system_prompt(None) == get_preset_agent_system_prompt("") + assert ( + text + == get_preset_agent_system_prompt( + user_prompt=None, baseline=False, previous=(), custom_dataset=False + ) + == get_preset_agent_system_prompt( + user_prompt="", baseline=False, previous=(), custom_dataset=False + ) + ) assert "## Additional instructions" not in text assert " stays." @@ -72,7 +99,9 @@ def test_rejects_unknown_variables_even_in_a_dropped_branch(self, tmp_path, monk # `previous` is unset, so the branch would be dropped; the typo inside # it must not hide behind that. with pytest.raises(CLIError, match="Unknown variable"): - get_preset_agent_system_prompt() + get_preset_agent_system_prompt( + user_prompt=None, baseline=False, previous=(), custom_dataset=False + ) def test_rejects_malformed_directives(self, tmp_path, monkeypatch): broken = tmp_path / "system_prompt.md" @@ -80,19 +109,33 @@ def test_rejects_malformed_directives(self, tmp_path, monkeypatch): monkeypatch.setattr(prompt_module, "_SYSTEM_PROMPT_PATH", broken) with pytest.raises(CLIError, match="Invalid directive"): - get_preset_agent_system_prompt() + get_preset_agent_system_prompt( + user_prompt=None, baseline=False, previous=(), custom_dataset=False + ) broken.write_text("An opener that never closes its comment: BCD\n") monkeypatch.setattr(prompt_module, "_SYSTEM_PROMPT_PATH", doc) - assert get_preset_agent_system_prompt(baseline=True) == "ABD" - assert get_preset_agent_system_prompt() == "ACD" + assert ( + get_preset_agent_system_prompt( + user_prompt=None, baseline=True, previous=(), custom_dataset=False + ) + == "ABD" + ) + assert ( + get_preset_agent_system_prompt( + user_prompt=None, baseline=False, previous=(), custom_dataset=False + ) + == "ACD" + ) def test_dedents_only_an_exactly_indented_body(self, tmp_path, monkeypatch): doc = tmp_path / "system_prompt.md" @@ -124,7 +167,9 @@ def test_dedents_only_an_exactly_indented_body(self, tmp_path, monkeypatch): for content, expected in cases: doc.write_text(content) previous = "x" if "previous" in content else None - rendered = get_preset_agent_system_prompt(baseline=True, previous=previous) + rendered = get_preset_agent_system_prompt( + user_prompt=None, baseline=True, previous=previous, custom_dataset=False + ) assert rendered.strip("\n") == expected, content def test_nested_blocks_render_one_branch(self, tmp_path, monkeypatch): @@ -144,10 +189,27 @@ def test_nested_blocks_render_one_branch(self, tmp_path, monkeypatch): ) monkeypatch.setattr(prompt_module, "_SYSTEM_PROMPT_PATH", nested) - assert get_preset_agent_system_prompt().strip() == "" - assert get_preset_agent_system_prompt(baseline=True).strip() == "SOLO" - assert get_preset_agent_system_prompt(previous="a1, b2").strip() == "IDS=a1, b2" - both = get_preset_agent_system_prompt(baseline=True, previous="a1") + assert ( + get_preset_agent_system_prompt( + user_prompt=None, baseline=False, previous=(), custom_dataset=False + ).strip() + == "" + ) + assert ( + get_preset_agent_system_prompt( + user_prompt=None, baseline=True, previous=(), custom_dataset=False + ).strip() + == "SOLO" + ) + assert ( + get_preset_agent_system_prompt( + user_prompt=None, baseline=False, previous=("a1", "b2"), custom_dataset=False + ).strip() + == "IDS=a1, b2" + ) + both = get_preset_agent_system_prompt( + user_prompt=None, baseline=True, previous=("a1",), custom_dataset=False + ) assert both.strip() == "IDS=a1\n\nSEEDED" def test_unbalanced_blocks_fail_loudly(self, tmp_path, monkeypatch): @@ -161,4 +223,6 @@ def test_unbalanced_blocks_fail_loudly(self, tmp_path, monkeypatch): broken.write_text(content) monkeypatch.setattr(prompt_module, "_SYSTEM_PROMPT_PATH", broken) with pytest.raises(CLIError, match=error): - get_preset_agent_system_prompt() + get_preset_agent_system_prompt( + user_prompt=None, baseline=False, previous=(), custom_dataset=False + ) diff --git a/src/tests/_internal/cli/services/presets/test_store.py b/src/tests/_internal/cli/services/presets/test_store.py index eb0cfa422..51726bfc4 100644 --- a/src/tests/_internal/cli/services/presets/test_store.py +++ b/src/tests/_internal/cli/services/presets/test_store.py @@ -11,7 +11,7 @@ from dstack._internal.core.errors import ConfigurationError from dstack._internal.core.models.envs import EnvSentinel from dstack._internal.core.models.files import FilePathMapping -from tests._internal.cli.preset_factories import get_preset +from tests._internal.cli.common import get_preset pytestmark = pytest.mark.windows @@ -23,12 +23,12 @@ def test_saves_and_lists_self_contained_preset(self, tmp_path: Path): path = store.save(preset) - assert path == (tmp_path / "presets" / "8f3a12c4" / "preset.yaml") + assert path == (tmp_path / "presets" / "8f3a12c4" / "preset.yml") 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 data["submitted_at"] == "2026-01-02T03:04:00Z" assert "presets" not in data assert store.list() == [preset] assert store.get(preset.id) == preset @@ -44,27 +44,15 @@ def test_saving_same_id_overwrites_existing_preset(self, tmp_path: Path): assert store.get(preset.id) == updated - def test_migrates_legacy_layout_and_deletes_permanently(self, tmp_path: Path): + def test_ignores_directories_without_a_preset_file(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, - ) - ) + # A creation-session directory, or anything else that is not a preset. + (root / "ab12cd34").mkdir(parents=True) + (root / "ab12cd34" / "session.json").write_text("{}") - 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 not (root / preset.id).exists() assert store.list() == [] + assert store.delete("ab12cd34") is False def test_skips_invalid_preset_on_list_but_keeps_it_deletable(self, tmp_path: Path, capsys): store = PresetStore(tmp_path / "presets") @@ -72,7 +60,7 @@ def test_skips_invalid_preset_on_list_but_keeps_it_deletable(self, tmp_path: Pat store.save(valid) path = store.save(get_preset()) data = yaml.safe_load(path.read_text()) - data["validations"][0].pop("benchmark") + data["benchmark"].pop("metrics") path.write_text(yaml.safe_dump(data, sort_keys=False)) # One corrupt file must not take down every read; the warning goes to @@ -83,6 +71,197 @@ def test_skips_invalid_preset_on_list_but_keeps_it_deletable(self, tmp_path: Pat assert store.delete(get_preset().id) is True assert [preset.id for preset in store.list()] == [valid.id] + def test_upgrades_pre_0_21_2_preset(self, tmp_path: Path): + store = PresetStore(tmp_path / "presets") + directory = tmp_path / "presets" / "523096d6" + directory.mkdir(parents=True) + # The pre-0.21.2 shape: request echoed at the top level, the service + # under `service`, and per-group verification under `validations`. + (directory / "preset.yml").write_text( + """ + base: deepseek-ai/DeepSeek-V4-Flash + id: 523096d6 + name: dsv4-flash + model: deepseek-ai/DeepSeek-V4-Flash-0731 + context_length: 1048576 + trial: 2 + min_context_length: 1048576 + max_ttft: 675 + created_at: '2026-08-13T15:11:13.334759+00:00' + service: + port: 80 + model: deepseek-ai/DeepSeek-V4-Flash + commands: [vllm serve] + image: vllm/vllm-openai-rocm + priority: 0 + resources: + cpu: 16.. + memory: 200GB.. + gpu: MI300X:192GB:1 + disk: 500GB.. + validations: + - replicas: + - resources: + - cpu: + arch: x86 + count: {min: 20, max: 20} + memory: {min: 240.0, max: 240.0} + gpu: + vendor: amd + name: [MI300X] + count: {min: 1, max: 1} + memory: {min: 192.0, max: 192.0} + disk: + size: {min: 720.0, max: 720.0} + benchmark: + tool: vllm bench serve + tool_version: 0.26.1 + command: vllm bench serve --dataset-name random + workload: + api: completions + num_requests: 16 + input_tokens: 10000 + output_tokens: 1500 + concurrency: 4 + shared_prefix_tokens: 0 + metrics: + successful_requests: 16 + failed_requests: 0 + duration_seconds: 137.42 + total_input_tokens: 160000 + total_output_tokens: 24000 + output_tok_per_s: 174.65 + per_user_tok_per_s: 43.66 + ttft_ms: {mean: 1691.88, p50: 1240.31, p99: 4168.05} + tpot_ms: {mean: 21.57, p50: 21.84, p99: 23.28} + target: + type: server-proxy + client: + type: local + """ + ) + + preset = store.get("523096d6") + + assert preset is not None + assert preset.name == "dsv4-flash" + assert preset.best_trial == 2 + assert preset.configuration.model.base == "deepseek-ai/DeepSeek-V4-Flash" + assert preset.configuration.min_context_length == 1048576 + assert preset.configuration.max_ttft == 675 + # The measured workload is the old format's only record of the + # workload constraints; without it the objective would show defaults. + assert preset.configuration.input_tokens == 10000 + assert preset.configuration.output_tokens == 1500 + assert preset.configuration.shared_prefix_tokens == 0 + assert preset.configuration.concurrency == 4 + # `priority` is not an excluded field, so the regrouping keeps it. + assert preset.service.priority == 0 + assert [group.name for group in preset.verified_on] == [ + group.name for group in preset.service.replica_groups + ] + assert preset.verified_on[0].replicas[0].gpu.name == ["MI300X"] + assert preset.benchmark.workload.dataset == "random" + assert store.list() == [preset] + + def test_upgrade_maps_validation_replicas_to_replica_groups(self, tmp_path: Path): + store = PresetStore(tmp_path / "presets") + directory = tmp_path / "presets" / "aa11bb22" + directory.mkdir(parents=True) + # One validation; its `replicas` entries follow the service's replica + # group order, and each entry holds that group's per-replica resources. + (directory / "preset.yml").write_text( + """ + base: Qwen/Qwen3-32B + id: aa11bb22 + name: null + model: Qwen/Qwen3-32B + context_length: 32768 + trial: 1 + created_at: '2026-08-13T15:11:13+00:00' + service: + port: 80 + model: Qwen/Qwen3-32B + image: vllm/vllm-openai + replicas: + - name: small + count: 1 + commands: [vllm serve] + resources: {gpu: H100:80GB:1, disk: 100GB..} + - name: large + count: 2 + commands: [vllm serve] + resources: {gpu: H200:141GB:1, disk: 100GB..} + validations: + - replicas: + - resources: + - cpu: {count: {min: 8, max: 8}} + memory: {min: 100.0, max: 100.0} + gpu: {name: [H100], count: {min: 1, max: 1}, memory: {min: 80.0, max: 80.0}} + disk: {size: {min: 100.0, max: 100.0}} + - resources: + - &r + cpu: {count: {min: 16, max: 16}} + memory: {min: 200.0, max: 200.0} + gpu: {name: [H200], count: {min: 1, max: 1}, memory: {min: 141.0, max: 141.0}} + disk: {size: {min: 200.0, max: 200.0}} + - *r + benchmark: + tool: vllm bench serve + tool_version: '1.0' + command: vllm bench serve + workload: + api: completions + num_requests: 8 + input_tokens: 1024 + output_tokens: 1024 + concurrency: 2 + shared_prefix_tokens: 0 + metrics: + successful_requests: 8 + failed_requests: 0 + duration_seconds: 10.0 + total_input_tokens: 8192 + total_output_tokens: 8192 + output_tok_per_s: 819.2 + per_user_tok_per_s: 100.0 + ttft_ms: {mean: 100.0, p50: 100.0, p99: 200.0} + tpot_ms: {mean: 10.0, p50: 10.0, p99: 20.0} + """ + ) + + preset = store.get("aa11bb22") + + assert preset is not None + assert [(group.name, len(group.replicas)) for group in preset.verified_on] == [ + ("small", 1), + ("large", 2), + ] + assert preset.verified_on[0].replicas[0].gpu.name == ["H100"] + assert preset.verified_on[1].replicas[0].gpu.name == ["H200"] + + def test_reports_pre_0_21_presets_as_one_summary_line(self, tmp_path: Path, capsys): + store = PresetStore(tmp_path / "presets") + bodies = { + # The 0.20.x shape: no `trial` recorded, so it cannot be upgraded. + "523096d6": "base: Qwen/Q\nid: 523096d6\nservice: {}\nvalidations: []\n", + # A truncated file that lost its `service` block must degrade to + # the same one-line message, not a traceback. + "7cc991bd": "base: Qwen/Q\nid: 7cc991bd\ntrial: 1\nservice: null\nvalidations: null\n", + } + for preset_id, body in bodies.items(): + directory = tmp_path / "presets" / preset_id + directory.mkdir(parents=True) + (directory / "preset.yml").write_text(body) + + assert store.list() == [] + + err = capsys.readouterr().err + assert "2 presets created before dstack 0.21 cannot be read" in err + assert "523096d6, 7cc991bd" in err + assert err.count("before dstack 0.21") == 1 + assert "validation error" not in err + def test_resolves_relative_file_paths_against_the_preset_directory(self, tmp_path: Path): store = PresetStore(tmp_path / "presets") preset = get_preset() @@ -115,7 +294,7 @@ def test_release_name_keeps_file_paths_relative(self, tmp_path: Path): # directory silently stops being portable. store.release_name("qwen") - data = yaml.safe_load((tmp_path / "presets" / preset.id / "preset.yaml").read_text()) + data = yaml.safe_load((tmp_path / "presets" / preset.id / "preset.yml").read_text()) assert data["service"]["files"][0]["local_path"] == "service/1/patches" def test_preserves_literal_env_values(self, tmp_path: Path): @@ -139,16 +318,20 @@ def test_preserves_literal_env_values(self, tmp_path: Path): class TestParsePresetConfiguration: @pytest.mark.parametrize("key", ["base", "repo"]) - def test_warns_on_nested_model_without_name(self, key: str): + def test_rejects_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) + with pytest.raises(ConfigurationError, match=f"`{key}:`"): + 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 + def test_accepts_nested_model_with_name(self): + stream = StringIO( + "type: preset\nmodel:\n repo: community/Qwen3.5-27B-GPTQ-Int4\n name: Qwen/Qwen3.5-27B\n" + ) + + configuration = store_module.parse_preset_configuration(stream) + + assert configuration.model.api_model_name == "Qwen/Qwen3.5-27B" @pytest.mark.parametrize( "body", @@ -163,7 +346,7 @@ 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) + configuration = store_module.parse_preset_configuration(stream) warn.assert_not_called() assert configuration.model is not None @@ -172,32 +355,30 @@ def test_does_not_warn_on_preferred_syntax(self, body: str): 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") + base = tmp_path 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(inline, base) == "Inline text." + assert store_module.resolve_preset_prompt(from_file, base) == "From a file." assert ( - store_module.resolve_preset_prompt( - PresetConfiguration(name="q", base="Q/M"), configuration_path - ) + store_module.resolve_preset_prompt(PresetConfiguration(name="q", base="Q/M"), base) is None ) def test_rejects_missing_and_empty_prompt_files(self, tmp_path: Path): - configuration_path = str(tmp_path / "preset.dstack.yml") + base = tmp_path (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, + base, ) with pytest.raises(ConfigurationError, match="is empty"): store_module.resolve_preset_prompt( PresetConfiguration(name="q", base="Q/M", prompt={"path": "empty.md"}), - configuration_path, + base, ) diff --git a/src/tests/_internal/cli/services/presets/test_verify.py b/src/tests/_internal/cli/services/presets/test_verify.py index 82bf96e00..70b9734dc 100644 --- a/src/tests/_internal/cli/services/presets/test_verify.py +++ b/src/tests/_internal/cli/services/presets/test_verify.py @@ -1,14 +1,14 @@ 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.models.preset_agent import AnyPresetAgentResult from dstack._internal.cli.services.presets.agent import ( PresetAgentProcessOutput, ) +from dstack._internal.cli.services.presets.build import build_preset from dstack._internal.cli.services.presets.verify import ( build_verified_preset, load_preset_agent_report, @@ -17,10 +17,13 @@ PresetAgentWorkspace, ) from dstack._internal.core.errors import CLIError +from dstack._internal.core.models.common import validate_extra_ignore from dstack._internal.core.models.envs import EnvSentinel from dstack._internal.core.models.files import FilePathMapping from dstack._internal.core.models.profiles import ProfileParams -from tests._internal.cli.preset_factories import ( +from tests._internal.cli.common import ( + get_preset, + get_preset_benchmark, get_running_service_run, get_successful_preset_report, ) @@ -29,48 +32,85 @@ class TestBuildVerifiedPreset: + def test_stores_only_the_creation_contract_not_this_machine_s_deployment(self): + # `apply` takes name, gateway, env and profile from the user's own + # configuration, so a shared preset must not carry ours. + configuration = PresetConfiguration.model_validate( + { + "type": "preset", + "base": "Qwen/Qwen3.5-27B", + "trials": 3, + "concurrency": 1, + "min_context_length": 32768, + "name": "qwen-build", + "gateway": "benchmark-gateway", + "env": {"MY-VAR": "secret", "HF_TOKEN": "hf_secret"}, + "spot_policy": "on-demand", + } + ) + base = get_preset() + + preset = build_preset( + service=base.service, + verification_replica_groups=base.verified_on, + base_model="Qwen/Qwen3.5-27B", + model="community/Qwen3.5-27B-GPTQ-Int4", + context_length=32768, + benchmark=get_preset_benchmark(), + configuration=configuration, + best_trial=1, + preset_id="8f3a12c4", + name=None, + submitted_at=datetime(2026, 1, 2, 3, 4, tzinfo=timezone.utc), + ) + + assert preset.configuration.min_context_length == 32768 + assert preset.configuration.name is None + assert preset.configuration.gateway is None + assert not preset.configuration.env + assert preset.configuration.spot_policy is None + assert "secret" not in preset.model_dump_json() + def test_successful_report_requires_benchmark(self): run = get_running_service_run() data = get_successful_preset_report(run).model_dump() data.pop("benchmark") with pytest.raises(ValidationError, match="benchmark"): - AgentFinalReport.model_validate(data) + validate_extra_ignore(AnyPresetAgentResult, data) - def test_builds_portable_self_contained_preset(self): + def test_builds_portable_self_contained_preset(self, tmp_path): 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"}, - min_context_length=8192, - gateway="benchmark-gateway", - env=["LICENSE", "TOKENIZERS_PARALLELISM=false"], - ), - report=get_successful_preset_report(run), - ) + preset = build_verified_preset( + run=run, + preset_configuration=PresetConfiguration( + name="qwen-build", + base="Qwen/Qwen3.5-27B", + min_context_length=8192, + gateway="benchmark-gateway", + env=["LICENSE", "TOKENIZERS_PARALLELISM=false"], + ), + report=get_successful_preset_report(run), + workspace_path=tmp_path, + session_path=tmp_path, + preset_id="ab12cd34", + name=None, + submitted_at=created_at, + ) 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.submitted_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.model_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" + assert preset.verified_on[0].replicas[0].gpu.name == ["A6000"] def test_rewrites_file_paths_onto_the_mirrored_session_copies(self, tmp_path): # `files` local paths resolve into the agent workspace at submission, and @@ -92,12 +132,13 @@ def test_rewrites_file_paths_onto_the_mirrored_session_copies(self, tmp_path): preset = build_verified_preset( run=run, - preset_configuration=PresetConfiguration( - name="qwen-build", model={"base": "Qwen/Qwen3.5-27B"} - ), + preset_configuration=PresetConfiguration(name="qwen-build", base="Qwen/Qwen3.5-27B"), report=get_successful_preset_report(run), workspace_path=workspace, session_path=session, + preset_id="ab12cd34", + name=None, + submitted_at=datetime(2026, 1, 2, 3, 4, tzinfo=timezone.utc), ) assert preset.service.files[0].local_path == "service/1/patches" @@ -119,29 +160,17 @@ def test_rejects_a_file_without_a_mirrored_copy(self, tmp_path): build_verified_preset( run=run, preset_configuration=PresetConfiguration( - name="qwen-build", model={"base": "Qwen/Qwen3.5-27B"} + name="qwen-build", base="Qwen/Qwen3.5-27B" ), report=get_successful_preset_report(run), workspace_path=workspace, session_path=session, + preset_id="ab12cd34", + name=None, + submitted_at=datetime(2026, 1, 2, 3, 4, tzinfo=timezone.utc), ) - def test_rejects_files_when_no_workspace_is_attached(self, tmp_path): - run = get_running_service_run() - run.run_spec.configuration.files = [ - FilePathMapping(local_path=str(tmp_path / "patches"), path="/patches") - ] - - with pytest.raises(CLIError, match="no workspace is attached"): - build_verified_preset( - run=run, - preset_configuration=PresetConfiguration( - name="qwen-build", model={"base": "Qwen/Qwen3.5-27B"} - ), - report=get_successful_preset_report(run), - ) - - def test_rejects_benchmark_on_a_different_dataset(self): + def test_rejects_benchmark_on_a_different_dataset(self, tmp_path): run = get_running_service_run() # The report's workload defaults to `random`, but the configuration @@ -151,13 +180,18 @@ def test_rejects_benchmark_on_a_different_dataset(self): run=run, preset_configuration=PresetConfiguration( name="qwen-build", - model={"base": "Qwen/Qwen3.5-27B"}, + base="Qwen/Qwen3.5-27B", dataset="spec_bench", ), report=get_successful_preset_report(run), + workspace_path=tmp_path, + session_path=tmp_path, + preset_id="ab12cd34", + name=None, + submitted_at=datetime(2026, 1, 2, 3, 4, tzinfo=timezone.utc), ) - def test_rejects_variant_for_exact_model_request(self): + def test_rejects_variant_for_exact_model_request(self, tmp_path): run = get_running_service_run() report = get_successful_preset_report(run).model_copy(update={"model": "other/model"}) @@ -172,6 +206,11 @@ def test_rejects_variant_for_exact_model_request(self): }, ), report=report, + workspace_path=tmp_path, + session_path=tmp_path, + preset_id="ab12cd34", + name=None, + submitted_at=datetime(2026, 1, 2, 3, 4, tzinfo=timezone.utc), ) diff --git a/src/tests/_internal/cli/services/presets/test_workspace.py b/src/tests/_internal/cli/services/presets/test_workspace.py index 06bd521e6..48a7c1aba 100644 --- a/src/tests/_internal/cli/services/presets/test_workspace.py +++ b/src/tests/_internal/cli/services/presets/test_workspace.py @@ -1,6 +1,6 @@ import pytest -from dstack._internal.cli.services.presets.session import PresetAgentSession +from dstack._internal.cli.services.presets.session import PresetSession from dstack._internal.cli.services.presets.workspace import ( PresetAgentWorkspace, install_previous_records, @@ -27,7 +27,7 @@ def _previous_session(tmp_path, preset_id="8d3b01aa"): (root / "runs.jsonl").write_text("{}") (root / "trials" / "not-a-trial").mkdir() (root / "trials" / "not-a-trial" / "trial.json").write_text("{}") - return PresetAgentSession(path=root, debug=False, preset_id=preset_id) + return PresetSession(path=root, debug=False, preset_id=preset_id) def _workspace(tmp_path): @@ -74,7 +74,7 @@ def test_a_session_without_records_warns(self, tmp_path, capsys): root = tmp_path / "store" / "empty000" root.mkdir(parents=True) (root / "session.json").write_text("{}") - session = PresetAgentSession(path=root, debug=False, preset_id="empty000") + session = PresetSession(path=root, debug=False, preset_id="empty000") workspace = _workspace(tmp_path) install_previous_records(workspace, [session])