diff --git a/skills/dstack-prototyping/SKILL.md b/skills/dstack-prototyping/SKILL.md index b70bb07a5..c41cf2465 100644 --- a/skills/dstack-prototyping/SKILL.md +++ b/skills/dstack-prototyping/SKILL.md @@ -22,10 +22,10 @@ through the dstack service URL. ## Choose Where To Run -Choose only VM-based backends, SSH fleets, or Kubernetes fleets because they support idle instances and/or instance volumes. That lets later runs reuse the provisioned/idle instance or instance volumes used by runs for caching model weights (and possibly other writes). You must follow this rule even if there are fleets/backends/offers that are cheaper. The only exception from this rule is when the required GPU class (regardless of the price) is not available through VM-based backend, SSH fleet, or Kubernetes fleet. +Pick the offer whose hardware best fits the goal at hand. Only when several offers fit comparably, choose a VM-based backend, an SSH fleet, or a Kubernetes fleet: they support idle instances and/or instance volumes, so later runs reuse the provisioned/idle instance or instance volumes for caching model weights (and possibly other writes), while container-based backends start clean on every run. -Read `https://dstack.ai/docs/concepts/backends.md` to know exactly which -backends are VM-based. +Fetch `https://dstack.ai/docs/concepts/backends.md` and classify backends +from the fetched document, not from memory. ## Check Serving Sources @@ -37,13 +37,13 @@ For vLLM and SGLang, use these as credible sources: - vLLM recipes and model index: `https://recipes.vllm.ai/` and `https://recipes.vllm.ai/models.json` -- vLLM recipe docs: `https://docs.vllm.ai/projects/recipes/en/stable/` -- SGLang docs and cookbook: `https://docs.sglang.ai/` and - `https://lmsysorg.mintlify.app/cookbook/intro` - -Use deeper serving-engine writeups, such as -`https://www.lmsys.org/blog/2026-07-02-agent-assisted-sglang-development`, when -these references do not explain the model, hardware, or serving failure. +- SGLang docs: `https://docs.sglang.io/` (fetch `/llms.txt` for the page + index) +- SGLang model recipes: `https://docs.sglang.io/cookbook/autoregressive/intro` +- Release notes: `https://github.com/vllm-project/vllm/releases` and + `https://github.com/sgl-project/sglang/releases` +- Performance-loop methodology (profiling, benchmark contracts): + `https://www.lmsys.org/blog/2026-07-02-agent-assisted-sglang-development` ## Use A Task Before Service diff --git a/src/dstack/_internal/cli/commands/endpoint.py b/src/dstack/_internal/cli/commands/endpoint.py index 0bf75b42b..c68d01568 100644 --- a/src/dstack/_internal/cli/commands/endpoint.py +++ b/src/dstack/_internal/cli/commands/endpoint.py @@ -1,19 +1,35 @@ import argparse import os +import time from pathlib import Path from argcomplete import FilesCompleter # type: ignore[attr-defined] from dstack._internal.cli.commands import BaseCommand -from dstack._internal.cli.models.endpoint_presets import EndpointPresetListOutput +from dstack._internal.cli.models.endpoint_presets import ( + EndpointPreset, + EndpointPresetListOutput, +) from dstack._internal.cli.models.endpoints import EndpointConfiguration from dstack._internal.cli.services.completion import ProjectNameCompleter +from dstack._internal.cli.services.endpoints.agent import ( + find_session_name_claims, + get_presets_dir, + list_agent_sessions, + load_resumable_agent_session, +) from dstack._internal.cli.services.endpoints.apply import apply_endpoint_preset -from dstack._internal.cli.services.endpoints.create import create_endpoint_preset +from dstack._internal.cli.services.endpoints.create import ( + attach_endpoint_preset, + create_endpoint_preset, + plan_endpoint_preset, + stop_endpoint_session, +) from dstack._internal.cli.services.endpoints.output import print_endpoint_presets from dstack._internal.cli.services.endpoints.store import ( EndpointPresetStore, load_endpoint_configuration, + resolve_endpoint_prompt, ) from dstack._internal.cli.services.profile import apply_profile_args, register_profile_args from dstack._internal.cli.utils.common import confirm_ask, console @@ -25,8 +41,8 @@ class EndpointCommand(BaseCommand): - NAME = "endpoint" - DESCRIPTION = "Manage model inference endpoints" + NAME = "preset" + DESCRIPTION = "Manage model serving presets" def _register(self) -> None: self._parser.add_argument( @@ -35,20 +51,13 @@ def _register(self) -> None: metavar="NAME", default=os.getenv("DSTACK_PROJECT"), ).completer = ProjectNameCompleter() # type: ignore[attr-defined] - self._parser.set_defaults(subfunc=lambda _: self._parser.print_help()) - subparsers = self._parser.add_subparsers(dest="action") - preset_parser = subparsers.add_parser( - "preset", - help="Manage endpoint presets", - formatter_class=self._parser.formatter_class, - ) - _add_list_args(preset_parser) - preset_parser.set_defaults(subfunc=self._list) - preset_subparsers = preset_parser.add_subparsers(dest="preset_action") + _add_list_args(self._parser) + self._parser.set_defaults(subfunc=self._list) + preset_subparsers = self._parser.add_subparsers(dest="action") list_parser = preset_subparsers.add_parser( "list", - help="List endpoint presets", + help="List presets", formatter_class=self._parser.formatter_class, ) _add_list_args(list_parser) @@ -56,10 +65,10 @@ def _register(self) -> None: get_parser = preset_subparsers.add_parser( "get", - help="Get an endpoint preset", + help="Get a preset", formatter_class=self._parser.formatter_class, ) - get_parser.add_argument("preset", metavar="ID", help="The preset ID") + get_parser.add_argument("preset", metavar="ID", help="The preset ID or name") get_parser.add_argument( "--json", action="store_true", @@ -70,7 +79,7 @@ def _register(self) -> None: create_parser = preset_subparsers.add_parser( "create", - help="Create an endpoint preset", + help="Create a preset", formatter_class=self._parser.formatter_class, ) _add_configuration_args(create_parser) @@ -80,21 +89,65 @@ def _register(self) -> None: action="store_true", help="Leave the verified service running", ) + create_parser.add_argument( + "--max-trials", + type=int, + metavar="N", + help="The maximum number of benchmarked trials before the best one is promoted", + ) create_parser.add_argument( "--debug", action="store_true", help="Save the agent prompt and raw trace", ) + create_parser.add_argument( + "--resume", + metavar="ID", + help="Resume an interrupted preset creation session by its preset ID", + ) + create_parser.add_argument( + "-y", "--yes", action="store_true", help="Do not ask for confirmation" + ) create_parser.set_defaults(subfunc=self._create) + attach_parser = preset_subparsers.add_parser( + "attach", + help="Attach to a detached preset creation session", + formatter_class=self._parser.formatter_class, + ) + attach_parser.add_argument("preset", metavar="ID", help="The preset ID or name") + attach_parser.add_argument( + "--keep-service", + action="store_true", + help="Leave the verified service running", + ) + attach_parser.set_defaults(subfunc=self._attach) + + stop_parser = preset_subparsers.add_parser( + "stop", + help="Stop a running preset creation session", + formatter_class=self._parser.formatter_class, + ) + stop_parser.add_argument("preset", metavar="ID", help="The preset ID or name") + stop_parser.add_argument( + "-y", "--yes", action="store_true", help="Do not ask for confirmation" + ) + stop_parser.set_defaults(subfunc=self._stop) + apply_parser = preset_subparsers.add_parser( "apply", - help="Apply an endpoint preset", + help="Apply a preset", formatter_class=self._parser.formatter_class, ) _add_configuration_args(apply_parser) register_profile_args(apply_parser) - apply_parser.add_argument("--preset", metavar="ID", help="The preset ID to use") + apply_parser.add_argument( + "--id", + action="append", + dest="preset_ids", + metavar="ID", + help="Deploy the best available preset among these IDs or names. Can be repeated", + ) apply_parser.add_argument( "-y", "--yes", action="store_true", help="Do not ask for confirmation" ) @@ -111,7 +164,7 @@ def _register(self) -> None: delete_parser = preset_subparsers.add_parser( "delete", - help="Delete endpoint presets", + help="Delete presets", formatter_class=self._parser.formatter_class, ) delete_target = delete_parser.add_mutually_exclusive_group(required=True) @@ -119,13 +172,18 @@ def _register(self) -> None: "preset", nargs="?", metavar="ID", - help="The preset ID", + help="The preset ID or name", ) delete_target.add_argument( - "--model", + "--base", metavar="MODEL", help="Delete all presets for a base model", ) + delete_target.add_argument( + "--repo", + metavar="REPO", + help="Delete all presets serving a model repo", + ) delete_parser.add_argument( "-y", "--yes", action="store_true", help="Do not ask for confirmation" ) @@ -140,33 +198,100 @@ def _command(self, args: argparse.Namespace) -> None: exit(0) def _list(self, args: argparse.Namespace) -> None: - presets = EndpointPresetStore().list() + base = getattr(args, "base", None) + repo = getattr(args, "repo", None) if getattr(args, "json", False): + presets = _filter_presets(EndpointPresetStore().list(), base=base, repo=repo) print(EndpointPresetListOutput(presets=presets).json()) return - print_endpoint_presets(presets, verbose=getattr(args, "verbose", False)) + verbose = getattr(args, "verbose", False) + while True: + presets = EndpointPresetStore().list() + sessions = list_agent_sessions() + if base or repo: + repo_to_base = {preset.model: preset.base for preset in presets} + presets = _filter_presets(presets, base=base, repo=repo) + sessions = [ + session + for session in sessions + if _session_matches_model( + session, base=base, repo=repo, repo_to_base=repo_to_base + ) + ] + if getattr(args, "watch", False): + console.clear() + print_endpoint_presets(presets, sessions=sessions, verbose=verbose) + if not getattr(args, "watch", False): + return + time.sleep(5) def _create(self, args: argparse.Namespace) -> None: - _, configuration = load_endpoint_configuration(args.configuration_file) - configuration = _get_effective_configuration(configuration, args) + configuration_path, configuration = load_endpoint_configuration(args.configuration_file) + configuration = _get_effective_configuration(configuration, args, require_name=False) + user_prompt = resolve_endpoint_prompt(configuration, configuration_path) + store = EndpointPresetStore() + resume_session = None + if getattr(args, "resume", None): + resume_session = load_resumable_agent_session(args.resume) + if getattr(args, "max_trials", None) is not None: + console.print( + "[warning]--max-trials is ignored when resuming: " + "session constraints are fixed at creation[/]" + ) + api = Client.from_config(project_name=args.project) + allowed_fleets = None + if resume_session is None: + allowed_fleets = plan_endpoint_preset(api=api, configuration=configuration) + if not _confirm_preset_creation(store, configuration.name, assume_yes=args.yes): + console.print("\nExiting...") + return result = create_endpoint_preset( - api=Client.from_config(project_name=args.project), + api=api, configuration=configuration, - store=EndpointPresetStore(), + store=store, keep_service=args.keep_service, debug=args.debug, + resume_session=resume_session, + user_prompt=user_prompt, + allowed_fleets=allowed_fleets, ) console.print( - f"Endpoint preset [code]{result.preset.id}[/] for " + f"Preset [code]{result.preset.id}[/] for " f"[code]{result.preset.base}[/] saved to [code]{result.path}[/]" ) if args.keep_service: console.print(f"Final service [code]{result.final_run_name}[/] kept running") + def _attach(self, args: argparse.Namespace) -> None: + result = attach_endpoint_preset( + api=Client.from_config(project_name=args.project), + store=EndpointPresetStore(), + preset_id=_resolve_session_ref(args.preset), + keep_service=args.keep_service, + ) + console.print( + f"Preset [code]{result.preset.id}[/] for " + f"[code]{result.preset.base}[/] saved to [code]{result.path}[/]" + ) + if args.keep_service: + console.print(f"Final service [code]{result.final_run_name}[/] kept running") + + def _stop(self, args: argparse.Namespace) -> None: + preset_id = _resolve_session_ref(args.preset) + if not args.yes and not confirm_ask( + f"Stop the preset creation session [code]{preset_id}[/]?" + ): + console.print("\nExiting...") + return + stop_endpoint_session( + Client.from_config(project_name=args.project), preset_id, assume_yes=args.yes + ) + def _get(self, args: argparse.Namespace) -> None: - preset = EndpointPresetStore().get(args.preset) + store = EndpointPresetStore() + preset = store.get(args.preset) or store.find_by_name(args.preset) if preset is None: - raise CLIError(f"Endpoint preset {args.preset!r} does not exist") + raise CLIError(f"Preset {args.preset!r} does not exist") print(preset.json()) def _apply(self, args: argparse.Namespace) -> None: @@ -176,7 +301,7 @@ def _apply(self, args: argparse.Namespace) -> None: api=Client.from_config(project_name=args.project), configuration=configuration, configuration_path=configuration_path, - preset_id=args.preset or configuration.preset, + preset_ids=args.preset_ids, profile_name=args.profile, command_args=args, store=EndpointPresetStore(), @@ -184,34 +309,35 @@ def _apply(self, args: argparse.Namespace) -> None: def _delete(self, args: argparse.Namespace) -> None: store = EndpointPresetStore() - preset = None if args.preset is not None: - preset = store.get(args.preset) + preset = store.get(args.preset) or store.find_by_name(args.preset) if preset is None: - raise CLIError(f"Endpoint preset {args.preset!r} does not exist") - message = f"Delete endpoint preset [code]{preset.id}[/] for [code]{preset.base}[/]?" + raise CLIError(f"Preset {args.preset!r} does not exist") + presets = [preset] + message = f"Delete preset [code]{preset.id}[/] for [code]{preset.base}[/]?" else: - presets = [preset for preset in store.list() if preset.base == args.model] + target = args.base or args.repo + presets = _filter_presets(store.list(), base=args.base, repo=args.repo) if not presets: - raise CLIError(f"No endpoint presets found for base model {args.model!r}") + kind = "base model" if args.base else "model repo" + raise CLIError(f"No presets found for {kind} {target!r}") message = ( - f"Delete {len(presets)} endpoint preset" - f"{'s' if len(presets) != 1 else ''} for [code]{args.model}[/]?" + f"Delete {len(presets)} preset" + f"{'s' if len(presets) != 1 else ''} for [code]{target}[/]?" ) if not args.yes and not confirm_ask(message): console.print("\nExiting...") return + for preset in presets: + store.delete(preset.id) if args.preset is not None: - assert preset is not None - store.delete(args.preset) console.print( - f"Endpoint preset [code]{preset.id}[/] for [code]{preset.base}[/] deleted" + f"Preset [code]{presets[0].id}[/] for [code]{presets[0].base}[/] deleted" ) else: - count = store.delete_for_base(args.model) console.print( - f"Deleted {count} endpoint preset{'s' if count != 1 else ''} " - f"for [code]{args.model}[/]" + f"Deleted {len(presets)} preset{'s' if len(presets) != 1 else ''} " + f"for [code]{args.base or args.repo}[/]" ) @@ -222,39 +348,137 @@ def _add_configuration_args(parser: argparse.ArgumentParser) -> None: required=True, metavar="FILE", dest="configuration_file", - help="The endpoint configuration file", + help="The preset configuration file", ).completer = FilesCompleter(allowednames=["*.yml", "*.yaml"]) # type: ignore[attr-defined] parser.add_argument( "-n", "--name", metavar="NAME", - help="The endpoint name. Required when the configuration omits name", + help="The service name. Required when the configuration omits name", ) def _add_list_args(parser: argparse.ArgumentParser) -> None: parser.add_argument("-v", "--verbose", action="store_true") + parser.add_argument( + "-w", + "--watch", + action="store_true", + help="Watch presets and sessions in realtime", + ) parser.add_argument( "--json", action="store_true", help="Output in JSON format", ) + model_filter = parser.add_mutually_exclusive_group() + model_filter.add_argument( + "--base", + metavar="MODEL", + help="Show only presets for a base model", + ) + model_filter.add_argument( + "--repo", + metavar="REPO", + help="Show only presets serving a model repo", + ) + + +def _filter_presets( + presets: list[EndpointPreset], + *, + base: str | None, + repo: str | None, +) -> list[EndpointPreset]: + return [ + preset + for preset in presets + if (base is None or preset.base == base) and (repo is None or preset.model == repo) + ] + +def _session_matches_model( + session: dict, + *, + base: str | None, + repo: str | None, + repo_to_base: dict[str, str], +) -> bool: + model = str(session.get("model") or "") + if not model: + return False + if repo is not None: + return model == repo + return model == base or repo_to_base.get(model) == base -def _apply_name(configuration: EndpointConfiguration, name: str | None) -> None: + +def _apply_name(configuration: EndpointConfiguration, name: str | None, *, required: bool) -> None: if name is not None: configuration.name = name if configuration.name is None: - raise CLIError("Endpoint name is required. Set `name` in the configuration or use --name") + if required: + raise CLIError( + "The service name is required. Set `name` in the configuration or use --name" + ) + return if not is_valid_dstack_resource_name(configuration.name): - raise CLIError("Endpoint name must match '^[a-z][a-z0-9-]{1,40}$'") + raise CLIError("The name must match '^[a-z][a-z0-9-]{1,40}$'") + + +def _resolve_session_ref(ref: str) -> str: + """A session reference may be a preset id or a claimed name.""" + if (get_presets_dir() / ref).is_dir(): + return ref + claims = find_session_name_claims(ref) + if len(claims) == 1: + return claims[0].preset_id + return ref + + +def _confirm_preset_creation( + store: EndpointPresetStore, name: str | None, *, assume_yes: bool +) -> bool: + """One apply-style confirmation; reassigns the name from any holder on yes.""" + preset_holder = None + session_holders = [] + if name is not None: + preset_holder = store.find_by_name(name) + session_holders = [ + session + for session in find_session_name_claims(name) + if preset_holder is None or session.preset_id != preset_holder.id + ] + holders = [] + if preset_holder is not None: + holders.append(f"preset [code]{preset_holder.id}[/]") + holders.extend(f"session [code]{session.preset_id}[/]" for session in session_holders) + if holders: + message = ( + f"The name [code]{name}[/] is already used by {', '.join(holders)}." + f" Reassign it to a new preset?" + ) + elif name is not None: + message = f"Create the preset [code]{name}[/]?" + else: + message = "Create the preset?" + if not assume_yes and not confirm_ask(message): + return False + if preset_holder is not None and name is not None: + store.release_name(name) + for session in session_holders: + session.update_manifest(name=None) + return True def _get_effective_configuration( configuration: EndpointConfiguration, args: argparse.Namespace, + *, + require_name: bool = True, ) -> EndpointConfiguration: - _apply_name(configuration, args.name) + _apply_name(configuration, args.name, required=require_name) + if getattr(args, "max_trials", None) is not None: + configuration.max_trials = args.max_trials profile = load_profile(Path.cwd(), args.profile) for field in ProfileParams.__fields__: if getattr(configuration, field) is None: diff --git a/src/dstack/_internal/cli/models/endpoint_agent.py b/src/dstack/_internal/cli/models/endpoint_agent.py index c31a1c854..05dad6602 100644 --- a/src/dstack/_internal/cli/models/endpoint_agent.py +++ b/src/dstack/_internal/cli/models/endpoint_agent.py @@ -1,5 +1,5 @@ import uuid -from typing import Optional +from typing import Any, Dict, Optional from pydantic import PositiveInt, root_validator @@ -118,3 +118,23 @@ def validate_report(cls, values: dict) -> dict: elif not values.get("failure_summary"): raise ValueError("failed agent report must include failure_summary") return values + + +class EndpointAgentInfo(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(EndpointAgentInfo): + """Claude agent runtime information, saved as `agent.json`.""" + + model: ClaudeModelParams + auth: Dict[str, Any] diff --git a/src/dstack/_internal/cli/models/endpoint_presets.py b/src/dstack/_internal/cli/models/endpoint_presets.py index aea33f499..75aa778e8 100644 --- a/src/dstack/_internal/cli/models/endpoint_presets.py +++ b/src/dstack/_internal/cli/models/endpoint_presets.py @@ -70,6 +70,10 @@ def validate_command_has_no_bearer_token(cls, value: str) -> str: token = match.group(1) if token.startswith("$") or "redacted" in token.lower() or set(token) == {"*"}: continue + # Prose such as "auth via bearer header from env" is not a + # credential: only credential-shaped values are rejected. + if len(token) < 16 or not any(char.isdigit() for char in token): + continue raise ValueError("command must not contain a bearer token value") return value @@ -77,6 +81,7 @@ def validate_command_has_no_bearer_token(cls, value: str) -> str: def validate_metrics(cls, values: dict) -> dict: metrics = values.get("metrics") workload = values.get("workload") + assert metrics is not None and workload is not None if metrics.failed_requests != 0: raise ValueError("benchmark must not include failed requests") if metrics.successful_requests != workload.num_requests: @@ -99,6 +104,8 @@ class EndpointPreset(CoreModel): base: str """Base model used for local preset lookup.""" id: str + name: Optional[str] = None + """Mutable human name; at most one preset or in-flight session holds it.""" model: str """Exact repo/path loaded by the service command.""" context_length: PositiveInt diff --git a/src/dstack/_internal/cli/models/endpoint_trials.py b/src/dstack/_internal/cli/models/endpoint_trials.py new file mode 100644 index 000000000..1dd137553 --- /dev/null +++ b/src/dstack/_internal/cli/models/endpoint_trials.py @@ -0,0 +1,16 @@ +from typing import Optional + +from dstack._internal.cli.models.endpoint_presets import EndpointBenchmark +from dstack._internal.core.models.common import CoreModel +from dstack._internal.core.models.configurations import TaskConfiguration +from dstack._internal.core.models.resources import ResourcesSpec + + +class EndpointPresetTrial(CoreModel): + # TODO: a single task (= single node) for now; revisit multi-task trials + # (P/D disaggregation) once tasks support node groups. + task: TaskConfiguration + resources: ResourcesSpec + """Exact instance resources the task ran on, as in preset validations.""" + benchmark: Optional[EndpointBenchmark] = None + """Null only for a failed trial (the configuration never served).""" diff --git a/src/dstack/_internal/cli/models/endpoints.py b/src/dstack/_internal/cli/models/endpoints.py index 8b2cb409e..03762c41b 100644 --- a/src/dstack/_internal/cli/models/endpoints.py +++ b/src/dstack/_internal/cli/models/endpoints.py @@ -1,6 +1,6 @@ from typing import Annotated, Any, Literal, Optional, Union -from pydantic import Field, PositiveInt, validator +from pydantic import Field, PositiveInt, root_validator, validator from dstack._internal.core.models.common import ( CoreModel, @@ -11,6 +11,9 @@ from dstack._internal.core.models.profiles import ProfileParams, ProfileParamsConfig from dstack._internal.utils.json_schema import add_extra_schema_types +DEFAULT_MAX_TRIALS = 3 +DEFAULT_CONCURRENCY = 8 + class EndpointModelRepo(CoreModel): repo: Annotated[str, Field(description="The exact model repo or path to deploy")] @@ -66,6 +69,21 @@ def validate_base(cls, value: str) -> str: EndpointModelSpec = Union[EndpointModelRepo, EndpointModelBase] +MAX_PROMPT_LENGTH = 10_000 + + +class EndpointPromptFile(CoreModel): + path: Annotated[ + str, + Field(description="The path to a prompt file, relative to the configuration file"), + ] + + @validator("path") + def validate_path(cls, value: str) -> str: + if not value.strip(): + raise ValueError("Prompt path must be a non-empty string") + return value + class EndpointConfigurationConfig(ProfileParamsConfig): @staticmethod @@ -81,25 +99,62 @@ class EndpointConfiguration( ProfileParams, generate_dual_core_model(EndpointConfigurationConfig), ): - type: Annotated[Literal["endpoint"], Field(description="The configuration type")] = "endpoint" + type: Annotated[Literal["preset"], Field(description="The configuration type")] = "preset" name: Annotated[ Optional[str], - Field(description="The endpoint name. Required unless passed with `--name`"), + Field(description="The service name. Required unless passed with `--name`"), ] = None model: Annotated[ EndpointModelSpec, Field( description=( "The model to serve. Use a string or `repo` for an exact repo/path, " - "or `base` to allow compatible model variants." + "or `base` to allow compatible model variants. " + "Prefer the top-level `base`/`repo` shorthand unless a custom " + "client-facing model name is needed" ) ), ] + base: Annotated[ + Optional[str], + Field( + description=( + "The base model repo; compatible variants are allowed. Shorthand for `model.base`" + ) + ), + ] = None + repo: Annotated[ + Optional[str], + Field(description="The exact model repo/path to serve. Shorthand for `model.repo`"), + ] = None + prompt: Annotated[ + Optional[Union[str, EndpointPromptFile]], + Field( + description=( + "Additional instructions for the preset creation agent, inline or as a file `path`" + ) + ), + ] = None context_length: Annotated[ Optional[PositiveInt], Field(description="The minimum required context length") ] = None - preset: Annotated[ - Optional[str], Field(description="The preset ID to use when applying the endpoint") + max_trials: Annotated[ + Optional[PositiveInt], + Field( + description=( + "The maximum number of benchmarked trials during preset creation" + f" before the best one is promoted. Defaults to `{DEFAULT_MAX_TRIALS}`" + ) + ), + ] = None + concurrency: Annotated[ + Optional[PositiveInt], + Field( + description=( + "The number of simultaneous requests used for benchmarks during" + f" preset creation. Defaults to `{DEFAULT_CONCURRENCY}`" + ) + ), ] = None gateway: Annotated[ Optional[Union[bool, EntityReference, str]], @@ -115,19 +170,59 @@ class EndpointConfiguration( Env() ) + @property + def effective_max_trials(self) -> int: + return self.max_trials if self.max_trials is not None else DEFAULT_MAX_TRIALS + + @property + def effective_concurrency(self) -> int: + return self.concurrency if self.concurrency is not None else DEFAULT_CONCURRENCY + + @root_validator(pre=True) + def apply_model_shorthand(cls, values: Any) -> Any: + if not isinstance(values, dict): + return values + base, repo = values.get("base"), values.get("repo") + if base and repo: + raise ValueError("`base` and `repo` are mutually exclusive") + if base or repo: + if values.get("model") is not None: + raise ValueError("`model` cannot be combined with the `base`/`repo` shorthand") + values = dict(values) + values.pop("base", None) + values.pop("repo", None) + values["model"] = {"base": base} if base else {"repo": repo} + return values + @validator("model", pre=True) def parse_model(cls, value: Any) -> Any: if isinstance(value, str): return {"repo": _validate_model(value, field="model")} return value - @validator("preset") - def validate_preset(cls, value: Optional[str]) -> Optional[str]: - if value is not None and not value.strip(): - raise ValueError("Endpoint preset must be a non-empty string") + @validator("prompt") + def validate_prompt(cls, value: Any) -> Any: + if isinstance(value, str): + if not value.strip(): + raise ValueError("Prompt must be a non-empty string") + if len(value) > MAX_PROMPT_LENGTH: + raise ValueError(f"Prompt must be at most {MAX_PROMPT_LENGTH} characters") return value +class EndpointPresetConstraints(CoreModel): + """The effective constraints for endpoint preset creation, saved as `constraints.json` + in the agent workspace. Field semantics are documented in the agent system prompt.""" + + run_name_prefix: str + model: EndpointModelSpec + context_length: Optional[PositiveInt] = None + max_trials: PositiveInt + concurrency: PositiveInt + fleets: list[str] = Field(min_items=1) + env: list[str] = [] + + def _validate_model(value: Any, *, field: str) -> str: if not isinstance(value, str) or not value.strip(): raise ValueError(f"Endpoint model {field} must be a non-empty string") diff --git a/src/dstack/_internal/cli/services/endpoints/agent.py b/src/dstack/_internal/cli/services/endpoints/agent.py index 546907b96..0c97fdddf 100644 --- a/src/dstack/_internal/cli/services/endpoints/agent.py +++ b/src/dstack/_internal/cli/services/endpoints/agent.py @@ -1,22 +1,27 @@ import asyncio import json import os +import secrets import shutil import signal import subprocess import sys import tempfile -from contextlib import contextmanager, suppress +import time +from contextlib import asynccontextmanager, suppress from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path -from typing import Any, Iterator, Optional, Sequence +from typing import Any, AsyncIterator, Callable, Optional, Protocol, Sequence import psutil import yaml from rich.text import Text -from dstack._internal.cli.models.endpoint_agent import AGENT_FINAL_REPORT_JSON_SCHEMA +from dstack._internal.cli.models.endpoint_agent import ( + AGENT_FINAL_REPORT_JSON_SCHEMA, + ClaudeAgentInfo, +) from dstack._internal.cli.models.endpoints import EndpointConfiguration from dstack._internal.cli.utils.common import console from dstack._internal.compat import IS_WINDOWS @@ -27,13 +32,22 @@ _SKILL_NAMES = ("dstack", "dstack-prototyping") _PROGRESS_FILENAME = "progress.jsonl" -_SUBMISSIONS_FILENAME = "submissions.jsonl" +_RUNS_FILENAME = "runs.jsonl" +_TRIALS_FILENAME = "trials.jsonl" +_CONSTRAINTS_FILENAME = "constraints.json" _FINAL_REPORT_FILENAME = "final_report.json" +_SESSION_FILENAME = "session.json" +_USER_PROMPT_FILENAME = "user_prompt.md" _PROGRESS_ENV = "DSTACK_ENDPOINT_PROGRESS_LOG" _REDACTION = "[redacted]" _CLAUDE_TOOLS = "Bash,Read,Write,Edit,WebFetch,WebSearch,StructuredOutput" _CLAUDE_EFFORT_LEVELS = ("low", "medium", "high", "xhigh", "max") -_CLAUDE_STREAM_LIMIT = 16 * 1024 * 1024 +_RESUME_DELAYS_SECONDS: tuple[int, ...] = (30, 60, 120) +_RESUME_PROMPT = ( + "The previous agent process was interrupted. Continue where you left off. " + "Re-check the states of your runs before relying on them: time may have " + "passed, and tasks or instances may have stopped in the meantime." +) _MAX_RUN_NAME_LENGTH = 41 _MAX_UNIX_SOCKET_PATH_BYTES = 103 _INHERITED_ENV_NAMES = ( @@ -91,19 +105,36 @@ def progress_path(self) -> Path: return self.path / _PROGRESS_FILENAME @property - def submissions_path(self) -> Path: - return self.path / _SUBMISSIONS_FILENAME + def runs_path(self) -> Path: + return self.path / _RUNS_FILENAME + + @property + def trials_path(self) -> Path: + return self.path / _TRIALS_FILENAME + + @property + def constraints_path(self) -> Path: + return self.path / _CONSTRAINTS_FILENAME @property def final_report_path(self) -> Path: return self.path / _FINAL_REPORT_FILENAME + @property + def agent_stdout_path(self) -> Path: + return self.path / ".agent-stdout.jsonl" + + @property + def agent_stderr_path(self) -> Path: + return self.path / ".agent-stderr.log" + @dataclass class EndpointAgentSession: path: Path timestamp: str debug: bool + preset_id: str = "" _log_enabled: bool = field(default=True, init=False, repr=False) @property @@ -114,9 +145,49 @@ def log_path(self) -> Path: def trace_path(self) -> Path: return self.path / "trace.jsonl" + @property + def runs_path(self) -> Path: + return self.path / _RUNS_FILENAME + + @property + def trials_path(self) -> Path: + return self.path / _TRIALS_FILENAME + def write_prompt(self, prompt: str) -> None: _write_private_text(self.path / "prompt.md", prompt + "\n") + def write_user_prompt(self, user_prompt: str) -> None: + _write_private_text(self.path / _USER_PROMPT_FILENAME, user_prompt + "\n") + + def read_user_prompt(self) -> Optional[str]: + try: + text = (self.path / _USER_PROMPT_FILENAME).read_text(encoding="utf-8").strip() + except OSError: + return None + return text or None + + def write_constraints(self, constraints_text: str) -> None: + _write_private_text(self.path / _CONSTRAINTS_FILENAME, constraints_text) + + def write_final_report(self, report_text: str) -> None: + _write_private_text(self.path / _FINAL_REPORT_FILENAME, report_text) + + def write_agent_info(self, auth: "ClaudeAuth") -> None: + info = ClaudeAgentInfo.parse_obj( + { + "executable": auth.executable, + "version": _get_claude_version(auth), + "model": { + "name": auth.model, + "effort": auth.effort or "default", + }, + "auth": _get_claude_auth_status(auth), + } + ) + _write_private_text( + self.path / "agent.json", json.dumps(json.loads(info.json()), indent=2) + "\n" + ) + def append_log(self, line: str) -> None: if not self._log_enabled: return @@ -128,22 +199,24 @@ def append_log(self, line: str) -> None: self._log_enabled = False console.print(f"[warning]Could not write agent log {self.log_path}: {e}[/]") - def finish(self, preset_id: Optional[str] = None) -> Path: - suffix = preset_id or "failed" - name = f"{self.timestamp}-{suffix}" - index = 0 - while True: - target = self.path.with_name(name if index == 0 else f"{name}-{index}") - if target.exists(): - index += 1 - continue - try: - self.path.rename(target) - except FileExistsError: - index += 1 - continue - self.path = target - return target + def read_manifest(self) -> dict[str, Any]: + try: + manifest = json.loads((self.path / _SESSION_FILENAME).read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return {} + return manifest if isinstance(manifest, dict) else {} + + def update_manifest(self, **fields: Any) -> None: + manifest = self.read_manifest() + manifest.update(fields) + _write_private_text(self.path / _SESSION_FILENAME, json.dumps(manifest, indent=2) + "\n") + + def record_claude_session_id(self, session_id: str) -> None: + self.update_manifest(claude_session_id=session_id) + + def finish(self, status: str) -> Path: + self.update_manifest(status=status) + return self.path @dataclass(frozen=True) @@ -158,16 +231,84 @@ class ClaudeAuth: class EndpointAgentProcessOutput: report_data: Optional[dict[str, Any]] = None error: Optional[str] = None + session_id: Optional[str] = None + made_progress: bool = False -@contextmanager -def endpoint_agent_workspace() -> Iterator[EndpointAgentWorkspace]: - with tempfile.TemporaryDirectory(prefix="dpe-", dir=_get_short_temp_dir()) as directory: - root = Path(directory) - _validate_control_socket_path(root) - workspace = EndpointAgentWorkspace(path=root / "w", dstack_home=root / "h") +def create_agent_workspace(session: EndpointAgentSession) -> EndpointAgentWorkspace: + real = session.path / "workspace" + try: + real.mkdir(mode=0o700) + if IS_WINDOWS: + # Windows has no Unix-socket path limit and symlinks require + # privileges, so the workspace is used directly. + alias = real + else: + alias = _create_workspace_alias(real) + _validate_control_socket_path(alias) + workspace = EndpointAgentWorkspace(path=alias / "w", dstack_home=alias / "h") _prepare_workspace(workspace) - yield workspace + except OSError as e: + raise CLIError(f"Could not create the agent workspace under {real}: {e}") from e + session.update_manifest(workspace=str(real), alias=str(alias)) + return workspace + + +def attach_agent_workspace(session: EndpointAgentSession) -> EndpointAgentWorkspace: + manifest = session.read_manifest() + real_value, alias_value = manifest.get("workspace"), manifest.get("alias") + if not real_value or not alias_value: + raise CLIError("The session has no workspace to resume") + real, alias = Path(real_value), Path(alias_value) + if not real.is_dir(): + raise CLIError( + f"The session workspace no longer exists: {real}. " + "Stop any leftover runs manually and start a new session." + ) + if alias != real: + _ensure_workspace_alias(alias, real) + return EndpointAgentWorkspace(path=alias / "w", dstack_home=alias / "h") + + +def remove_agent_workspace(session: EndpointAgentSession) -> None: + manifest = session.read_manifest() + alias = manifest.get("alias") + workspace = manifest.get("workspace") + if alias and alias != workspace and Path(alias).is_symlink(): + with suppress(OSError): + os.unlink(alias) + if workspace: + shutil.rmtree(workspace, ignore_errors=True) + + +def _create_workspace_alias(real: Path) -> Path: + base = _get_short_temp_dir() or tempfile.gettempdir() + while True: + alias = Path(base) / f"dpe-{secrets.token_hex(4)}" + try: + os.symlink(real, alias) + except FileExistsError: + continue + return alias + + +def _ensure_workspace_alias(alias: Path, real: Path) -> None: + if os.path.lexists(alias): + if ( + alias.is_symlink() + and os.readlink(alias) == str(real) + and (IS_WINDOWS or os.lstat(alias).st_uid == os.getuid()) + ): + return + raise CLIError( + f"The workspace alias path cannot be used safely: {alias}. " + "Remove it manually if it is yours, or start a new session." + ) + os.symlink(real, alias) + + +def get_presets_dir() -> Path: + return get_dstack_dir() / "presets" def create_endpoint_agent_session( @@ -176,44 +317,249 @@ def create_endpoint_agent_session( debug: bool = False, ) -> EndpointAgentSession: if configuration.name is None: - raise CLIError("Endpoint name is required to save agent output") + raise CLIError("The service name is required to save agent output") timestamp = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S-%fZ") - parent = get_dstack_dir() / "agent" / configuration.name + parent = get_presets_dir() path: Optional[Path] = None try: parent.mkdir(mode=0o700, parents=True, exist_ok=True) - path = _create_agent_session_directory(parent, timestamp) + while True: + preset_id = secrets.token_hex(4) + path = parent / preset_id + try: + path.mkdir(mode=0o700) + except FileExistsError: + continue + break _write_private_text(path / "agent.log", "") + manifest = { + "id": preset_id, + "status": "running", + "pid": os.getpid(), + "endpoint": configuration.name, + "name": configuration.name, + "model": getattr(configuration.model, "base", None) + or getattr(configuration.model, "repo", None), + "max_trials": configuration.effective_max_trials, + "created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"), + "debug": debug, + } + _write_private_text(path / _SESSION_FILENAME, json.dumps(manifest, indent=2) + "\n") + data = json.loads(configuration.json(exclude_none=True)) + if configuration.env: + data["env"] = list(configuration.env) + else: + data.pop("env", None) + _write_private_text( + path / "preset.dstack.yml", + yaml.safe_dump(data, sort_keys=False), + ) if debug: - data = json.loads(configuration.json(exclude_none=True)) - if configuration.env: - data["env"] = list(configuration.env) - else: - data.pop("env", None) - _write_private_text( - path / "endpoint.dstack.yml", - yaml.safe_dump(data, sort_keys=False), - ) _write_private_text(path / "trace.jsonl", "") except OSError as e: if path is not None: shutil.rmtree(path, ignore_errors=True) raise CLIError(f"Could not create agent output under {parent}: {e}") from e assert path is not None - return EndpointAgentSession(path=path, timestamp=timestamp, debug=debug) + return EndpointAgentSession(path=path, timestamp=timestamp, debug=debug, preset_id=preset_id) -def _create_agent_session_directory(parent: Path, timestamp: str) -> Path: - index = 0 - while True: - name = f"{timestamp}-running" if index == 0 else f"{timestamp}-{index}-running" - path = parent / name +def _get_claude_version(auth: "ClaudeAuth") -> Optional[str]: + try: + result = subprocess.run( + [auth.executable, "--version"], + capture_output=True, + text=True, + timeout=15, + ) + return result.stdout.strip() or None + except (OSError, subprocess.SubprocessError): + return None + + +def _get_claude_auth_status(auth: "ClaudeAuth") -> dict[str, Any]: + if auth.api_key: + return {"authMethod": "api-key"} + try: + result = subprocess.run( + [auth.executable, "auth", "status", "--json"], + capture_output=True, + text=True, + timeout=15, + ) + status = json.loads(result.stdout) + if isinstance(status, dict): + return status + except (OSError, subprocess.SubprocessError, json.JSONDecodeError): + pass + return {"authMethod": "unknown"} + + +def load_resumable_agent_session(preset_id: str) -> EndpointAgentSession: + path = get_presets_dir() / preset_id + session = EndpointAgentSession(path=path, timestamp="", debug=False, preset_id=preset_id) + manifest = session.read_manifest() + if not path.is_dir() or not manifest: + raise CLIError(f"Unknown preset creation session: {preset_id}") + status = manifest.get("status") + if status == "success": + raise CLIError(f"Session {preset_id} completed successfully; nothing to resume") + if status == "failed": + raise CLIError(f"Session {preset_id} failed and cannot be resumed") + if status == "running" and session_process_alive(manifest): + raise CLIError( + f"Session {preset_id} appears to be running already;" + f" attach with [code]dstack preset attach {preset_id}[/]" + ) + if not manifest.get("claude_session_id"): + raise CLIError( + f"Session {preset_id} stopped before the agent started; start a new session" + ) + session.debug = bool(manifest.get("debug")) + session.timestamp = str(manifest.get("created_at") or "") + return session + + +def _pid_alive(pid: Any, started_at: Any = None) -> bool: + if not isinstance(pid, int) or pid <= 0 or not psutil.pid_exists(pid): + return False + if isinstance(started_at, (int, float)): + create_time = _process_started_at(pid) + # A recycled pid has a different start time. + if create_time is not None and abs(create_time - started_at) > 1.0: + return False + return True + + +def session_process_alive(manifest: dict[str, Any]) -> bool: + """Whether the session is still worked on: a live agent (possibly + detached) or a live CLI (possibly between agent retries).""" + if _pid_alive(manifest.get("agent_pid"), manifest.get("agent_started_at")): + return True + pid = manifest.get("pid") + return isinstance(pid, int) and pid > 0 and pid != os.getpid() and psutil.pid_exists(pid) + + +def load_attachable_agent_session(preset_id: str) -> EndpointAgentSession: + path = get_presets_dir() / preset_id + session = EndpointAgentSession(path=path, timestamp="", debug=False, preset_id=preset_id) + manifest = session.read_manifest() + if not path.is_dir() or not manifest: + raise CLIError(f"Unknown preset creation session: {preset_id}") + status = manifest.get("status") + if status == "success": + raise CLIError(f"Session {preset_id} completed successfully; nothing to attach") + if status == "failed": + raise CLIError(f"Session {preset_id} failed and cannot be attached") + if status == "interrupted": + raise CLIError( + f"Session {preset_id} is interrupted; resume it with" + f" [code]dstack preset create -f --resume {preset_id}[/]" + ) + pid = manifest.get("pid") + if isinstance(pid, int) and pid > 0 and pid != os.getpid() and psutil.pid_exists(pid): + raise CLIError( + f"Session {preset_id} is already attached by another CLI (pid {pid});" + f" stop or detach it there with Ctrl+C" + ) + session.debug = bool(manifest.get("debug")) + session.timestamp = str(manifest.get("created_at") or "") + return session + + +def claimed_session_name(manifest: dict[str, Any]) -> Optional[str]: + """The name this session holds; pre-claim manifests fall back to the endpoint.""" + if "name" in manifest: + value = manifest.get("name") + return value if isinstance(value, str) and value else None + value = manifest.get("endpoint") + return value if isinstance(value, str) and value else None + + +def find_session_name_claims(name: str) -> list[EndpointAgentSession]: + """Sessions of any status holding `name`, including failed ones.""" + root = get_presets_dir() + claims = [] + if not root.is_dir(): + return claims + for path in sorted(root.iterdir()): + if not path.is_dir() or path.name.startswith((".", "models--")): + continue + session = EndpointAgentSession(path=path, timestamp="", debug=False, preset_id=path.name) + if claimed_session_name(session.read_manifest()) == name: + claims.append(session) + return claims + + +def list_agent_sessions() -> list[dict[str, Any]]: + root = get_presets_dir() + entries = [] + candidates = sorted(root.iterdir()) if root.is_dir() else [] + for path in candidates: + if not path.is_dir() or path.name.startswith((".", "models--")): + continue + session = EndpointAgentSession(path=path, timestamp="", debug=False, preset_id=path.name) + manifest = session.read_manifest() + status = manifest.get("status") + if status not in ("running", "interrupted", "success"): + continue + if status == "running" and not session_process_alive(manifest): + status = "interrupted" + entry = dict(manifest) + entry["id"] = path.name + entry["name"] = claimed_session_name(manifest) + entry["status"] = status + entry["trials"] = _summarize_session_trials(path / _TRIALS_FILENAME) + entries.append(entry) + return entries + + +def _summarize_session_trials(path: Path) -> Optional[dict[str, Any]]: + """Best-so-far summary from a session's mirrored trial records.""" + try: + lines = path.read_text(encoding="utf-8").splitlines() + except OSError: + lines = [] + count = 0 + best: Optional[dict[str, Any]] = None + for line in lines: try: - path.mkdir(mode=0o700) - except FileExistsError: - index += 1 + record = json.loads(line) + except json.JSONDecodeError: + continue + if not isinstance(record, dict): + continue + # One record per trial (the agent contract); trials may share a task, + # so task names must not be deduplicated. + count += 1 + benchmark = record.get("benchmark") + if not isinstance(benchmark, dict): + continue + metrics = benchmark.get("metrics") or {} + workload = benchmark.get("workload") or {} + duration = metrics.get("duration_seconds") + tokens = metrics.get("total_output_tokens") + if not isinstance(duration, (int, float)) or duration <= 0: continue - return path + if not isinstance(tokens, (int, float)): + continue + tok_s = tokens / duration + if best is None or tok_s > best["tok_s"]: + resources = record.get("resources") or {} + gpu = resources.get("gpu") if isinstance(resources, dict) else None + gpu_text = None + if isinstance(gpu, dict) and gpu.get("name"): + gpu_text = str(gpu["name"]) + if gpu.get("memory"): + gpu_text += f":{gpu['memory']}" + if gpu.get("count"): + gpu_text += f":{gpu['count']}" + best = { + "tok_s": tok_s, + "concurrency": workload.get("concurrency"), + "gpu": gpu_text, + } + return {"count": count, "best": best} def get_claude_auth() -> ClaudeAuth: @@ -261,8 +607,6 @@ def build_endpoint_agent_env( env["DSTACK_SERVER_URL"] = api.client.base_url env["DSTACK_PROJECT"] = api.project env["DSTACK_TOKEN"] = token - env["DSTACK_ENDPOINT_SERVER_URL"] = api.client.base_url - env["DSTACK_ENDPOINT_BEARER_TOKEN"] = token env[_PROGRESS_ENV] = str(workspace.progress_path) for name in ["TMPDIR", "TEMP", "TMP"]: env[name] = str(workspace.temp_path) @@ -286,74 +630,114 @@ async def run_endpoint_agent( auth: ClaudeAuth, redacted_values: Sequence[str], agent_session: EndpointAgentSession, + initial_resume_session_id: Optional[str] = None, ) -> EndpointAgentProcessOutput: - command = _prepare_subprocess_command(_build_claude_command(auth=auth)) - progress_tailer = _ProgressTailer( - path=workspace.progress_path, - redacted_values=redacted_values, - agent_session=agent_session, - ) - progress_task = asyncio.create_task(progress_tailer.run()) + async with _session_tailers( + workspace=workspace, agent_session=agent_session, redacted_values=redacted_values + ): + resume_session_id: Optional[str] = initial_resume_session_id + attempt_prompt = prompt if resume_session_id is None else _RESUME_PROMPT + retry_delays = list(_RESUME_DELAYS_SECONDS) + while True: + command = _prepare_subprocess_command( + _build_claude_command(auth=auth, resume_session_id=resume_session_id) + ) + output, returncode = await _run_claude_process( + command=command, + prompt=attempt_prompt, + env=env, + workspace=workspace, + redacted_values=redacted_values, + agent_session=agent_session, + ) + if output.report_data is None and returncode != 0: + output.error = output.error or f"Claude exited with return code {returncode}" + # Retry any process death without a submitted report; a terminal + # failure report from the agent returns immediately. + if output.report_data is not None or output.error is None: + return output + # A failed attempt that produced agent work is a new outage, not a + # continuation of the previous one: restore the full retry budget. + # Attempts that fail without any work drain it, so the loop always + # terminates when the network stays down. + if output.made_progress: + retry_delays = list(_RESUME_DELAYS_SECONDS) + if not retry_delays: + return output + delay = retry_delays.pop(0) + session_id = output.session_id or resume_session_id + if session_id is not None: + resume_session_id = session_id + attempt_prompt = _RESUME_PROMPT + action = "resuming the session" + else: + action = "retrying" + print_endpoint_progress( + f"Agent process exited without a report; {action} in {delay}s.", + agent_session=agent_session, + ) + await asyncio.sleep(delay) + + +async def _run_claude_process( + *, + command: list[str], + prompt: str, + env: dict[str, str], + workspace: EndpointAgentWorkspace, + redacted_values: Sequence[str], + agent_session: EndpointAgentSession, +) -> tuple[EndpointAgentProcessOutput, int]: proc: Optional[asyncio.subprocess.Process] = None try: - proc = await asyncio.create_subprocess_exec( - *command, - cwd=workspace.path, - env=env, - stdin=asyncio.subprocess.PIPE, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - start_new_session=not IS_WINDOWS, - limit=_CLAUDE_STREAM_LIMIT, + # The agent's streams go to workspace files rather than pipes, so the + # agent survives CLI death (detach) and a later attach can continue + # parsing from the persisted offsets. + with ( + workspace.agent_stdout_path.open("ab") as stdout_file, + workspace.agent_stderr_path.open("ab") as stderr_file, + ): + proc = await asyncio.create_subprocess_exec( + *command, + cwd=workspace.path, + env=env, + stdin=asyncio.subprocess.PIPE, + stdout=stdout_file, + stderr=stderr_file, + start_new_session=not IS_WINDOWS, + ) + agent_session.update_manifest( + agent_pid=proc.pid, agent_started_at=_process_started_at(proc.pid) ) assert proc.stdin is not None - assert proc.stdout is not None - assert proc.stderr is not None proc.stdin.write(prompt.encode()) with suppress(BrokenPipeError, ConnectionResetError): await proc.stdin.drain() proc.stdin.close() - stdout_task = asyncio.create_task( - _read_process_stream( - stream=proc.stdout, - stream_name="stdout", - parse_result=True, - redacted_values=redacted_values, + + def agent_alive() -> bool: + return proc.returncode is None + + collect_task = asyncio.create_task( + _collect_agent_output( + workspace=workspace, agent_session=agent_session, - ) - ) - stderr_task = asyncio.create_task( - _read_process_stream( - stream=proc.stderr, - stream_name="stderr", - parse_result=False, redacted_values=redacted_values, - agent_session=agent_session, + is_alive=agent_alive, ) ) - stdout_output, stderr_output, returncode = await asyncio.gather( - stdout_task, - stderr_task, - proc.wait(), - ) + returncode = await proc.wait() + output = await collect_task + except (KeyboardInterrupt, asyncio.CancelledError): + # The stop-or-detach decision belongs to the interrupt handler; the + # agent must stay alive here in case the user detaches. + raise except BaseException: if proc is not None and proc.returncode is None: await _terminate_process(proc) raise - finally: - progress_task.cancel() - with suppress(asyncio.CancelledError): - await progress_task - progress_tailer.flush() - output = stdout_output - if output.report_data is None: - output.report_data = stderr_output.report_data - if output.error is None: - output.error = stderr_output.error - if output.report_data is None and returncode != 0: - output.error = output.error or f"Claude exited with return code {returncode}" - return output + return output, returncode def print_endpoint_progress(message: str, *, agent_session: EndpointAgentSession) -> None: @@ -412,7 +796,8 @@ def _prepare_workspace(workspace: EndpointAgentWorkspace) -> None: workspace.temp_path.mkdir(mode=0o700) for path in [ workspace.progress_path, - workspace.submissions_path, + workspace.runs_path, + workspace.trials_path, ]: path.touch() workspace.bin_path.mkdir() @@ -445,7 +830,7 @@ def _install_home_wrapper(bin_dir: Path, command: str, home: Path) -> None: script = f"""#!{sys.executable} import sys -print("Endpoint preset creation could not find the {command} executable.", file=sys.stderr) +print("Preset creation could not find the {command} executable.", file=sys.stderr) raise SystemExit(127) """ else: @@ -512,7 +897,7 @@ def _install_skills(workspace: Path) -> None: for skill_name in _SKILL_NAMES: source = source_dir / skill_name if not (source / "SKILL.md").is_file(): - raise CLIError(f"Missing endpoint agent skill: {skill_name}") + raise CLIError(f"Missing preset agent skill: {skill_name}") shutil.copytree(source, target_dir / skill_name) @@ -528,7 +913,9 @@ def _get_skills_dir() -> Path: raise CLIError("Could not find packaged dstack skills") -def _build_claude_command(*, auth: ClaudeAuth) -> list[str]: +def _build_claude_command( + *, auth: ClaudeAuth, resume_session_id: Optional[str] = None +) -> list[str]: command = [ auth.executable, "-p", @@ -554,6 +941,8 @@ def _build_claude_command(*, auth: ClaudeAuth) -> list[str]: command[2:2] = ["--bare"] if auth.effort is not None: command[2:2] = ["--effort", auth.effort] + if resume_session_id is not None: + command += ["--resume", resume_session_id] return command @@ -610,9 +999,124 @@ def _redact_trace_value(value: Any, redacted_values: Sequence[str]) -> Any: return value +@asynccontextmanager +async def _session_tailers( + *, + workspace: EndpointAgentWorkspace, + agent_session: EndpointAgentSession, + redacted_values: Sequence[str], +) -> AsyncIterator[None]: + """Mirrors the session's progress and record files while the body runs.""" + offset_store = _OffsetStore(agent_session.path / ".offsets.json") + progress_tailer = _ProgressTailer( + path=workspace.progress_path, + redacted_values=redacted_values, + agent_session=agent_session, + offset_store=offset_store, + ) + record_mirrors = [ + _RecordMirror( + source=workspace.runs_path, + target=agent_session.runs_path, + redacted_values=redacted_values, + offset_store=offset_store, + offset_key="runs", + ), + _RecordMirror( + source=workspace.trials_path, + target=agent_session.trials_path, + redacted_values=redacted_values, + offset_store=offset_store, + offset_key="trials", + ), + ] + tailer_tasks = [ + asyncio.create_task(tailer.run()) for tailer in [progress_tailer, *record_mirrors] + ] + try: + yield + finally: + for task in tailer_tasks: + task.cancel() + for task in tailer_tasks: + with suppress(asyncio.CancelledError): + await task + progress_tailer.flush() + for mirror in record_mirrors: + mirror.flush() + + +async def _collect_agent_output( + *, + workspace: EndpointAgentWorkspace, + agent_session: EndpointAgentSession, + redacted_values: Sequence[str], + is_alive: Callable[[], bool], +) -> EndpointAgentProcessOutput: + """Parses the agent's stream files until it exits; safe alongside a live + process or over the remains of a finished one.""" + offset_store = _OffsetStore(agent_session.path / ".offsets.json") + stdout_output, stderr_output = await asyncio.gather( + _read_process_stream( + stream=_FileLineReader( + workspace.agent_stdout_path, + offset_store=offset_store, + offset_key="agent_stdout", + is_alive=is_alive, + ), + stream_name="stdout", + parse_result=True, + redacted_values=redacted_values, + agent_session=agent_session, + ), + _read_process_stream( + stream=_FileLineReader( + workspace.agent_stderr_path, + offset_store=offset_store, + offset_key="agent_stderr", + is_alive=is_alive, + ), + stream_name="stderr", + parse_result=False, + redacted_values=redacted_values, + agent_session=agent_session, + ), + ) + output = stdout_output + if output.report_data is None: + output.report_data = stderr_output.report_data + if output.error is None: + output.error = stderr_output.error + return output + + +async def attach_endpoint_agent( + *, + workspace: EndpointAgentWorkspace, + redacted_values: Sequence[str], + agent_session: EndpointAgentSession, +) -> EndpointAgentProcessOutput: + """Follows a detached session's agent to completion, like + `run_endpoint_agent` without owning the process.""" + async with _session_tailers( + workspace=workspace, agent_session=agent_session, redacted_values=redacted_values + ): + manifest = agent_session.read_manifest() + + def agent_alive() -> bool: + return _pid_alive(manifest.get("agent_pid"), manifest.get("agent_started_at")) + + return await _collect_agent_output( + workspace=workspace, + agent_session=agent_session, + redacted_values=redacted_values, + is_alive=agent_alive, + ) + + async def _read_process_stream( *, - stream: asyncio.StreamReader, + stream: "_LineStream", stream_name: str, parse_result: bool, redacted_values: Sequence[str], @@ -637,7 +1141,16 @@ async def _read_process_stream( message = json.loads(text) except json.JSONDecodeError: continue - if not isinstance(message, dict) or message.get("type") != "result": + if not isinstance(message, dict): + continue + if output.session_id is None: + session_id = message.get("session_id") + if isinstance(session_id, str) and session_id: + output.session_id = session_id + agent_session.record_claude_session_id(session_id) + if message.get("type") == "assistant": + output.made_progress = True + if message.get("type") != "result": continue if message.get("is_error"): error = message.get("result") or "Claude failed" @@ -681,6 +1194,40 @@ async def _terminate_process(proc: asyncio.subprocess.Process) -> None: await proc.wait() +def terminate_agent_process(manifest: dict[str, Any]) -> None: + """Terminates the session's agent process tree, if alive.""" + agent_pid = manifest.get("agent_pid") + if not isinstance(agent_pid, int) or not _pid_alive( + agent_pid, manifest.get("agent_started_at") + ): + return + if IS_WINDOWS: + _terminate_windows_process_tree(agent_pid) + return + with suppress(OSError): + os.killpg(agent_pid, signal.SIGTERM) # pyright: ignore[reportAttributeAccessIssue] + for _ in range(30): + if not _pid_running(agent_pid): + return + time.sleep(0.1) + with suppress(OSError): + os.killpg(agent_pid, signal.SIGKILL) # pyright: ignore[reportAttributeAccessIssue] + + +def _pid_running(pid: int) -> bool: + try: + return psutil.Process(pid).status() != psutil.STATUS_ZOMBIE + except psutil.Error: + return False + + +def _process_started_at(pid: int) -> Optional[float]: + try: + return psutil.Process(pid).create_time() + except psutil.Error: + return None + + def _terminate_windows_process_tree(pid: int) -> None: try: root = psutil.Process(pid) @@ -697,6 +1244,94 @@ def _terminate_windows_process_tree(pid: int) -> None: psutil.wait_procs(alive, timeout=3) +class _LineStream(Protocol): + async def readline(self) -> bytes: ... + + +class _FileLineReader: + """`readline()` over a growing file, so stream parsing survives CLI + restarts: offsets persist, and a later attach continues exactly where the + previous reader stopped.""" + + _POLL_SECONDS = 0.2 + _MAX_CHUNK = 1024 * 1024 + + def __init__( + self, + path: Path, + *, + offset_store: "_OffsetStore", + offset_key: str, + is_alive: Callable[[], bool], + ) -> None: + self._path = path + self._offset_store = offset_store + self._offset_key = offset_key + self._is_alive = is_alive + self._offset = offset_store.get(offset_key) + self._buffer = b"" + self._drained = False + + def _read_chunk(self) -> bytes: + try: + with self._path.open("rb") as f: + f.seek(self._offset) + return f.read(self._MAX_CHUNK) + except OSError: + return b"" + + async def readline(self) -> bytes: + while True: + if b"\n" in self._buffer: + line, self._buffer = self._buffer.split(b"\n", 1) + return line + b"\n" + chunk = self._read_chunk() + if chunk: + self._buffer += chunk + self._offset += len(chunk) + self._offset_store.set(self._offset_key, self._offset) + continue + if not self._is_alive(): + if self._drained: + # A final partial line without a newline, then EOF. + line, self._buffer = self._buffer, b"" + return line + # One more read after death to catch the last flush. + self._drained = True + continue + await asyncio.sleep(self._POLL_SECONDS) + + +class _OffsetStore: + """Persists tailer/mirror byte offsets so resumed sessions do not repeat output.""" + + def __init__(self, path: Path) -> None: + self._path = path + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + data = {} + self._data: dict[str, Any] = data if isinstance(data, dict) else {} + + def get(self, key: str) -> int: + value = self._data.get(key, 0) + return value if isinstance(value, int) and value >= 0 else 0 + + def set(self, key: str, value: int) -> None: + self._data[key] = value + # Re-read and update only this key: several stores (stream readers and + # record mirrors) share one file with disjoint keys, and writing the + # whole in-memory view would clobber the others' offsets. + try: + on_disk = json.loads(self._path.read_text(encoding="utf-8")) + data = on_disk if isinstance(on_disk, dict) else {} + except (OSError, json.JSONDecodeError): + data = {} + data[key] = value + with suppress(OSError): + _write_private_text(self._path, json.dumps(data) + "\n") + + class _ProgressTailer: def __init__( self, @@ -704,11 +1339,15 @@ def __init__( path: Path, redacted_values: Sequence[str], agent_session: EndpointAgentSession, + offset_store: Optional[_OffsetStore] = None, + offset_key: str = "progress", ) -> None: self._path = path self._redacted_values = redacted_values self._agent_session = agent_session - self._offset = 0 + self._offset_store = offset_store + self._offset_key = offset_key + self._offset = offset_store.get(offset_key) if offset_store else 0 async def run(self) -> None: while True: @@ -722,6 +1361,8 @@ def flush(self) -> None: f.seek(self._offset) lines = f.readlines() self._offset = f.tell() + if lines and self._offset_store is not None: + self._offset_store.set(self._offset_key, self._offset) for line in lines: message = _parse_progress(line) if message is not None: @@ -731,6 +1372,56 @@ def flush(self) -> None: ) +class _RecordMirror: + """Mirrors a workspace record file into the persistent session directory, redacted.""" + + def __init__( + self, + *, + source: Path, + target: Path, + redacted_values: Sequence[str], + offset_store: Optional[_OffsetStore] = None, + offset_key: str = "", + ) -> None: + self._source = source + self._target = target + self._redacted_values = redacted_values + self._offset_store = offset_store + self._offset_key = offset_key + self._offset = offset_store.get(offset_key) if offset_store and offset_key else 0 + self._enabled = True + + async def run(self) -> None: + while True: + self.flush() + await asyncio.sleep(1) + + def flush(self) -> None: + if not self._enabled or not self._source.exists(): + return + with self._source.open("rb") as f: + f.seek(self._offset) + data = f.read() + # Mirror complete lines only; a partial line is kept for the next flush. + end = data.rfind(b"\n") + if end < 0: + return + chunk = data[: end + 1].decode("utf-8", errors="replace") + self._offset += end + 1 + if self._offset_store is not None and self._offset_key: + self._offset_store.set(self._offset_key, self._offset) + try: + if not self._target.exists(): + _write_private_text(self._target, "") + with self._target.open("a", encoding="utf-8") as f: + f.write(redact(chunk, self._redacted_values)) + f.flush() + except OSError as e: + self._enabled = False + console.print(f"[warning]Could not mirror {self._target.name}: {e}[/]") + + def _parse_progress(line: str) -> Optional[str]: try: value = json.loads(line) diff --git a/src/dstack/_internal/cli/services/endpoints/apply.py b/src/dstack/_internal/cli/services/endpoints/apply.py index 350a47dd9..8aed115b7 100644 --- a/src/dstack/_internal/cli/services/endpoints/apply.py +++ b/src/dstack/_internal/cli/services/endpoints/apply.py @@ -1,6 +1,6 @@ import argparse from dataclasses import dataclass -from typing import Optional +from typing import Optional, Sequence from rich.markup import escape @@ -9,7 +9,6 @@ from dstack._internal.cli.services.configurators.run import ServiceConfigurator from dstack._internal.cli.services.endpoints.output import ( format_endpoint_benchmark, - format_endpoint_context_length, ) from dstack._internal.cli.services.endpoints.store import EndpointPresetStore from dstack._internal.core.errors import CLIError @@ -32,21 +31,18 @@ def apply_endpoint_preset( api: Client, configuration: EndpointConfiguration, configuration_path: str, - preset_id: Optional[str], + preset_ids: Optional[Sequence[str]], profile_name: Optional[str], command_args: argparse.Namespace, store: EndpointPresetStore, ) -> None: - presets = _get_matching_presets( - store.list(), - configuration=configuration, - preset_id=preset_id, - ) + candidates = _get_candidate_presets(store.list(), preset_ids=preset_ids) + presets = _get_matching_presets(candidates, configuration=configuration) if not presets: - qualifier = f" preset {preset_id!r}" if preset_id else "" - raise CLIError( - f"No matching endpoint preset{qualifier} for {configuration.model.api_model_name}" - ) + qualifier = "" + if preset_ids: + qualifier = f" among {', '.join(preset_ids)}" + raise CLIError(f"No matching preset{qualifier} for {configuration.model.api_model_name}") configurator = ServiceConfigurator(api_client=api) service_args = configurator.get_parser().parse_args([]) @@ -70,18 +66,40 @@ def apply_endpoint_preset( ) +def _get_candidate_presets( + presets: list[EndpointPreset], + *, + preset_ids: Optional[Sequence[str]], +) -> list[EndpointPreset]: + if not preset_ids: + return presets + presets_by_ref = {preset.id: preset for preset in presets} + for preset in presets: + if preset.name is not None: + presets_by_ref.setdefault(preset.name, preset) + missing = [ref for ref in preset_ids if ref not in presets_by_ref] + if len(missing) == 1: + raise CLIError(f"Preset {missing[0]} does not exist") + if missing: + raise CLIError(f"Presets {', '.join(missing)} do not exist") + # Preserve the order given: capacity-aware selection tries candidates in turn. + candidates = [] + for ref in preset_ids: + preset = presets_by_ref[ref] + if preset not in candidates: + candidates.append(preset) + return candidates + + def _get_matching_presets( presets: list[EndpointPreset], *, configuration: EndpointConfiguration, - preset_id: Optional[str], ) -> list[EndpointPreset]: model_name = configuration.model.api_model_name matches = [] for preset in presets: service_model = preset.service.model - if preset_id is not None and preset.id != preset_id: - continue if service_model is None or service_model.name.lower() != model_name.lower(): continue if configuration.context_length is not None: @@ -151,7 +169,5 @@ def _format_requested_model(configuration: EndpointConfiguration) -> str: def _format_selected_preset(preset: EndpointPreset) -> str: - details = ( - f"context={format_endpoint_context_length(preset)}, {format_endpoint_benchmark(preset)}" - ) + details = format_endpoint_benchmark(preset, verbose=True) return f"{escape(preset.id)} ([secondary]{details}[/])" diff --git a/src/dstack/_internal/cli/services/endpoints/create.py b/src/dstack/_internal/cli/services/endpoints/create.py index 9c1ef179e..763c2a7bf 100644 --- a/src/dstack/_internal/cli/services/endpoints/create.py +++ b/src/dstack/_internal/cli/services/endpoints/create.py @@ -1,43 +1,58 @@ import asyncio +import dataclasses import json import os -import secrets +import re import uuid from contextlib import suppress from dataclasses import dataclass from pathlib import Path from typing import Optional, Sequence +import yaml +from rich.table import Table + from dstack._internal.cli.models.endpoint_agent import AgentFinalReport from dstack._internal.cli.models.endpoint_presets import EndpointPreset -from dstack._internal.cli.models.endpoints import EndpointConfiguration +from dstack._internal.cli.models.endpoints import ( + EndpointConfiguration, + EndpointPresetConstraints, +) from dstack._internal.cli.services.endpoints.agent import ( EndpointAgentSession, EndpointAgentWorkspace, + attach_agent_workspace, + attach_endpoint_agent, build_endpoint_agent_env, + claimed_session_name, contains_redacted_value, + create_agent_workspace, create_endpoint_agent_session, - endpoint_agent_workspace, get_claude_auth, get_redacted_values, get_sensitive_inherited_env_values, + load_attachable_agent_session, print_endpoint_progress, + redact, + remove_agent_workspace, run_endpoint_agent, + session_process_alive, + terminate_agent_process, ) from dstack._internal.cli.services.endpoints.presets import endpoint_preset_to_data -from dstack._internal.cli.services.endpoints.prompt import ( - format_endpoint_constraints, - get_endpoint_agent_system_prompt, -) +from dstack._internal.cli.services.endpoints.prompt import get_endpoint_agent_system_prompt from dstack._internal.cli.services.endpoints.store import EndpointPresetStore from dstack._internal.cli.services.endpoints.verify import ( build_verified_endpoint_preset, load_endpoint_agent_report, ) -from dstack._internal.cli.utils.common import console, warn +from dstack._internal.cli.utils.common import confirm_ask, console, warn +from dstack._internal.cli.utils.run import print_offers from dstack._internal.core.errors import CLIError, ConfigurationError -from dstack._internal.core.models.envs import EnvSentinel +from dstack._internal.core.models.configurations import TaskConfiguration +from dstack._internal.core.models.envs import Env, EnvSentinel from dstack._internal.core.models.fleets import FleetStatus +from dstack._internal.core.models.runs import RunSpec from dstack.api import Client _RUN_STOP_TIMEOUT_SECONDS = 10 * 60 @@ -51,6 +66,127 @@ class EndpointPresetCreateResult: final_run_name: str +class AgentExitedWithoutReport(Exception): + """A detached agent died without submitting a report; the session is + resumable rather than failed.""" + + def __init__(self, error: Optional[str]) -> None: + super().__init__(error or "The agent exited without a report") + self.error = error + + +def attach_endpoint_preset( + *, + api: Client, + store: EndpointPresetStore, + preset_id: str, + keep_service: bool = False, +) -> EndpointPresetCreateResult: + agent_session = load_attachable_agent_session(preset_id) + configuration_path = agent_session.path / "preset.dstack.yml" + if not configuration_path.is_file(): + raise CLIError( + f"Session {preset_id} has no configuration copy and cannot be attached;" + f" resume it with [code]--resume {preset_id}[/] instead" + ) + # The session copy is canonical output, not user input: parse it without + # the user-facing deprecation warnings. + try: + configuration = EndpointConfiguration.parse_obj( + yaml.safe_load(configuration_path.read_text(encoding="utf-8")) + ) + except (OSError, ValueError) as e: + raise CLIError(f"Could not read the session configuration: {e}") from e + try: + result = asyncio.run( + _create_endpoint_preset( + api=api, + configuration=_resolve_endpoint_env_best_effort(configuration), + source_configuration=configuration, + store=store, + keep_service=keep_service, + agent_session=agent_session, + attach=True, + ) + ) + except KeyboardInterrupt: + _stop_or_detach_agent_session(agent_session, api) + raise + except AgentExitedWithoutReport as e: + runs_stopped = _stop_active_session_runs(api, agent_session, assume_yes=False) + _suspend_agent_session(agent_session, runs_left_active=not runs_stopped) + raise CLIError(str(e)) from e + except BaseException: + _finish_agent_session(agent_session, "failed") + remove_agent_workspace(agent_session) + raise + _finish_agent_session(agent_session, "success") + remove_agent_workspace(agent_session) + return result + + +def stop_endpoint_session(api: Client, preset_id: str, *, assume_yes: bool = False) -> None: + session = load_attachable_agent_session(preset_id) + terminate_agent_process(session.read_manifest()) + runs_stopped = _stop_active_session_runs(api, session, assume_yes=assume_yes) + _suspend_agent_session(session, runs_left_active=not runs_stopped) + + +def _stop_active_session_runs( + api: Client, session: EndpointAgentSession, *, assume_yes: bool +) -> bool: + """Offers to stop the session's non-terminal runs; returns whether none + are left active.""" + try: + lines = session.runs_path.read_text(encoding="utf-8").splitlines() + except OSError: + return False + names = [] + for line in lines: + try: + name = json.loads(line).get("name") + except json.JSONDecodeError: + continue + if isinstance(name, str) and name: + names.append(name) + active = [] + for name in names: + try: + run = api.client.runs.get(api.project, name) + except Exception: # noqa: BLE001 + continue + if not run.status.is_finished(): + active.append(name) + if not active: + return True + plural = "s" if len(active) != 1 else "" + if not assume_yes and not confirm_ask( + f"Stop {len(active)} active preset creation run{plural} ([code]{', '.join(active)}[/])?" + ): + return False + with console.status("Stopping runs..."): + api.client.runs.stop(api.project, active, abort=False) + return True + + +def _resolve_endpoint_env_best_effort( + configuration: EndpointConfiguration, +) -> EndpointConfiguration: + """For attach: env values only feed redaction, and the agent already runs, + so missing variables are tolerable.""" + configuration = configuration.copy(deep=True) + kept: dict[str, str] = {} + for key, value in configuration.env.items(): + if isinstance(value, EnvSentinel): + resolved = os.environ.get(key) + if resolved: + kept[key] = resolved + else: + kept[key] = value + configuration.env = Env.parse_obj(kept) + return configuration + + def create_endpoint_preset( *, api: Client, @@ -59,8 +195,11 @@ def create_endpoint_preset( keep_service: bool = False, build_name: Optional[str] = None, debug: bool = False, + resume_session: Optional[EndpointAgentSession] = None, + user_prompt: Optional[str] = None, + allowed_fleets: Optional[tuple[str, ...]] = None, ) -> EndpointPresetCreateResult: - agent_session = create_endpoint_agent_session(configuration, debug=debug) + agent_session = resume_session or create_endpoint_agent_session(configuration, debug=debug) try: resolved_configuration = _resolve_endpoint_env(configuration) result = asyncio.run( @@ -72,12 +211,20 @@ def create_endpoint_preset( keep_service=keep_service, build_name=build_name, agent_session=agent_session, + resume=resume_session is not None, + user_prompt=user_prompt, + allowed_fleets=allowed_fleets, ) ) + except KeyboardInterrupt: + _stop_or_detach_agent_session(agent_session, api) + raise except BaseException: - _finish_agent_session(agent_session) + _finish_agent_session(agent_session, "failed") + remove_agent_workspace(agent_session) raise - _finish_agent_session(agent_session, result.preset.id) + _finish_agent_session(agent_session, "success") + remove_agent_workspace(agent_session) return result @@ -90,13 +237,51 @@ async def _create_endpoint_preset( keep_service: bool = False, build_name: Optional[str] = None, agent_session: EndpointAgentSession, + resume: bool = False, + attach: bool = False, + user_prompt: Optional[str] = None, + allowed_fleets: Optional[tuple[str, ...]] = None, ) -> EndpointPresetCreateResult: source_configuration = source_configuration or configuration - build_name = build_name or _get_build_name(configuration.name) - allowed_fleets = _get_allowed_fleets(api, configuration) - if not allowed_fleets: - raise CLIError("The project has no active fleets available for preset creation") - auth = get_claude_auth() + initial_resume_session_id: Optional[str] = None + if attach: + auth = None + workspace = attach_agent_workspace(agent_session) + build_name = build_name or _load_build_name(workspace) + allowed_fleets = () + elif resume: + # The prompt is fixed at session creation, like the constraints. + pinned_prompt = agent_session.read_user_prompt() + if user_prompt is not None and user_prompt != pinned_prompt: + warn( + "The configuration prompt is ignored when resuming: the session 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) + 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 + build_name = build_name or _load_build_name(workspace) + allowed_fleets = () + else: + if allowed_fleets is None: + allowed_fleets = _get_allowed_fleets(api, configuration) + if not allowed_fleets: + raise CLIError("The project has no active fleets available for preset creation") + auth = get_claude_auth() + workspace = create_agent_workspace(agent_session) + build_name = build_name or _get_build_name( + configuration.name, configuration.model.api_model_name, agent_session.preset_id + ) + if auth is not None: + agent_session.update_manifest(status="running", pid=os.getpid(), claude_model=auth.model) + else: + agent_session.update_manifest(status="running", pid=os.getpid()) endpoint_env = configuration.env.as_dict() token = getattr(api.client, "_token", None) @@ -105,17 +290,19 @@ async def _create_endpoint_preset( redacted_values = get_redacted_values( [ token, - auth.api_key or "", + (auth.api_key if auth is not None else None) or "", *endpoint_env.values(), *get_sensitive_inherited_env_values(), ] ) + env: dict[str, str] = {} report: Optional[AgentFinalReport] = None preset: Optional[EndpointPreset] = None preset_path: Optional[Path] = None creation_succeeded = False + interrupted = False cleanup_error: Optional[str] = None - with endpoint_agent_workspace() as workspace: + if auth is not None: env = build_endpoint_agent_env( api=api, endpoint_env=endpoint_env, @@ -123,19 +310,32 @@ async def _create_endpoint_preset( workspace=workspace, token=token, ) - prompt = _build_prompt( + prompt = get_endpoint_agent_system_prompt(user_prompt=user_prompt) + if not resume and not attach: + if user_prompt: + agent_session.write_user_prompt(user_prompt) + constraints_text = _build_constraints( configuration=configuration, build_name=build_name, allowed_fleets=allowed_fleets, ) + workspace.constraints_path.write_text(constraints_text, encoding="utf-8") if agent_session.debug: agent_session.write_prompt(prompt) - print_endpoint_progress( - f"Starting endpoint preset creation for {configuration.model.api_model_name}. " - f"Allowed fleets: {', '.join(allowed_fleets)}.", - agent_session=agent_session, - ) - try: + agent_session.write_constraints(constraints_text) + if auth is not None: + agent_session.write_agent_info(auth) + try: + if attach: + process_output = await attach_endpoint_agent( + workspace=workspace, + redacted_values=redacted_values, + agent_session=agent_session, + ) + if process_output.report_data is None and not workspace.final_report_path.exists(): + raise AgentExitedWithoutReport(process_output.error) + else: + assert auth is not None process_output = await run_endpoint_agent( prompt=prompt, env=env, @@ -143,27 +343,40 @@ async def _create_endpoint_preset( auth=auth, redacted_values=redacted_values, agent_session=agent_session, + initial_resume_session_id=initial_resume_session_id, ) - report = load_endpoint_agent_report( - output=process_output, + report = load_endpoint_agent_report( + output=process_output, + workspace=workspace, + redacted_values=redacted_values, + ) + run = api.client.runs.get(api.project, report.run_name) + preset = build_verified_endpoint_preset( + run=run, + endpoint_configuration=source_configuration, + report=report, + preset_id=agent_session.preset_id or None, + name=claimed_session_name(agent_session.read_manifest()), + ) + if contains_redacted_value(endpoint_preset_to_data(preset), redacted_values): + raise CLIError("Generated preset contains a secret value") + preset_path = store.save(preset) + print_endpoint_progress( + f"Saved preset {preset.id} for {preset.base} at {preset_path}.", + agent_session=agent_session, + ) + creation_succeeded = True + except (KeyboardInterrupt, asyncio.CancelledError): + interrupted = True + raise + finally: + if agent_session.debug: + _save_final_report_copy( workspace=workspace, - redacted_values=redacted_values, - ) - run = api.client.runs.get(api.project, report.run_name) - preset = build_verified_endpoint_preset( - run=run, - endpoint_configuration=source_configuration, - report=report, - ) - if contains_redacted_value(endpoint_preset_to_data(preset), redacted_values): - raise CLIError("Generated endpoint preset contains a secret value") - preset_path = store.save(preset) - print_endpoint_progress( - f"Saved endpoint preset {preset.id} for {preset.base} at {preset_path}.", agent_session=agent_session, + redacted_values=redacted_values, ) - creation_succeeded = True - finally: + if not interrupted: keep_final_service = keep_service and creation_succeeded try: await _cleanup_runs( @@ -215,25 +428,128 @@ def _resolve_endpoint_env(configuration: EndpointConfiguration) -> EndpointConfi def _finish_agent_session( session: EndpointAgentSession, - preset_id: Optional[str] = None, + status: str, ) -> None: try: - path = session.finish(preset_id) + path = session.finish(status) except OSError as e: path = session.path warn(f"Could not finalize agent output. Files remain at {path}: {e}") console.print(f"Agent log saved to [code]{path / 'agent.log'}[/]") -def _get_build_name(endpoint_name: Optional[str]) -> str: - if endpoint_name is None: - raise CLIError("Endpoint name is required. Set `name` in the configuration or use --name") - suffix = secrets.token_hex(3) - # Leave room for the numeric submission suffix while retaining a recognizable prefix. - prefix = endpoint_name[:28].rstrip("-") +def _stop_or_detach_agent_session( + session: EndpointAgentSession, api: Optional[Client] = None +) -> None: + """Apply-style interrupt: stop the session, or detach and leave the agent + working — it stays visible as a running session in `dstack preset`.""" + manifest = session.read_manifest() + agent_alive = session_process_alive({**manifest, "pid": None}) + stop = True + if agent_alive: + try: + stop = confirm_ask(f"Stop the preset creation session [code]{session.preset_id}[/]?") + except (KeyboardInterrupt, EOFError): + stop = True + if stop: + terminate_agent_process(manifest) + runs_stopped = False + if api is not None: + runs_stopped = _stop_active_session_runs(api, session, assume_yes=False) + _suspend_agent_session(session, runs_left_active=not runs_stopped) + return + session.update_manifest(pid=None) + console.print( + f"\nDetached. The session keeps running; attach with" + f" [code]dstack preset attach {session.preset_id}[/]" + f" or stop it with [code]dstack preset stop {session.preset_id}[/]." + ) + + +def _suspend_agent_session( + session: EndpointAgentSession, *, runs_left_active: bool = True +) -> None: + try: + session.finish("interrupted") + except OSError as e: + warn(f"Could not record the interrupted session state: {e}") + if runs_left_active: + console.print( + f"\nSession [code]{session.preset_id}[/] interrupted. " + "Its runs may still be active and accruing cost." + ) + console.print( + f"Resume with [code]dstack preset create -f " + f"--resume {session.preset_id}[/], or stop the runs with [code]dstack stop[/]." + ) + else: + console.print(f"\nSession [code]{session.preset_id}[/] interrupted.") + console.print( + f"Resume with [code]dstack preset create -f " + f"--resume {session.preset_id}[/]." + ) + + +def _get_build_name(name: Optional[str], model_name: str, suffix: str) -> str: + base = name or _model_slug(model_name) + # Leave room for the preset id and numeric submission suffix while retaining + # a recognizable prefix. + prefix = base[:26].rstrip("-") return f"{prefix}-{suffix}" +def _model_slug(model_name: str) -> str: + """A run-name-safe slug for name-less presets, from the model's basename.""" + basename = model_name.rsplit("/", 1)[-1] + slug = re.sub(r"[^a-z0-9]+", "-", basename.lower()).strip("-") + if not slug or not slug[0].isalpha(): + slug = f"model-{slug}".strip("-") + return slug + + +def _load_build_name(workspace: EndpointAgentWorkspace) -> str: + try: + data = json.loads(workspace.constraints_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as e: + raise CLIError(f"The session constraints cannot be read: {e}") from e + prefix = data.get("run_name_prefix") if isinstance(data, dict) else None + if not isinstance(prefix, str) or not prefix: + raise CLIError("The session constraints do not contain a run name prefix") + return prefix + + +def plan_endpoint_preset(*, api: Client, configuration: EndpointConfiguration) -> tuple[str, ...]: + """Resolves the allowed fleets and shows what the agent will have to work + with — Project, User, the effective fleets, and their offers. Agent-free.""" + allowed_fleets = _get_allowed_fleets(api, configuration) + if not allowed_fleets: + raise CLIError("The project has no active fleets available for preset creation") + _print_fleet_offers(api, allowed_fleets) + return allowed_fleets + + +def _print_fleet_offers(api: Client, allowed_fleets: tuple[str, ...]) -> None: + try: + # Image and user are set so the server neither defaults gpu.vendor to + # nvidia nor pulls image config from a registry (as in `dstack offer`). + offer_configuration = TaskConfiguration(commands=[":"], image="scratch", user="root") + offer_configuration.fleets = list(allowed_fleets) + run_spec = RunSpec(configuration=offer_configuration, profile=None) + with console.status("Getting offers..."): + run_plan = api.client.runs.get_plan(api.project, run_spec, max_offers=10) + props = Table(box=None, show_header=False) + props.add_column(no_wrap=True) + props.add_column() + props.add_row("[bold]Project[/bold]", run_plan.project_name) + props.add_row("[bold]User[/bold]", run_plan.user) + props.add_row("[bold]Fleets[/bold]", ", ".join(allowed_fleets)) + console.print(props) + console.print() + print_offers(run_plan.job_plans[0], dim_after_first=False) + except Exception as e: # noqa: BLE001 + warn(f"Could not list offers for the allowed fleets: {e}") + + def _get_allowed_fleets(api: Client, configuration: EndpointConfiguration) -> tuple[str, ...]: if configuration.fleets is not None: return tuple( @@ -248,33 +564,40 @@ def _get_allowed_fleets(api: Client, configuration: EndpointConfiguration) -> tu ) -def _build_prompt( +def _build_constraints( *, configuration: EndpointConfiguration, build_name: str, allowed_fleets: Sequence[str], ) -> str: - context_lines = [f"- service_model_name: {configuration.model.api_model_name}"] - if configuration.model.allows_variant_selection: - context_lines.append(f"- base_model: {configuration.model.api_model_name}") - else: - context_lines.append(f"- model_repo: {configuration.model.exact_repo}") - if configuration.context_length is not None: - context_lines.append(f"- context_length: {configuration.context_length}") - return f"""{get_endpoint_agent_system_prompt()} - -Endpoint context: -- endpoint_name: {build_name} -{chr(10).join(context_lines)} - -{ - format_endpoint_constraints( - configuration, - configuration.env.as_dict(), - allowed_fleets=allowed_fleets, - ) - } -""" + constraints = EndpointPresetConstraints.parse_obj( + { + "run_name_prefix": build_name, + "model": json.loads(configuration.model.json(exclude_none=True)), + "context_length": configuration.context_length, + "max_trials": configuration.effective_max_trials, + "concurrency": configuration.effective_concurrency, + "fleets": list(allowed_fleets), + "env": list(configuration.env), + } + ) + # All fields are always present; unset optional constraints render as null. + return json.dumps(json.loads(constraints.json()), indent=2) + "\n" + + +def _save_final_report_copy( + *, + workspace: EndpointAgentWorkspace, + agent_session: EndpointAgentSession, + redacted_values: Sequence[str], +) -> None: + if not workspace.final_report_path.exists(): + return + try: + report_text = workspace.final_report_path.read_text(encoding="utf-8", errors="replace") + agent_session.write_final_report(redact(report_text, redacted_values)) + except OSError as e: + warn(f"Could not save a final report copy: {e}") async def _cleanup_runs( @@ -286,7 +609,7 @@ async def _cleanup_runs( agent_session: EndpointAgentSession, keep_final_service: bool = False, ) -> None: - run_names = _load_submitted_run_names(workspace.submissions_path) + run_names = _load_submitted_run_names(workspace.runs_path) if final_run_name is not None: run_names.append(final_run_name) run_names = list(dict.fromkeys(run_names)) @@ -301,10 +624,6 @@ async def _cleanup_runs( active_names.append(name) if not active_names: return - print_endpoint_progress( - f"Stopping preset creation runs: {', '.join(active_names)}.", - agent_session=agent_session, - ) api.client.runs.stop(api.project, active_names, abort=False) deadline = asyncio.get_running_loop().time() + _RUN_STOP_TIMEOUT_SECONDS pending = set(active_names) @@ -317,7 +636,8 @@ async def _cleanup_runs( pending.remove(name) if pending: await asyncio.sleep(2) - print_endpoint_progress("All preset creation runs stopped.", agent_session=agent_session) + if agent_session.debug: + print_endpoint_progress("All preset creation runs stopped.", agent_session=agent_session) def _load_submitted_run_names(path: Path) -> list[str]: diff --git a/src/dstack/_internal/cli/services/endpoints/output.py b/src/dstack/_internal/cli/services/endpoints/output.py index 3d11c02e8..c0e23fde5 100644 --- a/src/dstack/_internal/cli/services/endpoints/output.py +++ b/src/dstack/_internal/cli/services/endpoints/output.py @@ -1,114 +1,187 @@ from collections import defaultdict +from datetime import datetime +from typing import Any, Optional from rich.table import Table from dstack._internal.cli.models.endpoint_presets import ( EndpointPreset, - EndpointPresetValidation, ) from dstack._internal.cli.utils.common import add_row_from_dict, console from dstack._internal.utils.common import pretty_date, pretty_resources +_STATUS_DISPLAY = { + "ready": ("success", "green"), + "running": ("clauding", "bold deep_sky_blue1"), + "verifying": ("verifying", "bold deep_sky_blue1"), + "interrupted": ("interrupted", "bold gold1"), + "failed": ("failed", "indian_red1"), +} -def print_endpoint_presets(presets: list[EndpointPreset], verbose: bool = False) -> None: + +def _format_status(status: str) -> str: + text, style = _STATUS_DISPLAY.get(status, (status, None)) + return f"[{style}]{text}[/]" if style else text + + +def _trials_exhausted(session: dict[str, Any]) -> bool: + trials = session.get("trials") + max_trials = session.get("max_trials") + return ( + isinstance(trials, dict) + and isinstance(max_trials, int) + and isinstance(trials.get("count"), int) + and trials["count"] >= max_trials + ) + + +def _format_trial_progress(session: Optional[dict[str, Any]]) -> str: + """The ` (N/M)` suffix; stays outside the status markup to render in the + default color.""" + if not isinstance(session, dict): + return "" + trials = session.get("trials") + max_trials = session.get("max_trials") + if not isinstance(trials, dict) or not (trials.get("count") or isinstance(max_trials, int)): + return "" + progress = str(trials.get("count") or 0) + if isinstance(max_trials, int): + progress += f"/{max_trials}" + return f" [secondary]({progress})[/]" + + +def print_endpoint_presets( + presets: list[EndpointPreset], + sessions: Optional[list[dict[str, Any]]] = None, + verbose: bool = False, +) -> None: table = Table(box=None) - table.add_column("MODEL", no_wrap=True) - table.add_column("RESOURCES" if verbose else "GPU") - table.add_column("CONTEXT", justify="right") - table.add_column("BENCHMARK", min_width=len("concurrency=1"), overflow="fold") - table.add_column("CREATED", no_wrap=True) + table.add_column("BASE", no_wrap=True) + table.add_column("ID", no_wrap=True) + table.add_column("NAME", no_wrap=True) + table.add_column("RESOURCES" if verbose else "GPU", style="secondary") + table.add_column("BENCHMARK", min_width=len("con=1"), overflow="fold") + table.add_column("STATUS", no_wrap=True) + table.add_column("SUBMITTED", no_wrap=True, style="secondary") presets_by_base: dict[str, list[EndpointPreset]] = defaultdict(list) + repo_to_base: dict[str, str] = {} for preset in presets: presets_by_base[preset.base].append(preset) - - for base, base_presets in presets_by_base.items(): - add_row_from_dict(table, {"MODEL": f"[bold]{base}[/]"}) - for preset in base_presets: - _add_preset(table, preset, verbose=verbose) + repo_to_base[preset.model] = preset.base + sessions_by_model: dict[str, list[dict[str, Any]]] = defaultdict(list) + creations_by_id: dict[str, dict[str, Any]] = {} + for session in sessions or []: + if str(session.get("status")) == "success": + # A completed creation session decorates its preset row. + creations_by_id[str(session.get("id"))] = session + continue + model = str(session.get("model") or "unknown") + sessions_by_model[repo_to_base.get(model, model)].append(session) + + for base in sorted({*presets_by_base, *sessions_by_model}, key=str.lower): + add_row_from_dict(table, {"BASE": base}) + # Newest first within a group, as in `dstack ps`. + for preset in sorted( + presets_by_base.get(base, []), key=lambda p: p.created_at, reverse=True + ): + _add_preset(table, preset, verbose=verbose, creation=creations_by_id.get(preset.id)) + for session in sorted( + sessions_by_model.get(base, []), + key=lambda s: str(s.get("created_at") or ""), + reverse=True, + ): + _add_session(table, session) console.print(table) console.print() -def _add_preset(table: Table, preset: EndpointPreset, *, verbose: bool) -> None: - groups = preset.service.replica_groups - column = "RESOURCES" if verbose else "GPU" +def _add_session(table: Table, session: dict[str, Any], base_label: Optional[str] = None) -> None: + created = "" + created_at = session.get("created_at") + if isinstance(created_at, str): + try: + created = pretty_date(datetime.fromisoformat(created_at)) + except ValueError: + created = created_at + benchmark = "" + gpu = "" + status_key = str(session.get("status", "")) + if status_key == "running" and _trials_exhausted(session): + # The trial budget is spent, so the agent is deploying and verifying + # the final service. + status_key = "verifying" + status = _format_status(status_key) + _format_trial_progress(session) + trials = session.get("trials") + if isinstance(trials, dict): + best = trials.get("best") + if isinstance(best, dict): + parts = ["best trial:"] + if best.get("concurrency"): + parts.append(f"con={best['concurrency']}") + parts.append(f"{_format_number(best['tok_s'])} tok/s") + benchmark = " ".join(parts) + gpu = best.get("gpu") or "" add_row_from_dict( table, { - "MODEL": f"[secondary] preset={preset.id}[/]", - column: _format_resources(groups[0].resources, verbose=verbose), - "CONTEXT": format_endpoint_context_length(preset), - "BENCHMARK": format_endpoint_benchmark(preset, verbose=verbose), - "CREATED": pretty_date(preset.created_at), + "ID": str(session.get("id", "")), + "NAME": str(session.get("name") or ""), + "GPU": gpu, + "RESOURCES": gpu, + "BENCHMARK": benchmark, + "STATUS": status, + "SUBMITTED": created, }, ) - if preset.model != preset.base: - add_row_from_dict( - table, - {"MODEL": f" repo={preset.model}"}, - style="secondary", - ) + + +def _add_preset( + table: Table, + preset: EndpointPreset, + *, + verbose: bool, + creation: Optional[dict[str, Any]] = None, +) -> None: + groups = preset.service.replica_groups + column = "RESOURCES" if verbose else "GPU" + row = { + "ID": preset.id, + "NAME": preset.name or "", + column: _format_resources(groups[0].resources, verbose=verbose), + "STATUS": _format_status("ready") + _format_trial_progress(creation), + "BENCHMARK": format_endpoint_benchmark(preset, verbose=verbose), + "SUBMITTED": pretty_date(preset.created_at), + } + if verbose and preset.model != preset.base: + row["BASE"] = f"[secondary] repo={preset.model}[/]" + add_row_from_dict(table, row) if len(groups) > 1: for group in groups: add_row_from_dict( table, { - "MODEL": f" group={group.name}", + "BASE": f" group={group.name}", column: _format_resources(group.resources, verbose=verbose), }, style="secondary", ) -def format_endpoint_context_length(preset: EndpointPreset) -> str: - return _format_token_count(preset.context_length) - - def format_endpoint_benchmark(preset: EndpointPreset, *, verbose: bool = False) -> str: - validation = preset.validations[0] - benchmark = validation.benchmark + benchmark = preset.validations[0].benchmark workload = benchmark.workload metrics = benchmark.metrics - requests_per_second = metrics.successful_requests / metrics.duration_seconds output_tokens_per_second = metrics.total_output_tokens / metrics.duration_seconds parts = [ - f"concurrency={workload.concurrency}", + f"con={workload.concurrency}", f"{_format_number(output_tokens_per_second)} tok/s", f"TTFT {_format_latency(metrics.ttft_ms.p50)}", ] if verbose: - ttft = _format_latency_summary( - metrics.ttft_ms.mean, metrics.ttft_ms.p50, metrics.ttft_ms.p99 - ) - tpot = _format_latency_summary( - metrics.tpot_ms.mean, metrics.tpot_ms.p50, metrics.tpot_ms.p99 - ) - parts.extend( - [ - f"hardware={_format_validation_gpus(validation)}", - f"api={workload.api}", - f"n={workload.num_requests}", - f"{_format_token_count(workload.input_tokens)}" - f"->{_format_token_count(workload.output_tokens)}", - f"{_format_number(requests_per_second)} req/s", - f"duration={_format_number(metrics.duration_seconds)}s", - f"TTFT mean/p50/p99={ttft}", - f"TPOT mean/p50/p99={tpot}", - f"{benchmark.tool} {benchmark.tool_version}", - ] - ) + parts.insert(0, f"ctx={_format_token_count(preset.context_length)}") return " ".join(parts) -def _format_validation_gpus(validation: EndpointPresetValidation) -> str: - gpus = [ - _format_resources(resources, verbose=False) - for replica_group in validation.replicas - for resources in replica_group.resources - ] - return "+".join(gpus) or "-" - - def _format_token_count(value: int) -> str: for divisor, suffix in ((1024 * 1024, "M"), (1024, "K")): if value >= divisor and value % divisor == 0: @@ -117,6 +190,9 @@ def _format_token_count(value: int) -> str: def _format_number(value: float) -> str: + # `g` switches to scientific notation for values >= 1000 (after rounding). + if abs(value) >= 999.5: + return f"{value:.0f}" return f"{value:.3g}" @@ -126,11 +202,6 @@ def _format_latency(value_ms: float) -> str: return f"{_format_number(value_ms)}ms" -def _format_latency_summary(*values_ms: float) -> str: - divisor, unit = (1000, "s") if max(values_ms) >= 1000 else (1, "ms") - return "/".join(_format_number(value / divisor) for value in values_ms) + unit - - def _format_resources(resources, *, verbose: bool) -> str: if resources is None: return "-" diff --git a/src/dstack/_internal/cli/services/endpoints/presets.py b/src/dstack/_internal/cli/services/endpoints/presets.py index 727133fa3..553730346 100644 --- a/src/dstack/_internal/cli/services/endpoints/presets.py +++ b/src/dstack/_internal/cli/services/endpoints/presets.py @@ -1,6 +1,6 @@ import hashlib import json -from typing import Any +from typing import Any, Optional import gpuhunt @@ -26,6 +26,8 @@ def build_endpoint_preset( model: str, context_length: int, benchmark: EndpointBenchmark, + preset_id: Optional[str] = None, + name: Optional[str] = None, ) -> EndpointPreset: service = service.copy(deep=True) service.name = None @@ -38,8 +40,9 @@ def build_endpoint_preset( ) set_service_gpu_vendors_from_validations(service, [validation]) return EndpointPreset( + name=name, base=base_model, - id=make_endpoint_preset_id(service, context_length=context_length), + id=preset_id or make_endpoint_preset_id(service, context_length=context_length), model=model, context_length=context_length, created_at=get_current_datetime(), @@ -67,6 +70,7 @@ def endpoint_preset_to_data(preset: EndpointPreset) -> dict[str, Any]: return { "base": preset.base, "id": preset.id, + **({"name": preset.name} if preset.name else {}), "model": preset.model, "context_length": preset.context_length, "created_at": preset.created_at.isoformat(), @@ -115,7 +119,7 @@ def resources_spec_from_instance_resources(resources: Resources) -> ResourcesSpe or gpu.vendor != first_gpu.vendor for gpu in resources.gpus ): - raise ValueError("endpoint preset cannot be built from mixed-GPU instances") + raise ValueError("preset cannot be built from mixed-GPU instances") data["gpu"] = { "name": first_gpu.name, "memory": format_mib_as_gb(first_gpu.memory_mib), diff --git a/src/dstack/_internal/cli/services/endpoints/prompt.py b/src/dstack/_internal/cli/services/endpoints/prompt.py index 89e66964d..0123047ee 100644 --- a/src/dstack/_internal/cli/services/endpoints/prompt.py +++ b/src/dstack/_internal/cli/services/endpoints/prompt.py @@ -1,58 +1,36 @@ -import enum -import json +import re from pathlib import Path -from typing import Any, Sequence +from typing import Optional -from dstack._internal.cli.models.endpoints import EndpointConfiguration -from dstack._internal.core.models.profiles import ProfileParams +from dstack._internal.core.errors import CLIError _SYSTEM_PROMPT_PATH = Path(__file__).resolve().parent / "resources" / "system_prompt.md" - -def get_endpoint_agent_system_prompt() -> str: - return _SYSTEM_PROMPT_PATH.read_text(encoding="utf-8").strip() - - -def format_endpoint_constraints( - configuration: EndpointConfiguration, - endpoint_env: dict[str, str], - *, - allowed_fleets: Sequence[str], -) -> str: - lines = [ - "Fixed endpoint constraints:", - "- Do not submit any task or service that conflicts with these values.", - ] - for field in ProfileParams.__fields__: - if field == "fleets": - continue - value = getattr(configuration, field) - if value is not None: - lines.append(f"- {field}: {_format_constraint_value(value)}") - if allowed_fleets: - lines.append(f"- fleets: {', '.join(allowed_fleets)}") - if configuration.gateway is not None: - lines.append(f"- gateway: {_format_constraint_value(configuration.gateway)}") - lines.append("- endpoint_env_keys: " + (", ".join(endpoint_env) if endpoint_env else "none")) - return "\n".join(lines) - - -def _format_constraint_value(value: Any) -> str: - data = _constraint_value_to_data(value) - if isinstance(data, (bool, dict, list)): - return json.dumps(data, sort_keys=True) - return str(data) - - -def _constraint_value_to_data(value: Any) -> Any: - if isinstance(value, enum.Enum): - return value.value - if hasattr(value, "format"): - return value.format() - if hasattr(value, "json"): - return json.loads(value.json(exclude_none=True)) - if isinstance(value, dict): - return {key: _constraint_value_to_data(item) for key, item in value.items()} - if isinstance(value, (list, tuple)): - return [_constraint_value_to_data(item) for item in value] - return value +# `` emits CONTENT (with `{NAME}` interpolated) when the +# variable NAME is set, and nothing otherwise. The conditional text lives in +# the document; this module only applies the rule. +_DIRECTIVE_PATTERN = re.compile(r"", re.DOTALL) + + +# TODO: reintroduce a `# Resume` section in system_prompt.md once session resume +# (seeded from `runs.jsonl` and `trials.jsonl`) is designed. +def get_endpoint_agent_system_prompt(user_prompt: Optional[str] = None) -> str: + text = _SYSTEM_PROMPT_PATH.read_text(encoding="utf-8").strip() + variables = {"prompt": user_prompt.strip() if user_prompt else None} + applied = 0 + + def substitute(match: re.Match) -> str: + nonlocal applied + name, content = match.group(1), match.group(2) + if name not in variables: + raise CLIError(f"Unknown variable {name!r} in the agent system prompt") + value = variables[name] + if not value: + return "" + applied += 1 + return content.replace("{" + name + "}", value) + + rendered = _DIRECTIVE_PATTERN.sub(substitute, text) + if variables["prompt"] and not applied: + raise CLIError("The agent system prompt has no place for the user prompt") + return re.sub(r"\n{3,}", "\n\n", rendered) diff --git a/src/dstack/_internal/cli/services/endpoints/resources/system_prompt.md b/src/dstack/_internal/cli/services/endpoints/resources/system_prompt.md index a6ff533eb..ab62a35b4 100644 --- a/src/dstack/_internal/cli/services/endpoints/resources/system_prompt.md +++ b/src/dstack/_internal/cli/services/endpoints/resources/system_prompt.md @@ -1,295 +1,336 @@ # Objective -You are the endpoint preset creation agent for dstack. Produce one final dstack -service that can be saved as a reusable endpoint preset. Report success only -after that service answers a real request using `service_model_name` from -`Endpoint context:` through the dstack service URL. +Your goal is to find a model serving configuration with the best +performance through sequential experimental trials. Once the best-performing +candidate is found, it is deployed as a `dstack` service for the final +benchmark and saved as a reusable endpoint preset. + +# Constraints + +`constraints.json` in the workspace root contains the effective constraints +for this session. No submitted run or final service may conflict with them. + +Field semantics: + +- `run_name_prefix`: the required prefix for all submitted run names, + followed by the run counter; see `# Runs`. +- `model`: the model to serve. If it has `repo`, deploy that repo/path + exactly. If it has `base`, choose a compatible repo/path that best fits + performance and hardware within the constraints: the base repo itself, a + different precision or quantization, or another trusted compatible repo. + The client-facing model name of the final service is `model.name` when + set, otherwise `model.repo` or `model.base`. +- `context_length`: the minimum context length the selected repo/path and + the final service must support. `null` means no minimum is required. +- `max_trials`: the maximum number of trials in this session. +- `concurrency`: the number of simultaneous requests for every benchmark in + this session. It is fixed so that benchmark results are comparable. +- `fleets`: use these existing `dstack` fleets only. Do not create, delete, + apply, or edit fleets. +- `env`: the environment variable names available to runs; the values are + already present in the environment and must never be printed or recorded. + +During the trials and experimentation aimed at the best performance, you may +pick the hardware (the best available within the allowed `dstack` fleets), +the model variant (only if `model` has `base`), the serving framework, the +Docker image and dependencies, the serving framework parameters, and +anything else within these constraints — except generating custom kernels, +patching drivers, or patching serving framework source code. + + +## CLI And Skills + +All trials and the final verification are done using `dstack`. This includes +fetching information about `dstack` fleets, offers, and runs, as well as +submitting `dstack` runs (e.g. tasks for trials and services for the final +verification). + +To do this, use the real `dstack` CLI and shell commands in this workspace. +Load and follow `/dstack` for `dstack` CLI/YAML syntax. Load and follow +`/dstack-prototyping` for how to test a model-serving configuration with +`dstack` tasks before verifying it as a `dstack` service. If a skill cannot +be loaded with its +slash command, read it from `.claude/skills//SKILL.md` and +follow it the same way. + +# Workspace Files -The helper appends to `progress.jsonl`; the caller shows only the message text. +Files provided to you in the workspace root (read them; never edit them): -Write progress before your first investigation action, whenever you submit a -run, when new evidence changes the plan, when model verification succeeds or -fails, and before the final report. Messages should explain what you tried, what -happened, and what you will do next. Do not put raw YAML, command output, long -tables, traces, or secrets in progress. +- `constraints.json`: the effective constraints; see `# Constraints`. -Write a short progress message for every meaningful choice or action. Each -message should say what you did or chose, why, what evidence you used, what -happened, and what you will do next. Include the run name when a run is -involved. Include the fleet or backend when a fleet or backend is involved. Do -not use generic phase labels as the message. +Files you are expected to maintain in the workspace root: -# Workspace Files +- `runs.jsonl`: the append-only record of submitted `dstack` runs; see + `# Runs`. +- `progress.jsonl`: progress messages, written through the `progress` helper; + see `# Progress`. +- `trials.jsonl`: the append-only record of completed trials; see `# Trials`. +- `final_report.json`: the final report; see `# Final Report`. -Create local files only in the current workspace. +You may create any other working files (run YAML files, benchmark output, +notes) in the workspace, and only there: do not deliberately save files +elsewhere on this machine. Incidental writes made by the tools you run +(caches, temporary files, SSH configuration) are fine wherever those tools +keep them. Files inside running `dstack` tasks or services are not subject +to this rule. -`submissions.jsonl` is append-only. For every dstack task or service you submit, -append one JSON line when you submit it and more JSON lines when you learn its -run id, status, URL, or final outcome. +# Runs -Use these fields: +If you submit `dstack` runs via the `dstack` CLI, e.g. to create tasks or +services, name them as described below and record them in `runs.jsonl`. -- `event`: `submit`, `update`, or `final` -- `name`: submitted dstack run name -- `type`: `task` or `service` -- `status`: current known status -- `config_path`: YAML file path, when applicable -- `run_id`: run id, when known -- `reason`: why this run exists or why its status changed -- `service_url`: service URL, when known +Use only `-` as the run name, with +`run_name_prefix` from `constraints.json` and `` counting all +runs submitted in this session: the first submitted run is +`-1`, the second is `-2`, and so on. The +next counter value is the number of lines in `runs.jsonl` plus one. Do not +add framework, hardware, role, or purpose suffixes to run names; record those +details in workspace files and progress. -Example: +`runs.jsonl` is append-only. Immediately after submitting a run, fetch its +id with `dstack run get --json` and append one JSON line: ```json -{"event":"submit","name":"qwen-endpoint-1","type":"task","status":"submitted","config_path":"qwen-endpoint-1.dstack.yml","reason":"test the serving image and local model request on an allowed fleet"} -{"event":"update","name":"qwen-endpoint-1","type":"task","status":"running","run_id":"..."} +{"name":"qwen-endpoint-1","id":""} ``` -Do not rewrite previous `submissions.jsonl` lines; the latest line for a run is -the current record. Stop runs you no longer need unless they are still needed for -attach/SSH debugging, logs, or backend diagnosis. +Never edit or delete existing lines. Stop runs you no longer need unless they +are still needed for attach/SSH debugging, logs, or backend diagnosis. -After stopping a task or service, follow `/dstack` structured status guidance -and confirm that the run reached a terminal status before continuing. +After stopping a `dstack` task or service, follow `/dstack` structured status +guidance and confirm that the run reached a terminal status before continuing. -# Run Names - -Do not use the endpoint name itself as a submitted run name. +# Progress -Use only `-` for dstack runs submitted for -this endpoint. The first submitted run is `-1`, the second is -`-2`, and so on. Do not add framework, hardware, role, or purpose -suffixes to run names; record those details in workspace files and progress. +From the very beginning to the very end — through every trial, the service +verification, and each transition between them — report all major +intentions, decisions, and results, concisely but clearly, to +`progress.jsonl`. Write each message with the `progress` helper; it appends +to `progress.jsonl`, and the caller shows only the message text: -# Endpoint Constraints +```bash +progress "" +``` -Use existing allowed fleets only. Do not create, delete, apply, or edit fleets, -including `nodes`, `target`, `idle_duration`, backends, resources, max nodes, or -ownership. +Make sure progress messages explain your choices. This explicitly includes, +but is not limited to, the choice of fleet, backend, hardware, model repo, +and serving framework: name what you chose, why, what evidence you used, and +what you rejected. + +Do not put raw YAML, command output, long tables, traces, or secrets in +progress messages. + +# Trials (Main Section) + +The end goal is to find a model serving configuration that matches the +constraints and delivers the best performance. The search is done via +so-called trials, where each trial is formed around a substantive idea on +how to get better performance than the previous trials. Do not consider +P/D disaggregation setups yet. + + + +Trial ideas must not rely only on what you already know. Research how to +get the best performance for the chosen model, serving framework, and +hardware in trustworthy sources. Start with these: + +- vLLM recipes: `https://recipes.vllm.ai/` (model index: + `https://recipes.vllm.ai/models.json`) +- SGLang docs: `https://docs.sglang.io/` (fetch `/llms.txt` for the page + index) +- SGLang model recipes: `https://docs.sglang.io/cookbook/autoregressive/intro` +- Release notes: `https://github.com/vllm-project/vllm/releases` and + `https://github.com/sgl-project/sglang/releases` +- Performance-loop methodology (profiling, benchmark contracts): + `https://www.lmsys.org/blog/2026-07-02-agent-assisted-sglang-development` + +Go beyond this list proactively — official docs, repo issues, and reputable +benchmarks — whenever that can help the trial. Research before the first +trial and whenever a benchmark exposes a bottleneck. + +For each trial, use `dstack` tasks (see `# Task Usage`). During a trial, run +commands interactively inside the task (over SSH) and measure the +performance when needed, following `## Benchmark` below. + +If a trial needs another Docker image, stop the task, re-submit it, and +continue the same trial. This applies equally when a new trial's idea needs +a different image or serving framework: the cost of re-submitting a task is +not a reason to keep the current image or framework. You decide when a +trial is complete: better performance was achieved, or the ideas within this +trial are exhausted. + +Once a trial is completed, compile the interactive commands that produced +the final performance into a complete `dstack` task configuration with exact +commands, and log it together with the corresponding benchmark results (see +`## Benchmark` for the structure) to `trials.jsonl`. The benchmark may +be skipped in one case only: you failed to make the configuration run at +all — a failed trial. + +Each `trials.jsonl` record is one JSON line with exactly three fields: -If the endpoint config lists fleets, use only those fleets. Otherwise, use the -existing project/imported fleets supplied in the request. +``` +{"task": {...}, "resources": {...}, "benchmark": {...}} +``` -The request also lists fixed endpoint constraints such as max price, spot -policy, backends, regions, instance types, fleets, env keys, tags, or backend -options. Do not submit a task or service that conflicts with those values. +- `task`: the compiled `dstack` task configuration described above, as JSON. + Its `name` is the run name of the trial's task, and its `commands` are the + exact final commands that led to the benchmark results — the commands that + are supposed to replicate the benchmark results exactly if the task is + submitted (not `sleep infinity`). +- `resources`: the exact resources of the instance the task ran on, in + `dstack` resources syntax, e.g. `{"cpu": "9", "memory": "50GB", "disk": + "200GB", "gpu": {"name": "A40", "memory": "48GB", "count": 1}}`. Read the + actual values from the latest job submission's + `job_runtime_data.offer.instance.resources` in + `dstack run get --json`, converting MiB values to GB and the + `gpus` list into one `gpu` object with the GPU `name`, per-GPU `memory`, + and `count`. +- `benchmark`: the trial benchmark (see `## Benchmark` for the structure); + `null` only for a failed trial. + +Continue the trials until `max_trials` from `constraints.json` is reached or +you are entirely out of ideas on how to reach better performance. Before +stopping early, step back and think once more about what could still improve +performance within the constraints (see `# Constraints`). An early stop must +be justified in +`progress.jsonl` (see `# Progress`): report what you considered and why none +of it is worth a trial. + +Once the trials are over, pick the best trial and deploy it as a `dstack` +service to verify that it works and benchmark it finally (see +`# Final Service`). If +all is good, write the final report `final_report.json` (see +`# Final Report`). If the trial cannot be reproduced, pick the best trial +you have not verified yet and try again; proceed until a trial succeeds or +no trials remain. In that case, log the failure to `final_report.json` (see +`# Final Report`). + +## Benchmark + +During trials, run benchmarks via SSH inside the task, directly against the +serving engine: use `concurrency` from `constraints.json` and measure all +trials the same way so that their results are comparable with each other. +In trial benchmarks too, all measured requests must succeed. + +Before any benchmark — a trial one or the final one — warm the engine up by +verifying that the model works as expected: send real requests and check +the responses, including reasoning output when the model supports it. These +verification requests are never part of the measured metrics. + +Record every benchmark using the following structure and field names — +trial benchmarks in `trials.jsonl`, the final benchmark as +`final_report.json.benchmark` (values are illustrative): -Put each fixed constraint that has a dstack service YAML field into the final -service YAML. For example, if the endpoint is limited to a fleet, max price, or -spot policy, the final `service_yaml` should include the corresponding -`service_yaml.fleets`, `service_yaml.max_price`, or `service_yaml.spot_policy` -fields. Use `/dstack` for exact field names. +```json +{ + "tool": "vllm bench serve", + "tool_version": "0.11.0", + "command": "vllm bench serve ...", + "workload": {"api": "chat_completions", "num_requests": 16, "input_tokens": 1024, "output_tokens": 128, "concurrency": 8}, + "metrics": { + "successful_requests": 16, "failed_requests": 0, "duration_seconds": 4.0, + "total_input_tokens": 16384, "total_output_tokens": 2048, + "ttft_ms": {"mean": 110.9, "p50": 108.2, "p99": 121.6}, + "tpot_ms": {"mean": 7.5, "p50": 7.4, "p99": 8.1} + } +} +``` -If you cannot submit a useful task or service within the allowed fleets and -constraints, write a failed `final_report.json`. The -`final_report.json.failure_summary` value should say which fleet or constraint -blocked the run and what the user/admin would need to change. +Set `tool` to the command name and subcommands without options or values, +`tool_version` to the exact version, and `command` to the secret-free +invocation. For the final benchmark, run it with streaming responses, set +`workload.concurrency` to `concurrency` from `constraints.json`, produce +every field of the structure, and calculate all metrics from the +`num_requests` measured requests only — exclude setup, health-check, and +warmup requests; all measured requests must succeed. Never invent missing +values. # Task Usage -Use `/dstack-prototyping` to learn how to use tasks. Keep in mind that using a -task is a must. This means `sleep infinity` and directly attaching inside the -task via SSH to run commands as the only way to use tasks. If there is any -ambiguity, follow `/dstack-prototyping` and do nothing that contradicts it. +Trials are done entirely using `dstack` tasks. For maximum efficiency, it is a +requirement that you always set the task `commands` to `sleep infinity` and +run commands inside the task interactively, via SSH. It is important that +you follow the `/dstack-prototyping` skill when working with tasks. -# Backend/Fleet Selection, Idle Duration, Instance Volumes +# Hardware -Use only the fleets allowed by the endpoint request. If the endpoint request -also has constraints such as `backends`, `regions`, `instance_types`, -`spot_policy`, or `max_price`, apply them when running `dstack offer` if the CLI -has matching flags, and always apply them when submitting runs. +Pick the hardware before the trials start. You are allowed to change the +hardware from one trial to another if this can help the outcome. -## Backend And Fleet Choice +The available hardware is defined by the allowed `dstack` fleets and their +offers (coming from the configured backends). Follow the +`/dstack-prototyping` skill on how to efficiently select offers among the +available backends or SSH fleets. Pick the offer whose hardware best fits the trial's idea. Only when several +offers fit comparably, prefer backends or SSH fleets that support idle +instances/instance volumes: later runs reuse the instance and cached model +weights, while container-based backends start clean on every re-submission. -Use `/dstack-prototyping` to learn how to select backends and fleets. +If the backend allows, use instance volumes to mount cache and model weights +between runs. -Follow `/dstack-prototyping` skill on using a VM-based backend, Kubernetes backend, or SSH fleet that supports idle instances/instance volumes if there is such an option for the required GPU class. +When submitting a `dstack` task or service, pass exact `fleets`, `backends`, +and an intentional `resources` range based on the choice made from offers, so +the run does not land outside the intended fleet/backend/hardware. -If backend allows (see above), use instance volumes to mount cache and model weights between runs. +## Fleet Offers -## Validating Offers +When selecting `dstack` backends/fleets and evaluating hardware that can be +used to run the model, use `dstack offer --json` and pass the fleets from +`constraints.json` with `--fleet` explicitly. If `--fleet` is not passed, +`dstack offer` can show offers that are not applicable to the allowed fleets. +Use these offers when selecting fleet, backend, and hardware. -When selecting backends/fleets and evaluating hardware that can be used to run -the model, use `dstack offer --json` and pass `--fleet` explicitly. If the -endpoint request has `fleets`, pass those fleets. Otherwise, pass every existing -fleet. If `--fleet` is not passed, `dstack offer` can show offers that are not -applicable to the fleets allowed by the endpoint request. Use these offers when -selecting fleet, backend, and hardware. - -Choosing fleet/backend is gated by classifying each against `https://dstack.ai/docs/concepts/backends.md`. VM-based backends are listed under `## VM-based` (they support idle instances and instance volumes). Kubernetes backend is listed under `## Container-based`, but supports instance volumes and thus is preferred over other container-based backends. +To classify each backend's capabilities, fetch `https://dstack.ai/docs/concepts/backends.md` and classify from the fetched document, not from memory. VM-based backends are listed under `## VM-based` (they support idle instances and instance volumes). Kubernetes backend is listed under `## Container-based`, but supports instance volumes and thus is preferred over other container-based backends. SSH fleets can be treated as VM-based backends as they support both idle instances (its equivalent) and instance volumes. -## Model, Image, Serving Configuration, And Compute Fit - -Choose the model variant when allowed, image, serving configuration, and compute -to optimize expected performance within endpoint constraints and current offers. - -If `Endpoint context:` contains `model_repo`, choose fleet/backend/hardware that -can run `model_repo` within endpoint constraints. - -If `Endpoint context:` contains `base_model`, choose the repo/path and compute -together. You may pick the base repo or a compatible variant that fits the -allowed fleets, backends, and offers. - -If a task or service shows that the selected repo/path is a bad fit and -`Endpoint context:` contains `base_model`, pick another compatible variant if -available and test it in a task before submitting another service. - -## Submitting run - -When submitting a task or service, pass exact `fleets`, `backends`, and an -intentional `resources` range based on the choice made from offers, so the run -does not land outside the intended fleet/backend/hardware. - -## Decision Progress - -Progress should explain meaningful decisions and actions. - -When choosing a repo/path for `base_model`, fleet, backend, or hardware, write a -progress message that includes: - -- `service_model_name` and the selected repo/path, when they differ; -- the selected fleet, backend, and resource range; -- the offer/docs evidence used for the choice; -- the viable alternatives not selected and the exact reason; -- the fleet, backend, and resources that will be used in the submitted YAML. -- how the selected and rejected fleets/backends were classified; - -Backend-choice progress example (provide the same level of explanation for -repo/path, fleet, and hardware choice): - -"I chose backend ... on fleet ... because ...; I did not choose ... because ...; -I will submit YAML with fleets=..., backends=..., resources=...." - # Final Service -The final `service_yaml` is used to build the endpoint preset. It must -contain the full service config: `type: service`, the final run name, the service -model name, image/commands/port, resources, env references, and the fixed -endpoint constraints that apply to service YAML. - -Set final `service_yaml.name` to the final service run name. -If final `service_yaml.model` is a string, set it to `service_model_name` from -`Endpoint context:`. If final `service_yaml.model` is an object, set -`service_yaml.model.name` to `service_model_name`. - -Before submitting the final service, choose service resources from the least -restrictive requirements supported by the evidence, not from the exact machine -that happened to run. For example, if the model worked on an A40 but the -evidence only says it needs an NVIDIA GPU with at least 16GB memory, use that -broader requirement. Use an exact GPU, region, backend, or instance type only -when the endpoint constraints require it or the tested configuration depends on -that exact choice. - -Use run status to know whether the final service is still starting, running, or -failed. Use logs to understand failures. When dstack exposes the final service -run's `service.url`, build its absolute URL using `DSTACK_ENDPOINT_SERVER_URL` -and send a real model request using `DSTACK_ENDPOINT_BEARER_TOKEN` as the bearer -token. - -Verify the context length that the final service actually supports and report -it as `final_report.json.context_length`. - -# Benchmark +Once the trials are over, pick the best trial that has not been verified yet +and submit its configuration as a `dstack` service. Make it work with only +minor tweaks if needed; do not change the important decisions made during +the trial. Set the service `model` name to the client-facing model name from +`constraints.json` (see `# Constraints`). -Send benchmark requests to the final service run's absolute dstack service URL -using the bearer authentication used for service verification. Do not send them -to a server used during task prototyping or an SSH-forwarded local port. +Before the final benchmark, verify the model through the service: send real +requests using the client-facing model name and check that the model works +as it should, including reasoning output when the model supports it. This +verification also warms the service up. Only then run the final benchmark +(see `## Benchmark`). When verifying or benchmarking the service, use its +`service.url` reported by `dstack run get --json`, along with +`DSTACK_TOKEN` as the bearer token. If `service.url` is a relative path, +prepend `DSTACK_SERVER_URL` to build the absolute URL. -Run one benchmark with streaming responses. Choose a benchmark tool and workload -that can produce every required field below. +During the service verification, test the context length the service +actually supports and report it as `final_report.json.context_length`. -Report the benchmark as `final_report.json.benchmark` using exactly the -following structure and field names (values are illustrative): +If the service or its benchmark cannot be completed, stop that service, +pick the next-best trial, and repeat, until a service is verified or there +are no unverified trials left. Report the result accordingly (see +`# Final Report`). -```json -{ - "tool": "vllm bench serve", - "tool_version": "0.11.0", - "command": "vllm bench serve ...", - "workload": {"api": "chat_completions", "num_requests": 16, "input_tokens": 1024, "output_tokens": 128, "concurrency": 1}, - "metrics": { - "successful_requests": 16, "failed_requests": 0, "duration_seconds": 4.0, - "total_input_tokens": 16384, "total_output_tokens": 2048, - "ttft_ms": {"mean": 110.9, "p50": 108.2, "p99": 121.6}, - "tpot_ms": {"mean": 7.5, "p50": 7.4, "p99": 8.1} - } -} -``` - -Set `tool` to the command name and subcommands, without options or values (for -example, `vllm bench serve`). Set `tool_version` to the exact version and -`command` to the secret-free invocation. `api` must be `chat_completions` or -`completions`. `num_requests` is the number of measured requests; -`input_tokens` and `output_tokens` are the selected per-request lengths; and -`concurrency` is the maximum number of simultaneous requests. - -Calculate all metrics from the `num_requests` benchmark requests only. Exclude -setup, health-check, and warmup requests. `successful_requests` and -`failed_requests` are request counts; all requests must succeed. -`duration_seconds` is the elapsed wall-clock time from starting the first -measured request until the last measured request completes. `total_input_tokens` -and `total_output_tokens` are actual measured token totals. `ttft_ms` is the -time-to-first-token distribution across requests. `tpot_ms` is the -time-per-output-token distribution: for each request, divide the time from the -first output token to the last by one less than the actual output token count. -For both distributions, `mean`, `p50`, and `p99` are the arithmetic mean, 50th -percentile, and 99th percentile. Do not invent missing values. - -If benchmarking fails, write a failed `final_report.json`. - -Do not write a successful `final_report.json` until the final service answers a -request using `service_model_name` from `Endpoint context:` through the dstack -service URL. +When you report success, leave the final `dstack` service running: the +caller verifies the live run and stops it afterwards. # Secrets -The dstack CLI is already configured. Do not inspect `~/.dstack/config.yml` or +The `dstack` CLI is already configured. Do not inspect `~/.dstack/config.yml` or print, copy, or summarize tokens, secrets, or environment variable values. -Do not expose the value of `DSTACK_TOKEN`, `DSTACK_ENDPOINT_BEARER_TOKEN`, or -the value of any environment variable listed under `endpoint_env_keys` in -`Fixed endpoint constraints:`. +Do not expose the value of `DSTACK_TOKEN` or the value of any environment +variable listed under `env` in `constraints.json`. Do not put secret values in `final_report.json` or print them. Use env references in `final_report.json.service_yaml`; use environment variable names or redacted values in `final_report.json.benchmark.command`. -# Resume - -On startup or resume, inspect `final_report.json` and `submissions.jsonl` -before submitting anything new. If `final_report.json` already reports success -for the current endpoint configuration, do not submit another run. Otherwise use -the next run number and record why the old run is not enough. - # Final Report `final_report.json` may contain only `success`, `run_id`, `run_name`, @@ -301,20 +342,20 @@ On success, include exactly: - `success`: `true` - `run_id`: the final verified service run ID - `run_name`: the final verified service run name -- `service_yaml`: the full reusable service YAML described in `# Final Service` +- `service_yaml`: the full YAML of the verified final service - `base`: the base model repo, determined by the rules below - `model`: the exact repo/path loaded by the final service command - `context_length`: the context length verified for the final service -- `benchmark`: the final service benchmark described in `# Benchmark` +- `benchmark`: the final service benchmark described in `## Benchmark` Set `final_report.json.base` as follows: -- If `Endpoint context:` contains `base_model`, set `final_report.json.base` to - `base_model`. -- If `Endpoint context:` contains `model_repo`, inspect the repo metadata, model - card, config, or another reliable source to identify the base model repo. -- If `model_repo` is itself the base model repo, set `final_report.json.base` to - `model_repo`. +- If `model` in `constraints.json` has `base`, set `final_report.json.base` to + that value. +- If `model` has `repo`, inspect the repo metadata, model card, config, or + another reliable source to identify the base model repo. +- If `model.repo` is itself the base model repo, set `final_report.json.base` + to `model.repo`. - Do not infer `final_report.json.base` only from the repo name. On failure, include exactly: @@ -328,5 +369,11 @@ through `StructuredOutput`. Verify that `final_report.json` is correct and matches the required schema. -Stop after one correct service is verified and benchmarked. P/D disaggregation -is not covered by the current endpoint agent or `/dstack-prototyping` skill. +Stop only after `final_report.json` is written and submitted: either one +final `dstack` service was verified and benchmarked, or the trials and +unverified candidates were exhausted (see `# Trials` and `# Final Service`). + +Ending your turn stops the session even while background commands are still +running. Wait for long-running work — weight downloads, engine startup, +benchmarks, provisioning — by polling it with short blocking commands, +never by ending your turn. diff --git a/src/dstack/_internal/cli/services/endpoints/store.py b/src/dstack/_internal/cli/services/endpoints/store.py index bcb7c24db..561d88a47 100644 --- a/src/dstack/_internal/cli/services/endpoints/store.py +++ b/src/dstack/_internal/cli/services/endpoints/store.py @@ -1,47 +1,63 @@ import os +import shutil import sys import tempfile +from contextlib import suppress from pathlib import Path -from typing import List, TextIO +from typing import TextIO import yaml from pydantic import ValidationError from dstack._internal.cli.models.endpoint_presets import EndpointPreset -from dstack._internal.cli.models.endpoints import EndpointConfiguration +from dstack._internal.cli.models.endpoints import ( + MAX_PROMPT_LENGTH, + EndpointConfiguration, + EndpointPromptFile, +) from dstack._internal.cli.services.endpoints.presets import endpoint_preset_to_data +from dstack._internal.cli.utils.common import warn from dstack._internal.core.errors import CLIError, ConfigurationError from dstack._internal.utils.common import get_dstack_dir class EndpointPresetStore: + """Presets live at `//` — one directory per preset holding + the artifact (`preset.yaml`) next to the creation session internals. + Deleted presets are archived under `/.archive/`.""" + def __init__(self, root: Path | None = None) -> None: self.root = root or get_dstack_dir() / "presets" def list(self) -> list[EndpointPreset]: if not self.root.exists(): return [] - presets = [self._load(path) for path in self.root.glob("models--*/*.yaml")] + self._migrate_legacy() + presets = [self._load(path) for path in self.root.glob("*/preset.yaml")] return sorted(presets, key=lambda preset: (preset.base.lower(), preset.id)) def get(self, preset_id: str) -> EndpointPreset | None: - paths = self._find_preset_paths(preset_id) - if not paths: + _validate_preset_id(preset_id) + if not self.root.exists(): + return None + self._migrate_legacy() + path = self.root / preset_id / "preset.yaml" + if not path.is_file(): return None - if len(paths) > 1: - raise CLIError(f"Endpoint preset ID {preset_id!r} is not unique") - path = paths[0] preset = self._load(path) if preset.id != preset_id: - raise CLIError(f"Endpoint preset file {path} does not match its path") + raise CLIError(f"Preset file {path} does not match its path") return preset def save(self, preset: EndpointPreset) -> Path: - path = self._path(preset.base, preset.id) - path.parent.mkdir(parents=True, exist_ok=True) + _validate_preset_id(preset.id) + self._migrate_legacy() + directory = self.root / preset.id + directory.mkdir(parents=True, exist_ok=True) + path = directory / "preset.yaml" content = yaml.safe_dump(endpoint_preset_to_data(preset), sort_keys=False) fd, temporary_path = tempfile.mkstemp( - dir=path.parent, + dir=directory, prefix=f".{preset.id}.", suffix=".tmp", ) @@ -58,56 +74,62 @@ def save(self, preset: EndpointPreset) -> Path: pass return path + def find_by_name(self, name: str) -> EndpointPreset | None: + for preset in self.list(): + if preset.name == name: + return preset + return None + + def release_name(self, name: str) -> EndpointPreset | None: + """Releases `name` from the preset holding it, keeping the preset.""" + preset = self.find_by_name(name) + if preset is None: + return None + detached = preset.copy(update={"name": None}) + self.save(detached) + return detached + def delete(self, preset_id: str) -> bool: preset = self.get(preset_id) if preset is None: return False - path = self._path(preset.base, preset.id) - path.unlink() - try: - path.parent.rmdir() - except OSError: - pass + self._archive(self.root / preset_id) return True - def delete_for_base(self, base: str) -> int: - directory = self._directory(base) - paths = list(directory.glob("*.yaml")) - presets = [self._load(path) for path in paths] - if any(preset.base != base for preset in presets): - raise CLIError(f"Endpoint preset directory {directory} contains another base model") - for path in paths: - path.unlink() - try: - directory.rmdir() - except OSError: - pass - return len(presets) + def _archive(self, directory: Path) -> None: + archive_root = self.root / ".archive" + archive_root.mkdir(mode=0o700, parents=True, exist_ok=True) + target = archive_root / directory.name + index = 0 + while target.exists(): + index += 1 + target = archive_root / f"{directory.name}-{index}" + shutil.move(str(directory), str(target)) + + def _migrate_legacy(self) -> None: + for legacy in list(self.root.glob("models--*/*.yaml")): + target_dir = self.root / legacy.stem + target_dir.mkdir(mode=0o700, parents=True, exist_ok=True) + target = target_dir / "preset.yaml" + if target.exists(): + legacy.unlink() + else: + legacy.replace(target) + for directory in self.root.glob("models--*"): + with suppress(OSError): + directory.rmdir() def _load(self, path: Path) -> EndpointPreset: try: with path.open(encoding="utf-8") as f: return EndpointPreset.parse_obj(yaml.safe_load(f)) except (OSError, ValidationError, yaml.YAMLError) as e: - raise CLIError(f"Invalid endpoint preset file {path}: {e}") from e + raise CLIError(f"Invalid preset file {path}: {e}") from e - def _path(self, base: str, preset_id: str) -> Path: - if not preset_id or any(char in preset_id for char in "/\\"): - raise CLIError("Endpoint preset ID must not contain path separators") - return self._directory(base) / f"{preset_id}.yaml" - def _find_preset_paths(self, preset_id: str) -> List[Path]: - if not preset_id or any(char in preset_id for char in "/\\"): - raise CLIError("Endpoint preset ID must not contain path separators") - return [ - path - for directory in self.root.glob("models--*") - if (path := directory / f"{preset_id}.yaml").is_file() - ] - - def _directory(self, base: str) -> Path: - directory = "models--" + base.replace("/", "--").replace("\\", "--") - return self.root / directory +def _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_endpoint_configuration(path: str) -> tuple[str, EndpointConfiguration]: @@ -128,9 +150,39 @@ def _parse_endpoint_configuration(stream: TextIO) -> EndpointConfiguration: try: data = yaml.safe_load(stream) if not isinstance(data, dict): - raise ConfigurationError("Endpoint configuration must be a YAML object") - return EndpointConfiguration.parse_obj(data) + raise ConfigurationError("Preset configuration must be a YAML object") + configuration = EndpointConfiguration.parse_obj(data) except ValidationError as e: raise ConfigurationError(e) from e except yaml.YAMLError as e: - raise ConfigurationError(f"Invalid endpoint configuration: {e}") from e + 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_endpoint_prompt( + configuration: EndpointConfiguration, configuration_path: str +) -> str | None: + """The resolved user prompt text; file paths are relative to the configuration file.""" + if configuration.prompt is None: + return None + if isinstance(configuration.prompt, str): + return configuration.prompt.strip() + assert isinstance(configuration.prompt, EndpointPromptFile) + base = Path.cwd() if configuration_path == "-" else Path(configuration_path).parent + path = base / configuration.prompt.path + try: + text = path.read_text(encoding="utf-8").strip() + except (OSError, UnicodeDecodeError) as e: + raise ConfigurationError(f"Failed to read the prompt file {path}: {e}") from e + if not text: + raise ConfigurationError(f"The prompt file {path} is empty") + if len(text) > MAX_PROMPT_LENGTH: + raise ConfigurationError(f"The prompt file {path} exceeds {MAX_PROMPT_LENGTH} characters") + return text diff --git a/src/dstack/_internal/cli/services/endpoints/verify.py b/src/dstack/_internal/cli/services/endpoints/verify.py index d5f2872ec..f1917b1c4 100644 --- a/src/dstack/_internal/cli/services/endpoints/verify.py +++ b/src/dstack/_internal/cli/services/endpoints/verify.py @@ -28,6 +28,16 @@ from dstack._internal.core.models.runs import JobStatus, Run, RunStatus +def _redact_structure(value: Any, redacted_values: Sequence[str]) -> Any: + if isinstance(value, str): + return redact(value, redacted_values) + if isinstance(value, dict): + return {key: _redact_structure(item, redacted_values) for key, item in value.items()} + if isinstance(value, list): + return [_redact_structure(item, redacted_values) for item in value] + return value + + def load_endpoint_agent_report( *, output: EndpointAgentProcessOutput, @@ -42,6 +52,10 @@ def load_endpoint_agent_report( redacted_values, ) ) + # Scrub known secret values before validation: an echoed secret must never + # be persisted, but it also must not cost the whole session — the bearer + # check below still rejects unknown leaked tokens. + report_data = _redact_structure(report_data, redacted_values) try: report = AgentFinalReport.parse_obj(report_data) except ValidationError as e: @@ -61,6 +75,8 @@ def build_verified_endpoint_preset( run: Run, endpoint_configuration: EndpointConfiguration, report: AgentFinalReport, + preset_id: Optional[str] = None, + name: Optional[str] = None, ) -> EndpointPreset: 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") @@ -70,14 +86,14 @@ def build_verified_endpoint_preset( if not isinstance(service, ServiceConfiguration) or service.model is None: raise CLIError("Claude final run is not a model service") if service.model.name != endpoint_configuration.model.api_model_name: - raise CLIError("Claude final service model name does not match the endpoint request") + raise CLIError("Claude final service model name does not match the requested model") assert report.base is not None assert report.model is not None assert report.context_length is not None assert report.benchmark is not None if endpoint_configuration.model.allows_variant_selection: if report.base != endpoint_configuration.model.api_model_name: - raise CLIError("Claude final report base does not match the endpoint request") + raise CLIError("Claude final report base does not match the requested model") elif report.model != endpoint_configuration.model.exact_repo: raise CLIError("Claude changed an exact model request") if ( @@ -101,12 +117,14 @@ def build_verified_endpoint_preset( if isinstance(value, EnvSentinel) and key in portable_service.env: portable_service.env[key] = value return build_endpoint_preset( + name=name, service=portable_service, validation_replicas=_get_validation_replicas(run, service), base_model=report.base, model=report.model, context_length=report.context_length, benchmark=benchmark, + preset_id=preset_id, ) diff --git a/src/dstack/_internal/cli/utils/run.py b/src/dstack/_internal/cli/utils/run.py index dab2d264f..1a0577197 100644 --- a/src/dstack/_internal/cli/utils/run.py +++ b/src/dstack/_internal/cli/utils/run.py @@ -29,6 +29,7 @@ from dstack._internal.core.models.runs import ( ImagePullProgress, Job, + JobPlan, JobStatus, JobSubmission, Probe, @@ -174,6 +175,25 @@ def th(s: str) -> str: for key, value in (extra_properties or {}).items(): props.add_row(th(key), value) + console.print(props) + console.print() + print_offers( + job_plan, + max_offers=max_offers, + dim_after_first=include_run_properties, + show_offer_fleet_hint=show_offer_fleet_hint, + no_fleets=no_fleets, + ) + + +def print_offers( + job_plan: JobPlan, + *, + max_offers: Optional[int] = None, + dim_after_first: bool = True, + show_offer_fleet_hint: bool = False, + no_fleets: bool = False, +): offers = Table(box=None, expand=shutil.get_terminal_size(fallback=(120, 40)).columns <= 110) offers.add_column("#") offers.add_column("BACKEND", style="grey58", ratio=2) @@ -197,13 +217,11 @@ def th(s: str) -> str: instance, f"${offer.price:.4f}".rstrip("0").rstrip("."), format_instance_availability(offer.availability), - style=None if i == 1 or not include_run_properties else "secondary", + style=None if i == 1 or not dim_after_first else "secondary", ) if job_plan.total_offers > len(displayed_offers): offers.add_row("", "...", style="secondary") - console.print(props) - console.print() if len(displayed_offers) > 0: show_offer_fleet_hint_before_table = ( show_offer_fleet_hint @@ -429,7 +447,7 @@ def _format_instance_type( def _format_run_name(run: CoreRun, show_deployment_num: bool) -> str: - parts: List[str] = [run.run_spec.run_name] + parts: List[str] = [run.run_spec.run_name or ""] if show_deployment_num: parts.append(f" [secondary]deployment={run.deployment_num}[/]") return "".join(parts) diff --git a/src/tests/_internal/cli/commands/test_endpoint.py b/src/tests/_internal/cli/commands/test_endpoint.py index ac843c910..e6216e0b5 100644 --- a/src/tests/_internal/cli/commands/test_endpoint.py +++ b/src/tests/_internal/cli/commands/test_endpoint.py @@ -30,8 +30,8 @@ def test_formats_second_scale_ttft_without_scientific_notation(self): output = endpoint_presets_utils.format_endpoint_benchmark(preset, verbose=True) + assert output.startswith("ctx=32K ") assert "TTFT 8.15s" in output - assert "TTFT mean/p50/p99=8.15/8.15/8.33s" in output assert "e+03" not in output def test_preserves_benchmark_concurrency_at_narrow_width(self, monkeypatch): @@ -49,9 +49,9 @@ def test_preserves_benchmark_concurrency_at_narrow_width(self, monkeypatch): endpoint_presets_utils.print_endpoint_presets([get_endpoint_preset()]) - assert "concurrency=1" in "".join(output.getvalue().split()) + assert "con=1" in "".join(output.getvalue().split()) - def test_prints_created_column(self, monkeypatch): + def test_prints_submitted_column(self, monkeypatch): output = StringIO() monkeypatch.setattr( endpoint_presets_utils, @@ -67,24 +67,28 @@ def test_prints_created_column(self, monkeypatch): endpoint_presets_utils.print_endpoint_presets([get_endpoint_preset()]) - assert "CREATED" in output.getvalue() + assert "SUBMITTED" in output.getvalue() assert "2 months ago" in output.getvalue() def test_handles_keyboard_interrupt(self, tmp_path, capsys): configuration_path = tmp_path / "endpoint.dstack.yml" configuration_path.write_text( - "type: endpoint\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\n" + "type: preset\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\n" ) with ( patch("dstack.api.Client.from_config"), + patch( + "dstack._internal.cli.commands.endpoint.plan_endpoint_preset", + return_value=("fleet-a",), + ), patch( "dstack._internal.cli.commands.endpoint.create_endpoint_preset", side_effect=KeyboardInterrupt, ), ): exit_code = run_dstack_cli( - ["endpoint", "preset", "create", "-f", str(configuration_path)], + ["preset", "create", "-y", "-f", str(configuration_path)], home_dir=tmp_path, ) @@ -103,7 +107,7 @@ def test_lists_and_deletes_preset_without_api_client(self, tmp_path, capsys): "console", Console( file=list_output, - width=160, + width=250, color_system=None, theme=Theme({"secondary": "grey58"}), ), @@ -112,37 +116,39 @@ def test_lists_and_deletes_preset_without_api_client(self, tmp_path, capsys): endpoint_presets_utils, "pretty_date", return_value="2 months ago" ) as pretty_date, ): - assert run_dstack_cli(["endpoint", "preset", "list"], home_dir=tmp_path) == 0 + assert run_dstack_cli(["preset", "list"], home_dir=tmp_path) == 0 from_config.assert_not_called() pretty_date.assert_called_once_with(preset.created_at) + output = list_output.getvalue() + assert run_dstack_cli(["preset", "list", "-v"], home_dir=tmp_path) == 0 + verbose_output = list_output.getvalue()[len(output) :] - output = list_output.getvalue() assert "Qwen/Qwen3.5-27B" in output - assert "preset=8f3a12c4" in output - assert "repo=community/Qwen3.5-27B-GPTQ-Int4" in output - assert "CONTEXT" in output + assert "8f3a12c4" in output + # The repo row is shown only in verbose mode. + assert "repo=community/Qwen3.5-27B-GPTQ-Int4" not in output + # The context column is shown only in verbose mode. + assert "CONTEXT" not in output assert "BENCHMARK" in output - assert "32K" in output + assert "32K" not in output assert "42.1" in output - assert "concurrency=1" in "".join(output.split()) + assert "con=1" in "".join(output.split()) assert "tok/s" in output assert "TTFT" in output assert "108ms" in output assert "A6000:48GB:1" not in output - assert run_dstack_cli(["endpoint", "preset", "list", "-v"], home_dir=tmp_path) == 0 - verbose_output = capsys.readouterr().out - assert "hardware=" in verbose_output - assert "api=" in verbose_output - assert "n=16" in verbose_output - assert "concurrency=1" in "".join(verbose_output.split()) - assert "1K->128" in verbose_output + joined_verbose = "".join(verbose_output.split()) + # Verbose adds only the repo and the ctx= benchmark prefix. + assert "repo=community/Qwen3.5-27B-GPTQ-Int4" in joined_verbose + assert "ctx=32K" in joined_verbose + assert "con=1" in joined_verbose + assert "hardware=" not in joined_verbose with patch("dstack.api.Client.from_config") as from_config: assert ( run_dstack_cli( [ - "endpoint", "preset", "delete", preset.id, @@ -164,7 +170,7 @@ def test_gets_complete_preset_as_json_without_api_client(self, tmp_path, capsys) with patch("dstack.api.Client.from_config") as from_config: assert ( run_dstack_cli( - ["endpoint", "preset", "get", preset.id, "--json"], + ["preset", "get", preset.id, "--json"], home_dir=tmp_path, ) == 0 @@ -180,8 +186,8 @@ def test_gets_complete_preset_as_json_without_api_client(self, tmp_path, capsys) @pytest.mark.parametrize( "args", [ - ["endpoint", "preset", "--json"], - ["endpoint", "preset", "list", "--json"], + ["preset", "--json"], + ["preset", "list", "--json"], ], ) def test_lists_complete_presets_as_json(self, tmp_path, capsys, args): @@ -198,7 +204,9 @@ def test_lists_complete_presets_as_json(self, tmp_path, capsys, args): assert data["context_length"] == 32768 assert data["validations"][0]["benchmark"]["metrics"]["total_output_tokens"] == 2048 - def test_deletes_preset_without_api_client(self, tmp_path): + @pytest.mark.parametrize("flag_attribute", [("--base", "base"), ("--repo", "model")]) + def test_deletes_presets_by_model_without_api_client(self, tmp_path, flag_attribute): + flag, attribute = flag_attribute preset = get_endpoint_preset() store = EndpointPresetStore(tmp_path / ".dstack" / "presets") store.save(preset) @@ -207,7 +215,7 @@ def test_deletes_preset_without_api_client(self, tmp_path): with patch("dstack.api.Client.from_config") as from_config: assert ( run_dstack_cli( - ["endpoint", "preset", "delete", "--model", preset.base, "-y"], + ["preset", "delete", flag, getattr(preset, attribute), "-y"], home_dir=tmp_path, ) == 0 @@ -216,6 +224,41 @@ def test_deletes_preset_without_api_client(self, tmp_path): assert store.list() == [] + def test_delete_by_model_keeps_other_presets(self, tmp_path): + preset = get_endpoint_preset() + store = EndpointPresetStore(tmp_path / ".dstack" / "presets") + store.save(preset) + other = preset.copy( + update={"id": "01234567", "base": "meta/Llama-4", "model": "meta/Llama-4"} + ) + store.save(other) + + assert ( + run_dstack_cli( + ["preset", "delete", "--base", preset.base, "-y"], + home_dir=tmp_path, + ) + == 0 + ) + + assert [remaining.id for remaining in store.list()] == ["01234567"] + + @pytest.mark.parametrize("flag_attribute", [("--base", "base"), ("--repo", "model")]) + def test_lists_presets_filtered_by_model(self, tmp_path, capsys, flag_attribute): + flag, attribute = flag_attribute + preset = get_endpoint_preset() + store = EndpointPresetStore(tmp_path / ".dstack" / "presets") + store.save(preset) + store.save( + preset.copy(update={"id": "01234567", "base": "meta/Llama-4", "model": "meta/Llama-4"}) + ) + + args = ["preset", "list", "--json", flag, getattr(preset, attribute)] + assert run_dstack_cli(args, home_dir=tmp_path) == 0 + + output = json.loads(capsys.readouterr().out) + assert [entry["id"] for entry in output["presets"]] == [preset.id] + def test_merges_profile_configuration_and_cli_args(self, tmp_path): (tmp_path / ".dstack").mkdir() (tmp_path / ".dstack" / "profiles.yml").write_text( @@ -230,7 +273,7 @@ def test_merges_profile_configuration_and_cli_args(self, tmp_path): ) configuration_path = tmp_path / "endpoint.dstack.yml" configuration_path.write_text( - """type: endpoint + """type: preset name: file-name model: base: Qwen/Qwen3.5-27B @@ -249,6 +292,10 @@ def test_merges_profile_configuration_and_cli_args(self, tmp_path): with ( patch("dstack.api.Client.from_config"), + patch( + "dstack._internal.cli.commands.endpoint.plan_endpoint_preset", + return_value=("fleet-a",), + ), patch( "dstack._internal.cli.commands.endpoint.create_endpoint_preset", return_value=result, @@ -256,9 +303,9 @@ def test_merges_profile_configuration_and_cli_args(self, tmp_path): ): exit_code = run_dstack_cli( [ - "endpoint", "preset", "create", + "-y", "-f", str(configuration_path), "--name", @@ -286,17 +333,23 @@ def test_merges_profile_configuration_and_cli_args(self, tmp_path): assert create.call_args.kwargs["debug"] is True @pytest.mark.parametrize( - ("extra_args", "expected_preset"), - [([], "file-preset"), (["--preset", "cli-preset"], "cli-preset")], + ("extra_args", "expected_ids"), + [ + ([], None), + (["--id", "cli-preset"], ["cli-preset"]), + (["--id", "a1", "--id", "b2"], ["a1", "b2"]), + ], ) - def test_apply_passes_selected_profile_and_preset(self, tmp_path, extra_args, expected_preset): + def test_apply_passes_selected_profile_and_preset_ids( + self, tmp_path, extra_args, expected_ids + ): (tmp_path / ".dstack").mkdir() (tmp_path / ".dstack" / "profiles.yml").write_text( "profiles:\n - name: gpu\n max_price: 0.5\n" ) - configuration_path = tmp_path / "endpoint.dstack.yml" + configuration_path = tmp_path / "preset.dstack.yml" configuration_path.write_text( - "type: endpoint\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\npreset: file-preset\n" + "type: preset\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\n" ) with ( @@ -304,7 +357,6 @@ def test_apply_passes_selected_profile_and_preset(self, tmp_path, extra_args, ex patch("dstack._internal.cli.commands.endpoint.apply_endpoint_preset") as apply, ): args = [ - "endpoint", "preset", "apply", "-f", @@ -321,5 +373,99 @@ def test_apply_passes_selected_profile_and_preset(self, tmp_path, extra_args, ex assert exit_code == 0 assert apply.call_args.kwargs["profile_name"] == "gpu" - assert apply.call_args.kwargs["preset_id"] == expected_preset + assert apply.call_args.kwargs["preset_ids"] == expected_ids assert apply.call_args.kwargs["configuration"].max_price == 0.5 + + +class TestPresetNameClaims: + def test_create_detaches_the_name_from_the_old_preset(self, tmp_path, monkeypatch): + preset = get_endpoint_preset().copy(update={"name": "qwen"}) + store = EndpointPresetStore(tmp_path / ".dstack" / "presets") + store.save(preset) + configuration_path = tmp_path / "preset.dstack.yml" + configuration_path.write_text("type: preset\nname: qwen\nbase: Qwen/Qwen3.5-27B\n") + result = SimpleNamespace( + preset=preset, path=tmp_path / "preset.yaml", final_run_name="qwen-1" + ) + + with ( + patch("dstack.api.Client.from_config"), + patch( + "dstack._internal.cli.commands.endpoint.plan_endpoint_preset", + return_value=("fleet-a",), + ), + patch( + "dstack._internal.cli.commands.endpoint.create_endpoint_preset", + return_value=result, + ) as create, + ): + exit_code = run_dstack_cli( + ["preset", "create", "-f", str(configuration_path), "-y"], + home_dir=tmp_path, + repo_dir=tmp_path, + ) + + assert exit_code == 0 + create.assert_called_once() + assert store.get(preset.id).name is None + + def test_create_without_confirmation_exits_before_creating(self, tmp_path, capsys): + preset = get_endpoint_preset().copy(update={"name": "qwen"}) + store = EndpointPresetStore(tmp_path / ".dstack" / "presets") + store.save(preset) + configuration_path = tmp_path / "preset.dstack.yml" + configuration_path.write_text("type: preset\nname: qwen\nbase: Qwen/Qwen3.5-27B\n") + + with ( + patch("dstack.api.Client.from_config"), + patch( + "dstack._internal.cli.commands.endpoint.plan_endpoint_preset", + return_value=("fleet-a",), + ), + patch("dstack._internal.cli.commands.endpoint.confirm_ask", return_value=False), + patch("dstack._internal.cli.commands.endpoint.create_endpoint_preset") as create, + ): + exit_code = run_dstack_cli( + ["preset", "create", "-f", str(configuration_path)], + home_dir=tmp_path, + repo_dir=tmp_path, + ) + + assert exit_code == 0 + create.assert_not_called() + assert store.get(preset.id).name == "qwen" + + def test_get_and_delete_resolve_names(self, tmp_path, capsys): + preset = get_endpoint_preset().copy(update={"name": "qwen"}) + EndpointPresetStore(tmp_path / ".dstack" / "presets").save(preset) + + assert run_dstack_cli(["preset", "get", "qwen", "--json"], home_dir=tmp_path) == 0 + assert json.loads(capsys.readouterr().out)["id"] == preset.id + + assert run_dstack_cli(["preset", "delete", "qwen", "-y"], home_dir=tmp_path) == 0 + assert EndpointPresetStore(tmp_path / ".dstack" / "presets").list() == [] + + def test_create_always_asks_even_without_a_name_conflict(self, tmp_path): + configuration_path = tmp_path / "preset.dstack.yml" + configuration_path.write_text("type: preset\nbase: Qwen/Qwen3.5-27B\n") + + with ( + patch("dstack.api.Client.from_config"), + patch( + "dstack._internal.cli.commands.endpoint.plan_endpoint_preset", + return_value=("fleet-a",), + ), + patch( + "dstack._internal.cli.commands.endpoint.confirm_ask", return_value=False + ) as confirm, + patch("dstack._internal.cli.commands.endpoint.create_endpoint_preset") as create, + ): + exit_code = run_dstack_cli( + ["preset", "create", "-f", str(configuration_path)], + home_dir=tmp_path, + repo_dir=tmp_path, + ) + + assert exit_code == 0 + confirm.assert_called_once_with("Create the preset?") + create.assert_not_called() diff --git a/src/tests/_internal/cli/models/test_endpoints.py b/src/tests/_internal/cli/models/test_endpoints.py index a507a1a01..5aa4b6cc1 100644 --- a/src/tests/_internal/cli/models/test_endpoints.py +++ b/src/tests/_internal/cli/models/test_endpoints.py @@ -49,3 +49,36 @@ def test_parses_exact_repo_with_client_facing_name(self): def test_rejects_ambiguous_model_object(self): with pytest.raises(ValidationError): EndpointConfiguration(model={"base": "Qwen/base", "repo": "Qwen/repo"}) + + def test_parses_top_level_base_shorthand(self): + configuration = EndpointConfiguration(base="Qwen/Qwen3.5-27B") + + assert isinstance(configuration.model, EndpointModelBase) + assert configuration.model.api_model_name == "Qwen/Qwen3.5-27B" + assert configuration.base is None + + def test_parses_top_level_repo_shorthand(self): + configuration = EndpointConfiguration(repo="community/Qwen3.5-27B-GPTQ-Int4") + + assert isinstance(configuration.model, EndpointModelRepo) + assert configuration.model.exact_repo == "community/Qwen3.5-27B-GPTQ-Int4" + assert configuration.repo is None + + def test_shorthand_round_trips_through_dict(self): + configuration = EndpointConfiguration(base="Qwen/Qwen3.5-27B") + + round_tripped = EndpointConfiguration.parse_obj(configuration.dict()) + + assert round_tripped.model == configuration.model + + def test_rejects_combined_base_and_repo_shorthand(self): + with pytest.raises(ValidationError): + EndpointConfiguration(base="Qwen/base", repo="Qwen/repo") + + def test_rejects_shorthand_combined_with_model(self): + with pytest.raises(ValidationError): + EndpointConfiguration(base="Qwen/base", model={"repo": "Qwen/repo"}) + + def test_requires_model(self): + with pytest.raises(ValidationError): + EndpointConfiguration() diff --git a/src/tests/_internal/cli/services/endpoints/test_agent.py b/src/tests/_internal/cli/services/endpoints/test_agent.py index c3c13ef52..a8fabd89b 100644 --- a/src/tests/_internal/cli/services/endpoints/test_agent.py +++ b/src/tests/_internal/cli/services/endpoints/test_agent.py @@ -2,8 +2,11 @@ import json import os import shutil +import signal import subprocess import sys +from contextlib import suppress +from pathlib import Path from types import SimpleNamespace import psutil @@ -16,17 +19,21 @@ EndpointAgentSession, EndpointAgentWorkspace, _build_claude_command, - _create_agent_session_directory, _prepare_subprocess_command, _ProgressTailer, + _RecordMirror, + _summarize_session_trials, _terminate_process, + attach_agent_workspace, build_endpoint_agent_env, contains_redacted_value, + create_agent_workspace, create_endpoint_agent_session, - endpoint_agent_workspace, get_claude_auth, + load_resumable_agent_session, print_endpoint_progress, redact, + remove_agent_workspace, run_endpoint_agent, ) from dstack._internal.compat import IS_WINDOWS @@ -120,8 +127,9 @@ def test_inherits_only_required_environment(self, tmp_path, monkeypatch): assert project.url == "http://127.0.0.1:3000" assert project.token == "dstack-secret" - def test_creates_private_cli_home_and_dstack_wrapper(self): - with endpoint_agent_workspace() as workspace: + def test_creates_private_cli_home_and_dstack_wrapper(self, tmp_path): + workspace = _session_workspace(tmp_path) + if True: if not IS_WINDOWS: assert (workspace.dstack_home / ".ssh").stat().st_mode & 0o777 == 0o700 assert not (workspace.dstack_home / ".dstack" / "config.yml").exists() @@ -136,8 +144,9 @@ def test_creates_private_cli_home_and_dstack_wrapper(self): assert result.returncode == 0 assert "Usage: dstack" in result.stdout - def test_keeps_control_socket_path_bounded(self): - with endpoint_agent_workspace() as workspace: + def test_keeps_control_socket_path_bounded(self, tmp_path): + workspace = _session_workspace(tmp_path) + if True: if not IS_WINDOWS: socket_path = ( workspace.dstack_home / ".dstack" / "ssh" / f"{'x' * 41}.control.sock" @@ -151,6 +160,15 @@ 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 = EndpointAgentSession( + path=session_dir, timestamp="t", debug=False, preset_id="abcd1234" + ) + return create_agent_workspace(session) + + class TestAgentSession: def test_saves_log_and_debug_files(self, tmp_path, monkeypatch, capsys): monkeypatch.setenv("HOME", str(tmp_path)) @@ -164,9 +182,20 @@ def test_saves_log_and_debug_files(self, tmp_path, monkeypatch, capsys): session = create_endpoint_agent_session(configuration) - assert session.path.parent == tmp_path / ".dstack" / "agent" / "qwen" - assert session.path.name.endswith("-running") - assert {path.name for path in session.path.iterdir()} == {"agent.log"} + assert session.path.parent == tmp_path / ".dstack" / "presets" + assert session.path.name == session.preset_id + assert len(session.preset_id) == 8 + assert {path.name for path in session.path.iterdir()} == { + "agent.log", + "session.json", + "preset.dstack.yml", + } + manifest = json.loads((session.path / "session.json").read_text()) + assert manifest["id"] == session.preset_id + assert manifest["status"] == "running" + assert manifest["endpoint"] == "qwen" + assert manifest["model"] == "Qwen/Qwen3.5-27B" + assert manifest["pid"] == os.getpid() print_endpoint_progress("creating preset", agent_session=session) assert "creating preset" in session.log_path.read_text() assert "creating preset" in capsys.readouterr().out @@ -175,55 +204,46 @@ def test_saves_log_and_debug_files(self, tmp_path, monkeypatch, capsys): assert session.log_path.stat().st_mode & 0o777 == 0o600 debug_session = create_endpoint_agent_session(configuration, debug=True) - data = yaml.safe_load((debug_session.path / "endpoint.dstack.yml").read_text()) + data = yaml.safe_load((debug_session.path / "preset.dstack.yml").read_text()) assert {path.name for path in debug_session.path.iterdir()} == { "agent.log", - "endpoint.dstack.yml", + "preset.dstack.yml", + "session.json", "trace.jsonl", } assert data["max_price"] == 0.5 assert data["env"] == ["HF_TOKEN", "TOKENIZERS_PARALLELISM"] - assert "false" not in (debug_session.path / "endpoint.dstack.yml").read_text() - success_path = debug_session.finish("8f3a12c4") - assert success_path.name.endswith("-8f3a12c4") + assert "false" not in (debug_session.path / "preset.dstack.yml").read_text() + success_path = debug_session.finish("success") + assert success_path == debug_session.path + assert json.loads((debug_session.path / "session.json").read_text())["status"] == "success" failed_session = create_endpoint_agent_session(configuration) - failed_path = failed_session.finish() - assert failed_path.name.endswith("-failed") - - def test_avoids_existing_session_directories(self, tmp_path): - timestamp = "20260714-120000-000000Z" - (tmp_path / f"{timestamp}-running").mkdir() - - path = _create_agent_session_directory(tmp_path, timestamp) - - assert path.name == f"{timestamp}-1-running" - - def test_finish_does_not_replace_existing_directory(self, tmp_path): - running = tmp_path / "20260714-120000-000000Z-running" - running.mkdir() - (running / "agent.log").touch() - existing = tmp_path / "20260714-120000-000000Z-failed" - existing.mkdir() - (existing / "keep").touch() + failed_path = failed_session.finish("failed") + assert failed_path == failed_session.path + assert json.loads((failed_session.path / "session.json").read_text())["status"] == "failed" + + 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 = EndpointAgentSession( - path=running, + path=session_dir, timestamp="20260714-120000-000000Z", debug=False, ) - path = session.finish() + path = session.finish("failed") - assert path.name == "20260714-120000-000000Z-failed-1" - assert (existing / "keep").is_file() + assert path == session_dir + assert json.loads((session_dir / "session.json").read_text()) == {"status": "failed"} def test_reports_invalid_existing_parent(self, tmp_path, monkeypatch): monkeypatch.setenv("HOME", str(tmp_path)) monkeypatch.setenv("USERPROFILE", str(tmp_path)) - parent = tmp_path / ".dstack" / "agent" - parent.mkdir(parents=True) - (parent / "qwen").write_text("not a directory") + (tmp_path / ".dstack").mkdir(parents=True) + (tmp_path / ".dstack" / "presets").write_text("not a directory") with pytest.raises(CLIError, match="Could not create agent output"): create_endpoint_agent_session( @@ -391,3 +411,516 @@ async def test_terminates_windows_process_tree(self): assert proc.returncode is not None assert not psutil.pid_exists(child_pid) + + +class TestRecordMirror: + def test_mirrors_complete_lines_redacted(self, tmp_path): + source = tmp_path / "runs.jsonl" + target = tmp_path / "mirror" / "runs.jsonl" + target.parent.mkdir() + mirror = _RecordMirror(source=source, target=target, redacted_values=["dstack-secret"]) + + source.write_text( + '{"name":"run-1","note":"dstack-secret"}\n{"name":"run-2"', encoding="utf-8" + ) + mirror.flush() + + assert target.read_text() == '{"name":"run-1","note":"[redacted]"}\n' + + with source.open("a", encoding="utf-8") as f: + f.write("}\n") + mirror.flush() + + assert target.read_text().splitlines() == [ + '{"name":"run-1","note":"[redacted]"}', + '{"name":"run-2"}', + ] + + def test_missing_source_is_no_op(self, tmp_path): + mirror = _RecordMirror( + source=tmp_path / "absent.jsonl", + target=tmp_path / "target.jsonl", + redacted_values=[], + ) + + mirror.flush() + + assert not (tmp_path / "target.jsonl").exists() + + +class TestWriteAgentInfo: + def test_writes_model_params_and_auth(self, tmp_path, monkeypatch): + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.agent._get_claude_version", + lambda auth: "2.1.0 (Claude Code)", + ) + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.agent._get_claude_auth_status", + lambda auth: {"authMethod": "claude.ai", "loggedIn": True}, + ) + session_dir = tmp_path / "session" + session_dir.mkdir() + session = EndpointAgentSession(path=session_dir, timestamp="t", debug=True) + + session.write_agent_info( + ClaudeAuth(api_key=None, executable="claude", effort=None, model="claude-opus-4-8") + ) + + assert json.loads((session_dir / "agent.json").read_text()) == { + "executable": "claude", + "version": "2.1.0 (Claude Code)", + "model": {"name": "claude-opus-4-8", "effort": "default"}, + "auth": {"authMethod": "claude.ai", "loggedIn": True}, + } + + +class TestConnectionResume: + def _write_flaky_claude(self, tmp_path): + script = tmp_path / "fake_claude.py" + script.write_text( + """import json +import sys +from pathlib import Path + +args = sys.argv[1:] +with Path("calls.jsonl").open("a") as f: + f.write(json.dumps({"args": args, "prompt": sys.stdin.read()}) + "\\n") +if "--resume" in args: + print(json.dumps({ + "type": "result", + "session_id": "sid-123", + "structured_output": {"resumed": True}, + })) +else: + print(json.dumps({"type": "system", "subtype": "init", "session_id": "sid-123"})) + print(json.dumps({ + "type": "result", + "is_error": True, + "result": "API Error: Unable to connect to API (ENOTFOUND)", + })) + sys.exit(1) +""" + ) + return script + + def _agent_setup(self, tmp_path): + (tmp_path / "progress.jsonl").touch() + workspace = EndpointAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home") + session_path = tmp_path / "session" + session_path.mkdir() + (session_path / "agent.log").touch() + agent_session = EndpointAgentSession( + path=session_path, + timestamp="20260714-120000Z", + debug=False, + ) + return workspace, agent_session + + @pytest.mark.asyncio + async def test_resumes_after_connection_error(self, tmp_path, monkeypatch, capsys): + script = self._write_flaky_claude(tmp_path) + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.agent._RESUME_DELAYS_SECONDS", (0,) + ) + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.agent._build_claude_command", + lambda **kwargs: [sys.executable, str(script)] + + ( + ["--resume", kwargs["resume_session_id"]] + if kwargs.get("resume_session_id") + else [] + ), + ) + workspace, agent_session = self._agent_setup(tmp_path) + + output = await run_endpoint_agent( + prompt="system prompt", + env={}, + workspace=workspace, + auth=ClaudeAuth(api_key=None, executable="claude", effort=None, model="m"), + redacted_values=(), + agent_session=agent_session, + ) + + assert output.report_data == {"resumed": True} + calls = [json.loads(line) for line in (tmp_path / "calls.jsonl").read_text().splitlines()] + assert len(calls) == 2 + assert calls[0]["args"] == [] + assert calls[0]["prompt"] == "system prompt" + assert calls[1]["args"] == ["--resume", "sid-123"] + assert "The previous agent process was interrupted" in calls[1]["prompt"] + assert "resuming the session" in capsys.readouterr().out + + @pytest.mark.asyncio + async def test_resumes_on_any_unreported_death(self, tmp_path, monkeypatch): + script = tmp_path / "fake_claude.py" + script.write_text( + """import json +import sys +from pathlib import Path + +with Path("calls.jsonl").open("a") as f: + f.write("call\\n") +if "--resume" in sys.argv[1:]: + print(json.dumps({"type": "result", "structured_output": {"recovered": True}})) +else: + print(json.dumps({"type": "system", "subtype": "init", "session_id": "sid-123"})) + print(json.dumps({"type": "result", "is_error": True, "result": "model refused"})) + sys.exit(1) +""" + ) + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.agent._RESUME_DELAYS_SECONDS", (0,) + ) + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.agent._build_claude_command", + lambda **kwargs: [sys.executable, str(script)] + + ( + ["--resume", kwargs["resume_session_id"]] + if kwargs.get("resume_session_id") + else [] + ), + ) + workspace, agent_session = self._agent_setup(tmp_path) + + output = await run_endpoint_agent( + prompt="system prompt", + env={}, + workspace=workspace, + auth=ClaudeAuth(api_key=None, executable="claude", effort=None, model="m"), + redacted_values=(), + agent_session=agent_session, + ) + + assert output.report_data == {"recovered": True} + assert len((tmp_path / "calls.jsonl").read_text().splitlines()) == 2 + + +class TestConnectionRetryExhaustion: + @pytest.mark.asyncio + async def test_gives_up_after_repeated_no_progress_failures(self, tmp_path, monkeypatch): + script = tmp_path / "fake_claude.py" + script.write_text( + """import json +import sys +from pathlib import Path + +with Path("calls.jsonl").open("a") as f: + f.write("call\\n") +print(json.dumps({"type": "system", "subtype": "init", "session_id": "sid-123"})) +print(json.dumps({ + "type": "result", + "is_error": True, + "result": "API Error: Unable to connect to API (ENOTFOUND)", +})) +sys.exit(1) +""" + ) + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.agent._RESUME_DELAYS_SECONDS", (0, 0) + ) + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.agent._build_claude_command", + lambda **kwargs: [sys.executable, str(script)], + ) + (tmp_path / "progress.jsonl").touch() + workspace = EndpointAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home") + session_path = tmp_path / "session" + session_path.mkdir() + (session_path / "agent.log").touch() + agent_session = EndpointAgentSession( + path=session_path, timestamp="20260714-120000Z", debug=False + ) + + output = await run_endpoint_agent( + prompt="system prompt", + env={}, + workspace=workspace, + auth=ClaudeAuth(api_key=None, executable="claude", effort=None, model="m"), + redacted_values=(), + agent_session=agent_session, + ) + + assert output.report_data is None + assert output.error is not None and "Unable to connect" in output.error + # Initial attempt + one retry per configured delay, then give up. + assert len((tmp_path / "calls.jsonl").read_text().splitlines()) == 3 + + +class TestWorkspaceLifecycle: + def _session(self, tmp_path): + session_dir = tmp_path / "sessions" / "ab12cd34" + session_dir.mkdir(parents=True) + return EndpointAgentSession( + path=session_dir, timestamp="t", debug=False, preset_id="ab12cd34" + ) + + @pytest.mark.skipif(IS_WINDOWS, reason="workspace alias symlinks are POSIX-only") + def test_create_attach_and_remove(self, tmp_path): + session = self._session(tmp_path) + + workspace = create_agent_workspace(session) + manifest = session.read_manifest() + alias = Path(manifest["alias"]) + assert Path(manifest["workspace"]) == session.path / "workspace" + assert alias.is_symlink() + (workspace.path / "note.txt").write_text("kept", encoding="utf-8") + + os.unlink(alias) + attached = attach_agent_workspace(session) + assert alias.is_symlink() + assert (attached.path / "note.txt").read_text() == "kept" + + remove_agent_workspace(session) + assert not os.path.lexists(alias) + assert not (session.path / "workspace").exists() + + @pytest.mark.skipif(IS_WINDOWS, reason="workspace alias symlinks are POSIX-only") + def test_attach_refuses_occupied_alias(self, tmp_path): + session = self._session(tmp_path) + create_agent_workspace(session) + manifest = session.read_manifest() + alias = Path(manifest["alias"]) + os.unlink(alias) + alias.mkdir() + try: + with pytest.raises(CLIError, match="cannot be used safely"): + attach_agent_workspace(session) + finally: + alias.rmdir() + + def test_attach_fails_when_workspace_is_gone(self, tmp_path): + session = self._session(tmp_path) + create_agent_workspace(session) + remove_agent_workspace(session) + with pytest.raises(CLIError, match="no longer exists"): + attach_agent_workspace(session) + + +class TestOffsetPersistence: + def test_mirror_does_not_duplicate_after_restart(self, tmp_path): + from dstack._internal.cli.services.endpoints.agent import _OffsetStore + + source = tmp_path / "runs.jsonl" + target = tmp_path / "mirror.jsonl" + state = tmp_path / ".offsets.json" + source.write_text('{"name":"one"}\n', encoding="utf-8") + + mirror = _RecordMirror( + source=source, + target=target, + redacted_values=(), + offset_store=_OffsetStore(state), + offset_key="runs", + ) + mirror.flush() + + with source.open("a", encoding="utf-8") as f: + f.write('{"name":"two"}\n') + restarted = _RecordMirror( + source=source, + target=target, + redacted_values=(), + offset_store=_OffsetStore(state), + offset_key="runs", + ) + restarted.flush() + + assert target.read_text().splitlines() == ['{"name":"one"}', '{"name":"two"}'] + + +class TestLoadResumableSession: + def _write_session(self, tmp_path, monkeypatch, manifest): + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + path = tmp_path / ".dstack" / "presets" / manifest["id"] + path.mkdir(parents=True) + (path / "session.json").write_text(json.dumps(manifest), encoding="utf-8") + return path + + def test_loads_interrupted_session(self, tmp_path, monkeypatch): + self._write_session( + tmp_path, + monkeypatch, + { + "id": "ab12cd34", + "status": "interrupted", + "claude_session_id": "sid-1", + "debug": True, + "created_at": "2026-07-20T10:00:00+00:00", + }, + ) + + session = load_resumable_agent_session("ab12cd34") + + assert session.preset_id == "ab12cd34" + assert session.debug is True + + def test_treats_dead_running_session_as_resumable(self, tmp_path, monkeypatch): + self._write_session( + tmp_path, + monkeypatch, + { + "id": "ab12cd34", + "status": "running", + "pid": 4242, + "claude_session_id": "sid-1", + }, + ) + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.agent.psutil.pid_exists", + lambda pid: False, + ) + + assert load_resumable_agent_session("ab12cd34").preset_id == "ab12cd34" + + def test_refusals(self, tmp_path, monkeypatch): + with pytest.raises(CLIError, match="Unknown preset creation session"): + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + load_resumable_agent_session("00000000") + + self._write_session(tmp_path, monkeypatch, {"id": "aa000001", "status": "success"}) + with pytest.raises(CLIError, match="nothing to resume"): + load_resumable_agent_session("aa000001") + + self._write_session(tmp_path, monkeypatch, {"id": "aa000002", "status": "failed"}) + with pytest.raises(CLIError, match="cannot be resumed"): + load_resumable_agent_session("aa000002") + + self._write_session( + tmp_path, + monkeypatch, + { + "id": "aa000003", + "status": "running", + "pid": 4242, + "claude_session_id": "sid-1", + }, + ) + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.agent.psutil.pid_exists", + lambda pid: True, + ) + with pytest.raises(CLIError, match="running already"): + load_resumable_agent_session("aa000003") + + self._write_session(tmp_path, monkeypatch, {"id": "aa000004", "status": "interrupted"}) + with pytest.raises(CLIError, match="stopped before the agent started"): + load_resumable_agent_session("aa000004") + + +class TestSummarizeSessionTrials: + def test_counts_records_even_when_trials_share_a_task(self, tmp_path): + record = { + "task": {"name": "qwen-ab12cd34-1"}, + "resources": {"gpu": {"name": "A40", "memory": "48GB", "count": 1}}, + "benchmark": { + "workload": {"concurrency": 8}, + "metrics": {"duration_seconds": 10.0, "total_output_tokens": 20000}, + }, + } + rerun = json.loads(json.dumps(record)) + rerun["benchmark"]["metrics"]["total_output_tokens"] = 23000 + second_trial = json.loads(json.dumps(record)) + second_trial["task"]["name"] = "qwen-ab12cd34-2" + nameless = {"benchmark": {"workload": {}, "metrics": {}}} + path = tmp_path / "trials.jsonl" + path.write_text( + "\n".join(json.dumps(entry) for entry in [record, rerun, second_trial, nameless]) + ) + + summary = _summarize_session_trials(path) + + # 4 records = 4 trials: one long-lived task commonly hosts several + # trials, so shared task names must not collapse the count. + assert summary["count"] == 4 + assert summary["best"] == {"tok_s": 2300.0, "concurrency": 8, "gpu": "A40:48GB:1"} + + +class TestFileLineReader: + @pytest.mark.asyncio + async def test_reads_lines_and_continues_from_persisted_offset(self, tmp_path): + from dstack._internal.cli.services.endpoints.agent import _FileLineReader, _OffsetStore + + stream = tmp_path / "stdout.jsonl" + state = tmp_path / ".offsets.json" + stream.write_bytes(b"first\nsecond\n") + alive = True + + reader = _FileLineReader( + stream, + offset_store=_OffsetStore(state), + offset_key="agent_stdout", + is_alive=lambda: alive, + ) + assert await reader.readline() == b"first\n" + assert await reader.readline() == b"second\n" + + # A new reader (a later attach) continues where the previous stopped. + with stream.open("ab") as f: + f.write(b"third\npartial") + alive = False + attached = _FileLineReader( + stream, + offset_store=_OffsetStore(state), + offset_key="agent_stdout", + is_alive=lambda: alive, + ) + assert await attached.readline() == b"third\n" + assert await attached.readline() == b"partial" + assert await attached.readline() == b"" + + +class TestStopOrDetach: + @pytest.mark.skipif(IS_WINDOWS, reason="exercises POSIX process groups") + def test_detach_keeps_the_agent_and_stop_terminates_it(self, tmp_path, monkeypatch, capsys): + from dstack._internal.cli.services.endpoints import create as create_module + from dstack._internal.cli.services.endpoints.create import _stop_or_detach_agent_session + + session_dir = tmp_path / "ab12cd34" + session_dir.mkdir() + (session_dir / "agent.log").touch() + session = EndpointAgentSession( + path=session_dir, timestamp="t", debug=False, preset_id="ab12cd34" + ) + agent = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(300)"], start_new_session=True + ) + try: + session.update_manifest(status="running", agent_pid=agent.pid) + + monkeypatch.setattr(create_module, "confirm_ask", lambda *_: False) + _stop_or_detach_agent_session(session) + assert agent.poll() is None + assert session.read_manifest()["status"] == "running" + assert "Detached" in capsys.readouterr().out + + monkeypatch.setattr(create_module, "confirm_ask", lambda *_: True) + _stop_or_detach_agent_session(session) + psutil.Process(agent.pid).wait(timeout=10) + assert session.read_manifest()["status"] == "interrupted" + finally: + with suppress(OSError): + os.killpg(agent.pid, signal.SIGKILL) + + +class TestOffsetStoreSharing: + def test_disjoint_stores_do_not_clobber_each_other(self, tmp_path): + from dstack._internal.cli.services.endpoints.agent import _OffsetStore + + state = tmp_path / ".offsets.json" + readers = _OffsetStore(state) + mirrors = _OffsetStore(state) + + # Two stores with disjoint keys, interleaved writes to the same file. + readers.set("agent_stdout", 10) + mirrors.set("runs", 20) + readers.set("agent_stderr", 30) + mirrors.set("trials", 40) + + reloaded = _OffsetStore(state) + assert reloaded.get("agent_stdout") == 10 + assert reloaded.get("agent_stderr") == 30 + assert reloaded.get("runs") == 20 + assert reloaded.get("trials") == 40 diff --git a/src/tests/_internal/cli/services/endpoints/test_apply.py b/src/tests/_internal/cli/services/endpoints/test_apply.py index c9962ffd3..fab87e900 100644 --- a/src/tests/_internal/cli/services/endpoints/test_apply.py +++ b/src/tests/_internal/cli/services/endpoints/test_apply.py @@ -6,10 +6,12 @@ from dstack._internal.cli.models.endpoints import EndpointConfiguration from dstack._internal.cli.services.endpoints.apply import ( _build_service, + _get_candidate_presets, _get_matching_presets, _select_plan, apply_endpoint_preset, ) +from dstack._internal.core.errors import CLIError from dstack._internal.core.models.instances import InstanceAvailability from tests._internal.cli.endpoint_presets import get_endpoint_preset @@ -28,13 +30,25 @@ def test_matches_base_model_context_and_preset(self): context_length=8192, ) - assert _get_matching_presets(presets, configuration=configuration, preset_id=None) == [ - presets[1] - ] - assert _get_matching_presets(presets, configuration=configuration, preset_id="large") == [ - presets[1] + assert _get_matching_presets(presets, configuration=configuration) == [presets[1]] + assert _get_matching_presets( + _get_candidate_presets(presets, preset_ids=["large"]), configuration=configuration + ) == [presets[1]] + assert not _get_matching_presets( + _get_candidate_presets(presets, preset_ids=["small"]), configuration=configuration + ) + + def test_candidates_preserve_id_order_and_reject_unknown_ids(self): + presets = [ + get_endpoint_preset(preset_id="aa11bb22"), + get_endpoint_preset(preset_id="cc33dd44"), ] - assert not _get_matching_presets(presets, configuration=configuration, preset_id="small") + + ordered = _get_candidate_presets(presets, preset_ids=["cc33dd44", "aa11bb22"]) + assert [preset.id for preset in ordered] == ["cc33dd44", "aa11bb22"] + assert _get_candidate_presets(presets, preset_ids=None) == presets + with pytest.raises(CLIError, match="does not exist"): + _get_candidate_presets(presets, preset_ids=["aa11bb22", "ee55ff66"]) def test_exact_request_matches_repo_and_client_facing_name(self): matching = get_endpoint_preset(preset_id="matching") @@ -46,13 +60,10 @@ def test_exact_request_matches_repo_and_client_facing_name(self): }, ) - assert _get_matching_presets([matching], configuration=configuration, preset_id=None) == [ - matching - ] + assert _get_matching_presets([matching], configuration=configuration) == [matching] assert not _get_matching_presets( [matching.copy(update={"model": "other/repo"})], configuration=configuration, - preset_id=None, ) @@ -124,8 +135,8 @@ def test_applies_the_selected_plan(self, monkeypatch): name="qwen", model={"base": "Qwen/Qwen3.5-27B"}, ), - configuration_path="endpoint.dstack.yml", - preset_id=None, + configuration_path="preset.dstack.yml", + preset_ids=None, profile_name="gpu", command_args=command_args, store=Mock(list=Mock(return_value=[preset])), @@ -139,8 +150,7 @@ def test_applies_the_selected_plan(self, monkeypatch): configurator_args=service_args, plan_properties={ "Model": "Qwen/Qwen3.5-27B ([secondary]base[/])", - "Preset": "8f3a12c4 " - "([secondary]context=32K, concurrency=1 42.1 tok/s TTFT 108ms[/])", + "Preset": "8f3a12c4 ([secondary]ctx=32K con=1 42.1 tok/s TTFT 108ms[/])", }, ) diff --git a/src/tests/_internal/cli/services/endpoints/test_create.py b/src/tests/_internal/cli/services/endpoints/test_create.py index 99edded5d..d8a64ecfc 100644 --- a/src/tests/_internal/cli/services/endpoints/test_create.py +++ b/src/tests/_internal/cli/services/endpoints/test_create.py @@ -11,14 +11,20 @@ EndpointAgentProcessOutput, EndpointAgentSession, EndpointAgentWorkspace, + create_agent_workspace, print_endpoint_progress, + remove_agent_workspace, ) from dstack._internal.cli.services.endpoints.create import ( EndpointPresetCreateResult, - _build_prompt, + _build_constraints, _cleanup_runs, _create_endpoint_preset, _get_build_name, + _print_fleet_offers, + _save_final_report_copy, + _stop_active_session_runs, + attach_endpoint_preset, create_endpoint_preset, ) from dstack._internal.cli.services.endpoints.store import EndpointPresetStore @@ -76,7 +82,7 @@ def creation_context(tmp_path, monkeypatch): ) monkeypatch.setattr( "dstack._internal.cli.services.endpoints.create._get_build_name", - lambda _: "qwen-build", + lambda *_: "qwen-build", ) return SimpleNamespace( api=api, @@ -117,10 +123,20 @@ async def create(**kwargs): store=EndpointPresetStore(tmp_path / "presets"), ) - paths = list((tmp_path / ".dstack" / "agent" / "qwen").iterdir()) + paths = [ + path + for path in (tmp_path / ".dstack" / "presets").iterdir() + if path.is_dir() and not path.name.startswith(".") + ] assert len(paths) == 1 - assert paths[0].name.endswith(f"-{preset.id}") - assert {path.name for path in paths[0].iterdir()} == {"agent.log"} + assert {path.name for path in paths[0].iterdir()} == { + "agent.log", + "session.json", + "preset.dstack.yml", + } + manifest = json.loads((paths[0] / "session.json").read_text()) + assert manifest["status"] == "success" + assert manifest["id"] == paths[0].name assert "testing preset" in (paths[0] / "agent.log").read_text() output = capsys.readouterr().out.replace("\n", "") assert f"Agent log saved to {paths[0] / 'agent.log'}" in output @@ -166,17 +182,22 @@ def fail_finish(self, preset_id=None): debug=True, ) - paths = list((tmp_path / ".dstack" / "agent" / "qwen").iterdir()) + paths = [ + path + for path in (tmp_path / ".dstack" / "presets").iterdir() + if path.is_dir() and not path.name.startswith(".") + ] assert len(paths) == 1 - assert paths[0].name.endswith("-running") assert result.preset == preset assert {path.name for path in paths[0].iterdir()} == { - "endpoint.dstack.yml", + "preset.dstack.yml", "agent.log", "prompt.md", + "session.json", "trace.jsonl", } - assert "hf-secret" not in (paths[0] / "endpoint.dstack.yml").read_text() + assert json.loads((paths[0] / "session.json").read_text())["status"] == "running" + assert "hf-secret" not in (paths[0] / "preset.dstack.yml").read_text() assert "Files remain at" in capsys.readouterr().out def test_debug_finalization_does_not_mask_creation_error(self, tmp_path, monkeypatch): @@ -229,7 +250,7 @@ async def test_checks_active_fleets_before_claude_auth(self, tmp_path, monkeypat ) @pytest.mark.asyncio - async def test_cleans_up_runs_when_cancelled(self, creation_context, monkeypatch, tmp_path): + async def test_skips_cleanup_when_cancelled(self, creation_context, monkeypatch, tmp_path): async def run_agent(**_): raise asyncio.CancelledError @@ -256,8 +277,7 @@ async def cleanup_runs(**kwargs): agent_session=_agent_session(tmp_path), ) - assert len(cleanup_calls) == 1 - assert cleanup_calls[0]["build_name"] == "qwen-build" + assert cleanup_calls == [] @pytest.mark.parametrize( ("keep_service", "stopped_names"), @@ -294,6 +314,7 @@ async def run_agent(**kwargs): source_configuration=creation_context.source_configuration, store=creation_context.store, keep_service=keep_service, + build_name="qwen-build", agent_session=agent_session, ) @@ -306,54 +327,21 @@ async def run_agent(**kwargs): class TestBuildName: - def test_requires_name_and_keeps_generated_prefix_bounded(self, monkeypatch): - with pytest.raises(CLIError, match="Endpoint name is required"): - _get_build_name(None) - monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.secrets.token_hex", - lambda _: "a1b2c3", - ) - - build_name = _get_build_name("qwen-endpoint-with-a-name-that-is-forty-one") + def test_derives_slug_for_nameless_and_keeps_prefix_bounded(self): + assert _get_build_name(None, "Qwen/Qwen3.5-27B", "a1b2c3d4") == "qwen3-5-27b-a1b2c3d4" - assert len(f"{build_name}-99999") <= 41 - - -class TestBuildPrompt: - def test_distinguishes_base_from_exact_model(self): - base_prompt = _build_prompt( - configuration=EndpointConfiguration( - name="qwen", - model={"base": "Qwen/Qwen3.5-27B"}, - context_length=8192, - ), - build_name="qwen-build", - allowed_fleets=("gpu-fleet",), - ) - exact_prompt = _build_prompt( - configuration=EndpointConfiguration( - name="qwen", - model={ - "repo": "community/Qwen3.5-27B-GPTQ-Int4", - "name": "Qwen/Qwen3.5-27B", - }, - ), - build_name="qwen-build", - allowed_fleets=("gpu-fleet",), + build_name = _get_build_name( + "qwen-endpoint-with-a-name-that-is-forty-one", "Qwen/Qwen3.5-27B", "a1b2c3d4" ) - assert "- base_model: Qwen/Qwen3.5-27B" in base_prompt - assert "- model_repo:" not in base_prompt - assert "- context_length: 8192" in base_prompt - assert "- model_repo: community/Qwen3.5-27B-GPTQ-Int4" in exact_prompt - assert "- base_model:" not in exact_prompt - assert "- context_length:" not in exact_prompt + assert build_name.endswith("-a1b2c3d4") + assert len(f"{build_name}-99999") <= 41 class TestCleanupRuns: @pytest.mark.asyncio async def test_stops_only_recorded_build_runs(self, tmp_path, monkeypatch): - (tmp_path / "submissions.jsonl").write_text( + (tmp_path / "runs.jsonl").write_text( '{"name":"qwen-build-1"}\n{"name":"unrelated-run"}\n{"name":"qwen-build-2"}\n' ) runs = _FakeRuns() @@ -392,6 +380,7 @@ def _agent_session(tmp_path, *, debug: bool = False) -> EndpointAgentSession: path=path, timestamp="20260714-120000Z", debug=debug, + preset_id="ab12cd34", ) @@ -423,3 +412,384 @@ def stop(self, project, names, abort): assert abort is False self.stopped_names.extend(names) self.run.status = RunStatus.TERMINATED + + +class TestBuildConstraints: + def test_renders_all_fields_with_explicit_nulls_and_defaults(self): + configuration = EndpointConfiguration( + name="qwen", + model={"base": "Qwen/Qwen3-32B"}, + env=["HF_TOKEN"], + ) + + text = _build_constraints( + configuration=configuration, + build_name="qwen-abc123", + allowed_fleets=("gpu-fleet",), + ) + + assert text.endswith("\n") + assert json.loads(text) == { + "run_name_prefix": "qwen-abc123", + "model": {"base": "Qwen/Qwen3-32B"}, + "context_length": None, + "max_trials": 3, + "concurrency": 8, + "fleets": ["gpu-fleet"], + "env": ["HF_TOKEN"], + } + + def test_renders_configured_values(self): + configuration = EndpointConfiguration( + name="qwen", + model={"repo": "Qwen/Qwen3-32B-AWQ", "name": "qwen3"}, + context_length=32768, + max_trials=10, + concurrency=16, + ) + + data = json.loads( + _build_constraints( + configuration=configuration, + build_name="qwen-abc123", + allowed_fleets=("a", "b"), + ) + ) + + assert data["model"] == {"repo": "Qwen/Qwen3-32B-AWQ", "name": "qwen3"} + assert data["context_length"] == 32768 + assert data["max_trials"] == 10 + assert data["concurrency"] == 16 + assert data["fleets"] == ["a", "b"] + + +class TestSaveFinalReportCopy: + def test_copies_report_redacted(self, tmp_path): + workspace = EndpointAgentWorkspace(path=tmp_path / "w", dstack_home=tmp_path / "h") + workspace.path.mkdir() + workspace.final_report_path.write_text( + '{"success": true, "note": "token dstack-secret"}', encoding="utf-8" + ) + session = _agent_session(tmp_path, debug=True) + + _save_final_report_copy( + workspace=workspace, + agent_session=session, + redacted_values=["dstack-secret"], + ) + + copied = (session.path / "final_report.json").read_text() + assert "dstack-secret" not in copied + assert "[redacted]" in copied + + def test_missing_report_is_no_op(self, tmp_path): + workspace = EndpointAgentWorkspace(path=tmp_path / "w", dstack_home=tmp_path / "h") + workspace.path.mkdir() + session = _agent_session(tmp_path, debug=True) + + _save_final_report_copy( + workspace=workspace, + agent_session=session, + redacted_values=["dstack-secret"], + ) + + assert not (session.path / "final_report.json").exists() + + +class TestInterruptAndResume: + def test_interrupt_suspends_session(self, tmp_path, monkeypatch, capsys): + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("USERPROFILE", str(tmp_path)) + + async def create(**kwargs): + raise KeyboardInterrupt + + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.create._create_endpoint_preset", + create, + ) + + with pytest.raises(KeyboardInterrupt): + create_endpoint_preset( + api=SimpleNamespace(), + configuration=EndpointConfiguration( + name="qwen", model={"base": "Qwen/Qwen3.5-27B"} + ), + store=EndpointPresetStore(tmp_path / "presets"), + ) + + sessions = [ + path + for path in (tmp_path / ".dstack" / "presets").iterdir() + if path.is_dir() and not path.name.startswith(".") + ] + assert len(sessions) == 1 + manifest = json.loads((sessions[0] / "session.json").read_text()) + assert manifest["status"] == "interrupted" + output = capsys.readouterr().out + assert "Resume with" in output + assert sessions[0].name in output + + @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 = EndpointAgentSession( + path=session_dir, timestamp="t", debug=False, preset_id="fe98dc76" + ) + workspace = create_agent_workspace(agent_session) + workspace.constraints_path.write_text( + '{"run_name_prefix": "qwen-build"}', encoding="utf-8" + ) + agent_session.update_manifest(claude_session_id="sid-xyz", claude_model="claude-pinned") + captured = {} + + async def run_agent(**kwargs): + captured.update(kwargs) + return EndpointAgentProcessOutput( + report_data=json.loads(get_successful_endpoint_report(creation_context.run).json()) + ) + + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.create.run_endpoint_agent", + run_agent, + ) + + result = await _create_endpoint_preset( + api=creation_context.api, + configuration=creation_context.configuration, + source_configuration=creation_context.source_configuration, + store=creation_context.store, + agent_session=agent_session, + resume=True, + ) + + assert captured["initial_resume_session_id"] == "sid-xyz" + assert captured["auth"].model == "claude-pinned" + assert result.preset.id == "fe98dc76" + assert (session_dir / "workspace").is_dir() + remove_agent_workspace(agent_session) + + @pytest.mark.asyncio + async def test_pins_user_prompt_on_create(self, creation_context, monkeypatch, tmp_path): + session_dir = tmp_path / "ab34ef12" + session_dir.mkdir() + (session_dir / "agent.log").touch() + agent_session = EndpointAgentSession( + path=session_dir, timestamp="t", debug=False, preset_id="ab34ef12" + ) + captured = {} + + async def run_agent(**kwargs): + captured.update(kwargs) + return EndpointAgentProcessOutput( + report_data=json.loads(get_successful_endpoint_report(creation_context.run).json()) + ) + + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.create.run_endpoint_agent", + run_agent, + ) + + await _create_endpoint_preset( + api=creation_context.api, + configuration=creation_context.configuration, + source_configuration=creation_context.source_configuration, + store=creation_context.store, + build_name="qwen-build", + agent_session=agent_session, + user_prompt="Optimize for RAG traffic.", + ) + + assert agent_session.read_user_prompt() == "Optimize for RAG traffic." + assert "## Additional instructions" in captured["prompt"] + assert "Optimize for RAG traffic." in captured["prompt"] + + @pytest.mark.asyncio + async def test_resume_keeps_the_pinned_user_prompt( + self, creation_context, monkeypatch, tmp_path, capsys + ): + session_dir = tmp_path / "ab34ef12" + session_dir.mkdir() + (session_dir / "agent.log").touch() + agent_session = EndpointAgentSession( + path=session_dir, timestamp="t", debug=False, preset_id="ab34ef12" + ) + workspace = create_agent_workspace(agent_session) + workspace.constraints_path.write_text( + '{"run_name_prefix": "qwen-build"}', encoding="utf-8" + ) + agent_session.update_manifest(claude_session_id="sid-abc") + agent_session.write_user_prompt("Optimize for RAG traffic.") + captured = {} + + async def run_agent(**kwargs): + captured.update(kwargs) + return EndpointAgentProcessOutput( + report_data=json.loads(get_successful_endpoint_report(creation_context.run).json()) + ) + + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.create.run_endpoint_agent", + run_agent, + ) + + await _create_endpoint_preset( + api=creation_context.api, + configuration=creation_context.configuration, + source_configuration=creation_context.source_configuration, + store=creation_context.store, + agent_session=agent_session, + resume=True, + user_prompt="A different prompt.", + ) + + # The session keeps its original prompt; the new one is ignored with a warning. + assert "Optimize for RAG traffic." in captured["prompt"] + assert "A different prompt." not in captured["prompt"] + assert "keepsitsoriginalprompt" in "".join(capsys.readouterr().out.split()) + remove_agent_workspace(agent_session) + + +class TestFleetOffersPreview: + def test_no_offers_shows_the_shared_warning_without_failing(self, capsys): + plan = SimpleNamespace( + project_name="main", + user="admin", + job_plans=[SimpleNamespace(offers=[], total_offers=0, max_price=None)], + ) + api = SimpleNamespace( + project="main", + client=SimpleNamespace(runs=SimpleNamespace(get_plan=lambda *a, **k: plan)), + ) + + _print_fleet_offers(api, ("arm-fleet",)) + + out = capsys.readouterr().out + assert "arm-fleet" in out + assert "No matching instance offers available" in out + + +class TestAttachEndpointPreset: + def _detached_session(self, tmp_path, configuration_yaml: str) -> EndpointAgentSession: + session_dir = tmp_path / "ab12cd34" + session_dir.mkdir() + (session_dir / "agent.log").touch() + (session_dir / "preset.dstack.yml").write_text(configuration_yaml) + session = EndpointAgentSession( + path=session_dir, timestamp="t", debug=False, preset_id="ab12cd34" + ) + workspace = create_agent_workspace(session) + workspace.constraints_path.write_text('{"run_name_prefix": "qwen-build"}') + session.update_manifest(status="running", agent_pid=987654321) + return session + + def test_finalizes_a_detached_session(self, creation_context, monkeypatch, tmp_path): + session = self._detached_session( + tmp_path, "type: preset\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\n" + ) + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.agent.get_presets_dir", + lambda: tmp_path, + ) + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.create.load_attachable_agent_session", + lambda preset_id: session, + ) + + async def fake_attach(**kwargs): + return EndpointAgentProcessOutput( + report_data=json.loads(get_successful_endpoint_report(creation_context.run).json()) + ) + + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.create.attach_endpoint_agent", + fake_attach, + ) + + result = attach_endpoint_preset( + api=creation_context.api, + store=creation_context.store, + preset_id="ab12cd34", + ) + + assert result.preset.id == "ab12cd34" + assert creation_context.store.get("ab12cd34") is not None + assert session.read_manifest()["status"] == "success" + + def test_agent_death_without_report_suspends_instead_of_failing( + self, creation_context, monkeypatch, tmp_path + ): + session = self._detached_session( + tmp_path, "type: preset\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\n" + ) + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.create.load_attachable_agent_session", + lambda preset_id: session, + ) + + async def fake_attach(**kwargs): + return EndpointAgentProcessOutput(error="agent died") + + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.create.attach_endpoint_agent", + fake_attach, + ) + + with pytest.raises(CLIError, match="agent died"): + attach_endpoint_preset( + api=creation_context.api, + store=creation_context.store, + preset_id="ab12cd34", + ) + + assert session.read_manifest()["status"] == "interrupted" + + +class TestStopActiveSessionRuns: + def _session(self, tmp_path) -> EndpointAgentSession: + 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 EndpointAgentSession( + path=session_dir, timestamp="t", debug=False, preset_id="ab12cd34" + ) + + def _api(self, statuses: dict, stopped: list) -> SimpleNamespace: + def get(project, name): + return SimpleNamespace(status=statuses[name]) + + def stop(project, names, abort): + stopped.extend(names) + + return SimpleNamespace( + project="main", + client=SimpleNamespace(runs=SimpleNamespace(get=get, stop=stop)), + ) + + def test_stops_only_non_terminal_runs(self, tmp_path): + from dstack._internal.core.models.runs import RunStatus + + stopped: list = [] + api = self._api( + {"qwen-build-1": RunStatus.DONE, "qwen-build-2": RunStatus.RUNNING}, stopped + ) + + assert _stop_active_session_runs(api, self._session(tmp_path), assume_yes=True) is True + assert stopped == ["qwen-build-2"] + + def test_declined_confirmation_leaves_runs_active(self, tmp_path, monkeypatch): + from dstack._internal.cli.services.endpoints import create as create_module + from dstack._internal.core.models.runs import RunStatus + + monkeypatch.setattr(create_module, "confirm_ask", lambda *_: False) + stopped: list = [] + api = self._api( + {"qwen-build-1": RunStatus.RUNNING, "qwen-build-2": RunStatus.DONE}, stopped + ) + + assert _stop_active_session_runs(api, self._session(tmp_path), assume_yes=False) is False + assert stopped == [] diff --git a/src/tests/_internal/cli/services/endpoints/test_output.py b/src/tests/_internal/cli/services/endpoints/test_output.py new file mode 100644 index 000000000..6bb188736 --- /dev/null +++ b/src/tests/_internal/cli/services/endpoints/test_output.py @@ -0,0 +1,149 @@ +from datetime import timedelta +from io import StringIO + +import pytest +from rich.console import Console +from rich.table import Table +from rich.theme import Theme + +from dstack._internal.cli.services.endpoints import output as output_module +from dstack._internal.cli.services.endpoints.output import _add_session, _format_number +from tests._internal.cli.endpoint_presets import get_endpoint_preset + +pytestmark = pytest.mark.windows + + +class TestFormatNumber: + def test_large_values_avoid_scientific_notation(self): + assert _format_number(1723.4) == "1723" + assert _format_number(999.6) == "1000" + + def test_small_values_keep_three_significant_digits(self): + assert _format_number(384.8) == "385" + assert _format_number(15.92) == "15.9" + assert _format_number(0.99) == "0.99" + + +def _session_row(session: dict) -> dict: + table = Table(box=None) + for column in ("BASE", "ID", "GPU", "BENCHMARK", "STATUS", "SUBMITTED"): + table.add_column(column) + _add_session(table, session) + return { + column.header: "".join(str(cell) for cell in column._cells) for column in table.columns + } + + +class TestSessionRow: + def test_shows_progress_after_status_and_best_benchmark(self): + row = _session_row( + { + "id": "c7e18d52", + "status": "running", + "max_trials": 3, + "trials": { + "count": 2, + "best": {"tok_s": 2339.0, "concurrency": 8, "gpu": "A40:48GB:1"}, + }, + } + ) + + assert row["STATUS"] == "[bold deep_sky_blue1]clauding[/] [secondary](2/3)[/]" + assert row["BENCHMARK"] == "best trial: con=8 2339 tok/s" + assert row["GPU"] == "A40:48GB:1" + + def test_shows_zero_progress_without_benchmark(self): + row = _session_row( + {"id": "ab12cd34", "status": "running", "max_trials": 3, "trials": {"count": 0}} + ) + + assert row["STATUS"] == "[bold deep_sky_blue1]clauding[/] [secondary](0/3)[/]" + assert row["BENCHMARK"] == "" + + def test_omits_progress_without_trials_data(self): + row = _session_row({"id": "ab12cd34", "status": "interrupted"}) + + assert row["STATUS"] == "[bold gold1]interrupted[/]" + + def test_counts_without_max_trials(self): + row = _session_row({"id": "ab12cd34", "status": "interrupted", "trials": {"count": 2}}) + + assert row["STATUS"] == "[bold gold1]interrupted[/] [secondary](2)[/]" + + +class TestGroupOrdering: + def test_sorts_presets_and_sessions_newest_first(self, monkeypatch): + buffer = StringIO() + monkeypatch.setattr( + output_module, + "console", + Console( + file=buffer, width=200, color_system=None, theme=Theme({"secondary": "grey58"}) + ), + ) + old = get_endpoint_preset() + new = old.copy(update={"id": "11aa22bb", "created_at": old.created_at + timedelta(days=2)}) + sessions = [ + { + "id": "aaaaaaaa", + "status": "interrupted", + "model": old.base, + "created_at": "2026-07-01T00:00:00+00:00", + }, + { + "id": "bbbbbbbb", + "status": "interrupted", + "model": old.base, + "created_at": "2026-07-02T00:00:00+00:00", + }, + ] + + output_module.print_endpoint_presets([old, new], sessions=sessions) + + text = buffer.getvalue() + assert text.index("11aa22bb") < text.index(old.id) + assert text.index(old.id) < text.index("bbbbbbbb") < text.index("aaaaaaaa") + + +class TestDoneProgress: + def test_completed_creation_decorates_preset_row_without_extra_session_row(self, monkeypatch): + buffer = StringIO() + monkeypatch.setattr( + output_module, + "console", + Console( + file=buffer, width=200, color_system=None, theme=Theme({"secondary": "grey58"}) + ), + ) + preset = get_endpoint_preset() + sessions = [ + { + "id": preset.id, + "status": "success", + "model": preset.base, + "max_trials": 4, + "trials": {"count": 3}, + } + ] + + output_module.print_endpoint_presets([preset], sessions=sessions) + + text = buffer.getvalue() + assert "success (3/4)" in text + assert text.count(preset.id) == 1 + + +class TestVerifyingStatus: + def test_running_session_with_exhausted_trials_shows_verifying(self): + row = _session_row( + {"id": "ab12cd34", "status": "running", "max_trials": 2, "trials": {"count": 2}} + ) + + assert row["STATUS"] == "[bold deep_sky_blue1]verifying[/] [secondary](2/2)[/]" + + def test_running_session_with_remaining_trials_stays_clauding(self): + row = _session_row( + {"id": "ab12cd34", "status": "running", "max_trials": 2, "trials": {"count": 1}} + ) + + assert row["STATUS"].startswith("[bold deep_sky_blue1]clauding[/]") diff --git a/src/tests/_internal/cli/services/endpoints/test_prompt.py b/src/tests/_internal/cli/services/endpoints/test_prompt.py new file mode 100644 index 000000000..01c4a30c3 --- /dev/null +++ b/src/tests/_internal/cli/services/endpoints/test_prompt.py @@ -0,0 +1,46 @@ +import pytest + +from dstack._internal.cli.services.endpoints import prompt as prompt_module +from dstack._internal.cli.services.endpoints.prompt import get_endpoint_agent_system_prompt +from dstack._internal.core.errors import CLIError + +pytestmark = pytest.mark.windows + + +class TestSystemPrompt: + def test_stays_byte_identical_without_user_prompt(self): + text = get_endpoint_agent_system_prompt() + + assert ( + text == get_endpoint_agent_system_prompt(None) == get_endpoint_agent_system_prompt("") + ) + assert "## Additional instructions" not in text + assert " more.\n") + monkeypatch.setattr(prompt_module, "_SYSTEM_PROMPT_PATH", broken) + + with pytest.raises(CLIError, match="Unknown variable"): + get_endpoint_agent_system_prompt() diff --git a/src/tests/_internal/cli/services/endpoints/test_store.py b/src/tests/_internal/cli/services/endpoints/test_store.py index 6088ec8cd..f205c6964 100644 --- a/src/tests/_internal/cli/services/endpoints/test_store.py +++ b/src/tests/_internal/cli/services/endpoints/test_store.py @@ -1,4 +1,6 @@ +from io import StringIO from pathlib import Path +from unittest.mock import patch import pytest import yaml @@ -11,8 +13,10 @@ EndpointBenchmarkMetrics, EndpointBenchmarkWorkload, ) +from dstack._internal.cli.models.endpoints import EndpointConfiguration +from dstack._internal.cli.services.endpoints import store as store_module from dstack._internal.cli.services.endpoints.store import EndpointPresetStore -from dstack._internal.core.errors import CLIError +from dstack._internal.core.errors import CLIError, ConfigurationError from dstack._internal.core.models.envs import EnvSentinel from tests._internal.cli.endpoint_presets import ( get_endpoint_benchmark, @@ -68,7 +72,7 @@ def test_saves_and_lists_self_contained_preset(self, tmp_path: Path): path = store.save(preset) - assert path == (tmp_path / "presets" / "models--Qwen--Qwen3.5-27B" / "8f3a12c4.yaml") + assert path == (tmp_path / "presets" / "8f3a12c4" / "preset.yaml") data = yaml.safe_load(path.read_text()) assert data["base"] == preset.base assert data["id"] == preset.id @@ -89,14 +93,37 @@ def test_replaces_same_preset_id_atomically(self, tmp_path: Path): assert store.get(updated.id) == updated - def test_rejects_duplicate_preset_id(self, tmp_path: Path): + def test_same_id_save_overwrites(self, tmp_path: Path): store = EndpointPresetStore(tmp_path / "presets") preset = get_endpoint_preset() store.save(preset) store.save(preset.copy(update={"base": "Qwen/Another-Model"})) - with pytest.raises(CLIError, match="is not unique"): - store.get(preset.id) + loaded = store.get(preset.id) + assert loaded is not None + assert loaded.base == "Qwen/Another-Model" + + def test_migrates_legacy_layout_and_archives_on_delete(self, tmp_path: Path): + root = tmp_path / "presets" + store = EndpointPresetStore(root) + preset = get_endpoint_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(EndpointPresetStore(tmp_path / "tmp").save(preset).read_text()), + sort_keys=False, + ) + ) + + assert store.list() == [preset] + assert (root / preset.id / "preset.yaml").is_file() + assert not legacy.exists() + + assert store.delete(preset.id) is True + assert store.get(preset.id) is None + assert (root / ".archive" / preset.id / "preset.yaml").is_file() + assert store.list() == [] def test_rejects_preset_without_successful_benchmark(self, tmp_path: Path): store = EndpointPresetStore(tmp_path / "presets") @@ -125,3 +152,86 @@ def test_preserves_literal_env_values(self, tmp_path: Path): assert "TOKENIZERS_PARALLELISM=false" in env assert "MODEL_LABEL=monkey" in env assert "HF_TOKEN" in env + + +class TestParseEndpointConfiguration: + @pytest.mark.parametrize("key", ["base", "repo"]) + def test_warns_on_nested_model_without_name(self, key: str): + stream = StringIO(f"type: preset\nmodel:\n {key}: Qwen/Qwen3.5-27B\n") + + with patch.object(store_module, "warn") as warn: + configuration = store_module._parse_endpoint_configuration(stream) + + warn.assert_called_once() + assert f"model.{key}" in warn.call_args.args[0] + assert f"`{key}:`" in warn.call_args.args[0] + assert configuration.model is not None + + @pytest.mark.parametrize( + "body", + [ + "base: Qwen/Qwen3.5-27B\n", + "repo: Qwen/Qwen3.5-27B\n", + "model: Qwen/Qwen3.5-27B\n", + "model:\n repo: community/Qwen3.5-27B-GPTQ-Int4\n name: Qwen/Qwen3.5-27B\n", + ], + ) + def test_does_not_warn_on_preferred_syntax(self, body: str): + stream = StringIO(f"type: preset\n{body}") + + with patch.object(store_module, "warn") as warn: + configuration = store_module._parse_endpoint_configuration(stream) + + warn.assert_not_called() + assert configuration.model is not None + + +class TestResolveEndpointPrompt: + def test_resolves_inline_and_file_relative_to_configuration(self, tmp_path: Path): + (tmp_path / "notes.md").write_text("From a file.\n") + configuration_path = str(tmp_path / "preset.dstack.yml") + inline = EndpointConfiguration(name="q", base="Q/M", prompt="Inline text.") + from_file = EndpointConfiguration(name="q", base="Q/M", prompt={"path": "notes.md"}) + + assert store_module.resolve_endpoint_prompt(inline, configuration_path) == "Inline text." + assert ( + store_module.resolve_endpoint_prompt(from_file, configuration_path) == "From a file." + ) + assert ( + store_module.resolve_endpoint_prompt( + EndpointConfiguration(name="q", base="Q/M"), configuration_path + ) + is None + ) + + def test_rejects_missing_and_empty_prompt_files(self, tmp_path: Path): + configuration_path = str(tmp_path / "preset.dstack.yml") + (tmp_path / "empty.md").write_text(" \n") + + with pytest.raises(ConfigurationError, match="Failed to read"): + store_module.resolve_endpoint_prompt( + EndpointConfiguration(name="q", base="Q/M", prompt={"path": "missing.md"}), + configuration_path, + ) + with pytest.raises(ConfigurationError, match="is empty"): + store_module.resolve_endpoint_prompt( + EndpointConfiguration(name="q", base="Q/M", prompt={"path": "empty.md"}), + configuration_path, + ) + + +class TestPresetNames: + def test_finds_and_detaches_names(self, tmp_path: Path): + store = EndpointPresetStore(tmp_path / "presets") + named = get_endpoint_preset().copy(update={"name": "qwen"}) + store.save(named) + store.save(get_endpoint_preset().copy(update={"id": "01234567"})) + + assert store.find_by_name("qwen").id == named.id + assert store.find_by_name("other") is None + + released = store.release_name("qwen") + + assert released.id == named.id + assert store.find_by_name("qwen") is None + assert store.get(named.id).name is None diff --git a/src/tests/_internal/cli/services/endpoints/test_trials.py b/src/tests/_internal/cli/services/endpoints/test_trials.py new file mode 100644 index 000000000..030c83943 --- /dev/null +++ b/src/tests/_internal/cli/services/endpoints/test_trials.py @@ -0,0 +1,64 @@ +import pytest + +from dstack._internal.cli.models.endpoint_trials import EndpointPresetTrial + +pytestmark = pytest.mark.windows + + +def get_trial_record(benchmark=True) -> dict: + record = { + "task": { + "type": "task", + "name": "qwen-endpoint-1", + "image": "vllm/vllm-openai:v0.11.0", + "commands": [ + "vllm serve Qwen/Qwen3-32B-AWQ --reasoning-parser qwen3 --port 8000", + ], + }, + "resources": { + "cpu": "16", + "memory": "93GB", + "disk": "100GB", + "gpu": {"name": "RTX5090", "memory": "32GB", "count": 1}, + }, + "benchmark": None, + } + if benchmark: + record["benchmark"] = { + "tool": "vllm bench serve", + "tool_version": "0.11.0", + "command": "vllm bench serve --model Qwen/Qwen3-32B-AWQ --num-prompts 64", + "workload": { + "api": "chat_completions", + "num_requests": 64, + "input_tokens": 1024, + "output_tokens": 512, + "concurrency": 8, + }, + "metrics": { + "successful_requests": 64, + "failed_requests": 0, + "duration_seconds": 78.1, + "total_input_tokens": 65536, + "total_output_tokens": 32768, + "ttft_ms": {"mean": 1660.0, "p50": 1571.0, "p99": 2010.0}, + "tpot_ms": {"mean": 17.57, "p50": 17.5, "p99": 18.1}, + }, + } + return record + + +class TestEndpointPresetTrial: + def test_parses_measured_trial(self): + trial = EndpointPresetTrial.parse_obj(get_trial_record()) + + assert trial.task.name == "qwen-endpoint-1" + assert trial.task.commands[0].startswith("vllm serve") + assert trial.benchmark is not None + assert trial.benchmark.workload.concurrency == 8 + assert trial.benchmark.metrics.tpot_ms.mean == 17.57 + + def test_failed_trial_allows_null_benchmark(self): + trial = EndpointPresetTrial.parse_obj(get_trial_record(benchmark=False)) + + assert trial.benchmark is None diff --git a/src/tests/_internal/cli/services/endpoints/test_verify.py b/src/tests/_internal/cli/services/endpoints/test_verify.py index 29452b4b8..ee598694b 100644 --- a/src/tests/_internal/cli/services/endpoints/test_verify.py +++ b/src/tests/_internal/cli/services/endpoints/test_verify.py @@ -6,8 +6,13 @@ from dstack._internal.cli.models.endpoint_agent import AgentFinalReport from dstack._internal.cli.models.endpoints import EndpointConfiguration +from dstack._internal.cli.services.endpoints.agent import ( + EndpointAgentProcessOutput, + EndpointAgentWorkspace, +) from dstack._internal.cli.services.endpoints.verify import ( build_verified_endpoint_preset, + load_endpoint_agent_report, ) from dstack._internal.core.errors import CLIError from dstack._internal.core.models.envs import EnvSentinel @@ -80,3 +85,53 @@ def test_rejects_variant_for_exact_model_request(self): ), report=report, ) + + +class TestLoadEndpointAgentReport: + def _load(self, tmp_path, report_data, redacted_values): + return load_endpoint_agent_report( + output=EndpointAgentProcessOutput(report_data=report_data), + workspace=EndpointAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home"), + redacted_values=redacted_values, + ) + + def test_redacts_known_secret_in_benchmark_command_instead_of_failing(self, tmp_path): + run = get_running_service_run() + data = get_successful_endpoint_report(run).dict() + data["run_id"] = str(data["run_id"]) + data["benchmark"]["command"] = ( + "python bench.py --header 'Authorization: Bearer sk-live-0123456789abcdef'" + ) + + report = self._load(tmp_path, data, redacted_values=("sk-live-0123456789abcdef",)) + + assert report.benchmark is not None + assert report.benchmark.command.endswith("Bearer [redacted]'") + assert "sk-live" not in report.benchmark.command + + def test_still_rejects_unknown_bearer_token(self, tmp_path): + run = get_running_service_run() + data = get_successful_endpoint_report(run).dict() + data["run_id"] = str(data["run_id"]) + data["benchmark"]["command"] = ( + "curl -H 'Authorization: Bearer sk-unknown-9876543210fedcba'" + ) + + with pytest.raises(CLIError, match="bearer token"): + self._load(tmp_path, data, redacted_values=("some-other-secret-value",)) + + def test_allows_bearer_prose_without_credential(self, tmp_path): + # Regression: "(auth via DSTACK_TOKEN bearer header from env)" failed + # two live sessions — the word after "bearer" is prose, not a token. + run = get_running_service_run() + data = get_successful_endpoint_report(run).dict() + data["run_id"] = str(data["run_id"]) + data["benchmark"]["command"] = ( + "./benchenv/bin/python bench_service.py --base $DSTACK_SERVER_URL/x" + " (auth via DSTACK_TOKEN bearer header from env)" + ) + + report = self._load(tmp_path, data, redacted_values=()) + + assert report.benchmark is not None + assert "bearer header" in report.benchmark.command