From 1c96e107f830c20e859d9f36d1d0a48cb3253215 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Mon, 20 Jul 2026 12:39:58 -0700 Subject: [PATCH 01/25] [Feature] Endpoint preset trials MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Sequential experimental trials for endpoint preset creation: the agent searches serving configurations across hardware, frameworks, and variants in `dstack` tasks, records reproducible trial records (trials.jsonl), promotes the best to a verified service, and saves the preset. Verified across nine live e2e sessions (Qwen2.5-0.5B, Qwen3-32B on H100, Qwen3-32B under $1/hr — best result 533 tok/s on a single RTX 5090 via NVFP4 + EAGLE3). Co-Authored-By: Claude Fable 5 --- skills/dstack-prototyping/SKILL.md | 20 +- src/dstack/_internal/cli/commands/endpoint.py | 8 + .../_internal/cli/models/endpoint_agent.py | 22 +- .../_internal/cli/models/endpoint_trials.py | 16 + src/dstack/_internal/cli/models/endpoints.py | 42 ++ .../_internal/cli/services/endpoints/agent.py | 179 +++++- .../cli/services/endpoints/create.py | 84 +-- .../cli/services/endpoints/output.py | 3 + .../cli/services/endpoints/prompt.py | 53 +- .../endpoints/resources/system_prompt.md | 536 ++++++++++-------- .../cli/services/endpoints/test_agent.py | 103 +++- .../cli/services/endpoints/test_create.py | 131 +++-- .../cli/services/endpoints/test_output.py | 16 + .../cli/services/endpoints/test_trials.py | 64 +++ 14 files changed, 851 insertions(+), 426 deletions(-) create mode 100644 src/dstack/_internal/cli/models/endpoint_trials.py create mode 100644 src/tests/_internal/cli/services/endpoints/test_output.py create mode 100644 src/tests/_internal/cli/services/endpoints/test_trials.py 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..046276a1e 100644 --- a/src/dstack/_internal/cli/commands/endpoint.py +++ b/src/dstack/_internal/cli/commands/endpoint.py @@ -80,6 +80,12 @@ 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", @@ -255,6 +261,8 @@ def _get_effective_configuration( args: argparse.Namespace, ) -> EndpointConfiguration: _apply_name(configuration, args.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_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..e902a5c40 100644 --- a/src/dstack/_internal/cli/models/endpoints.py +++ b/src/dstack/_internal/cli/models/endpoints.py @@ -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")] @@ -101,6 +104,24 @@ class EndpointConfiguration( preset: Annotated[ Optional[str], Field(description="The preset ID to use when applying the endpoint") ] = None + 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]], Field( @@ -115,6 +136,14 @@ 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 + @validator("model", pre=True) def parse_model(cls, value: Any) -> Any: if isinstance(value, str): @@ -128,6 +157,19 @@ def validate_preset(cls, value: Optional[str]) -> Optional[str]: 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..5f124fd9e 100644 --- a/src/dstack/_internal/cli/services/endpoints/agent.py +++ b/src/dstack/_internal/cli/services/endpoints/agent.py @@ -16,7 +16,10 @@ 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,8 +30,11 @@ _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" _PROGRESS_ENV = "DSTACK_ENDPOINT_PROGRESS_LOG" _REDACTION = "[redacted]" _CLAUDE_TOOLS = "Bash,Read,Write,Edit,WebFetch,WebSearch,StructuredOutput" @@ -91,8 +97,16 @@ 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: @@ -114,9 +128,39 @@ 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_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 @@ -129,21 +173,12 @@ def append_log(self, line: str) -> None: 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 + status = { + "status": "success" if preset_id is not None else "failed", + "preset_id": preset_id, + } + _write_private_text(self.path / _SESSION_FILENAME, json.dumps(status, indent=2) + "\n") + return self.path @dataclass(frozen=True) @@ -184,6 +219,10 @@ def create_endpoint_agent_session( parent.mkdir(mode=0o700, parents=True, exist_ok=True) path = _create_agent_session_directory(parent, timestamp) _write_private_text(path / "agent.log", "") + _write_private_text( + path / _SESSION_FILENAME, + json.dumps({"status": "running", "preset_id": None}, indent=2) + "\n", + ) if debug: data = json.loads(configuration.json(exclude_none=True)) if configuration.env: @@ -203,10 +242,41 @@ def create_endpoint_agent_session( return EndpointAgentSession(path=path, timestamp=timestamp, debug=debug) +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 _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" + name = timestamp if index == 0 else f"{timestamp}-{index}" path = parent / name try: path.mkdir(mode=0o700) @@ -261,8 +331,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) @@ -293,7 +361,21 @@ async def run_endpoint_agent( redacted_values=redacted_values, agent_session=agent_session, ) - progress_task = asyncio.create_task(progress_tailer.run()) + record_mirrors = [ + _RecordMirror( + source=workspace.runs_path, + target=agent_session.runs_path, + redacted_values=redacted_values, + ), + _RecordMirror( + source=workspace.trials_path, + target=agent_session.trials_path, + redacted_values=redacted_values, + ), + ] + tailer_tasks = [ + asyncio.create_task(tailer.run()) for tailer in [progress_tailer, *record_mirrors] + ] proc: Optional[asyncio.subprocess.Process] = None try: proc = await asyncio.create_subprocess_exec( @@ -341,10 +423,14 @@ async def run_endpoint_agent( await _terminate_process(proc) raise finally: - progress_task.cancel() - with suppress(asyncio.CancelledError): - await progress_task + 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() output = stdout_output if output.report_data is None: @@ -412,7 +498,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() @@ -731,6 +818,44 @@ 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]) -> None: + self._source = source + self._target = target + self._redacted_values = redacted_values + self._offset = 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 + 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/create.py b/src/dstack/_internal/cli/services/endpoints/create.py index 9c1ef179e..7b353ff80 100644 --- a/src/dstack/_internal/cli/services/endpoints/create.py +++ b/src/dstack/_internal/cli/services/endpoints/create.py @@ -10,7 +10,10 @@ 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, @@ -22,13 +25,11 @@ get_redacted_values, get_sensitive_inherited_env_values, print_endpoint_progress, + redact, run_endpoint_agent, ) from dstack._internal.cli.services.endpoints.presets import endpoint_preset_to_data -from dstack._internal.cli.services.endpoints.prompt import ( - format_endpoint_constraints, - get_endpoint_agent_system_prompt, -) +from dstack._internal.cli.services.endpoints.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, @@ -123,18 +124,17 @@ async def _create_endpoint_preset( workspace=workspace, token=token, ) - prompt = _build_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") + prompt = get_endpoint_agent_system_prompt() 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, - ) + agent_session.write_constraints(constraints_text) + agent_session.write_agent_info(auth) try: process_output = await run_endpoint_agent( prompt=prompt, @@ -164,6 +164,12 @@ async def _create_endpoint_preset( ) creation_succeeded = True finally: + if agent_session.debug: + _save_final_report_copy( + workspace=workspace, + agent_session=agent_session, + redacted_values=redacted_values, + ) keep_final_service = keep_service and creation_succeeded try: await _cleanup_runs( @@ -248,33 +254,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()} + 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" -Endpoint context: -- endpoint_name: {build_name} -{chr(10).join(context_lines)} -{ - format_endpoint_constraints( - configuration, - configuration.env.as_dict(), - allowed_fleets=allowed_fleets, - ) - } -""" +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 +299,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 +314,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 +326,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..d040978c3 100644 --- a/src/dstack/_internal/cli/services/endpoints/output.py +++ b/src/dstack/_internal/cli/services/endpoints/output.py @@ -117,6 +117,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}" diff --git a/src/dstack/_internal/cli/services/endpoints/prompt.py b/src/dstack/_internal/cli/services/endpoints/prompt.py index 89e66964d..4d3fbd628 100644 --- a/src/dstack/_internal/cli/services/endpoints/prompt.py +++ b/src/dstack/_internal/cli/services/endpoints/prompt.py @@ -1,58 +1,9 @@ -import enum -import json from pathlib import Path -from typing import Any, Sequence - -from dstack._internal.cli.models.endpoints import EndpointConfiguration -from dstack._internal.core.models.profiles import ProfileParams _SYSTEM_PROMPT_PATH = Path(__file__).resolve().parent / "resources" / "system_prompt.md" +# 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() -> str: return _SYSTEM_PROMPT_PATH.read_text(encoding="utf-8").strip() - - -def format_endpoint_constraints( - configuration: EndpointConfiguration, - endpoint_env: dict[str, str], - *, - allowed_fleets: Sequence[str], -) -> str: - lines = [ - "Fixed endpoint constraints:", - "- Do not submit any task or service that conflicts with these values.", - ] - for field in ProfileParams.__fields__: - if field == "fleets": - continue - value = getattr(configuration, field) - if value is not None: - lines.append(f"- {field}: {_format_constraint_value(value)}") - if allowed_fleets: - lines.append(f"- fleets: {', '.join(allowed_fleets)}") - if configuration.gateway is not None: - lines.append(f"- gateway: {_format_constraint_value(configuration.gateway)}") - lines.append("- endpoint_env_keys: " + (", ".join(endpoint_env) if endpoint_env else "none")) - return "\n".join(lines) - - -def _format_constraint_value(value: Any) -> str: - data = _constraint_value_to_data(value) - if isinstance(data, (bool, dict, list)): - return json.dumps(data, sort_keys=True) - return str(data) - - -def _constraint_value_to_data(value: Any) -> Any: - if isinstance(value, enum.Enum): - return value.value - if hasattr(value, "format"): - return value.format() - if hasattr(value, "json"): - return json.loads(value.json(exclude_none=True)) - if isinstance(value, dict): - return {key: _constraint_value_to_data(item) for key, item in value.items()} - if isinstance(value, (list, tuple)): - return [_constraint_value_to_data(item) for item in value] - return value diff --git a/src/dstack/_internal/cli/services/endpoints/resources/system_prompt.md b/src/dstack/_internal/cli/services/endpoints/resources/system_prompt.md index a6ff533eb..2f0f2c08e 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,329 @@ # 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. -# Requested Model - -The `Endpoint context:` block contains either `model_repo` or `base_model`. - -- If it contains `model_repo`, deploy that repo/path exactly. -- If it contains `base_model`, choose a repo/path compatible with `base_model` - that best fits performance and hardware within the endpoint constraints, - allowed fleets, backends, and offers. A variant can be the base repo itself, a - different precision or quantization, or another trusted repo compatible with - `base_model`. - -If `Endpoint context:` contains `context_length`, the selected repo/path and -final service must support at least that context length. - -Use the real `dstack` CLI and shell commands in this workspace. Load and follow -`/dstack` for dstack CLI/YAML syntax. Load and follow `/dstack-prototyping` for -how to test a model-serving configuration with tasks before verifying it as a -service. The skill files are installed under `.claude/skills` if you need to -inspect them directly. - -# Progress - -Write endpoint progress with: - -```bash -progress "message" -``` +# 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. - -After stopping a task or service, follow `/dstack` structured status guidance -and confirm that the run reached a terminal status before continuing. +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. -# Run Names +After stopping a `dstack` task or service, follow `/dstack` structured status +guidance and confirm that the run reached a terminal status before continuing. -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. - -# Backend/Fleet Selection, Idle Duration, Instance Volumes +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. -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. +# Hardware -## Backend And Fleet Choice +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. -Use `/dstack-prototyping` to learn how to select backends and fleets. +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. -Follow `/dstack-prototyping` skill on using a VM-based backend, Kubernetes backend, or SSH fleet that supports idle instances/instance volumes if there is such an option for the required GPU class. +If the backend allows, use instance volumes to mount cache and model weights +between runs. -If backend allows (see above), use instance volumes to mount cache and model weights between runs. +When submitting a `dstack` task or service, pass exact `fleets`, `backends`, +and an intentional `resources` range based on the choice made from offers, so +the run does not land outside the intended fleet/backend/hardware. -## Validating Offers +## Fleet Offers -When selecting backends/fleets and evaluating hardware that can be used to run -the model, use `dstack offer --json` and pass `--fleet` explicitly. If the -endpoint request has `fleets`, pass those fleets. Otherwise, pass every existing -fleet. If `--fleet` is not passed, `dstack offer` can show offers that are not -applicable to the fleets allowed by the endpoint request. Use these offers when -selecting fleet, backend, and hardware. +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. -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. +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`). -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. +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. -Verify the context length that the final service actually supports and report -it as `final_report.json.context_length`. +During the service verification, test the context length the service +actually supports and report it as `final_report.json.context_length`. -# Benchmark +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`). -Send benchmark requests to the final service run's absolute dstack service URL -using the bearer authentication used for service verification. Do not send them -to a server used during task prototyping or an SSH-forwarded local port. - -Run one benchmark with streaming responses. Choose a benchmark tool and workload -that can produce every required field below. - -Report the benchmark as `final_report.json.benchmark` using exactly the -following structure and field names (values are illustrative): - -```json -{ - "tool": "vllm bench serve", - "tool_version": "0.11.0", - "command": "vllm bench serve ...", - "workload": {"api": "chat_completions", "num_requests": 16, "input_tokens": 1024, "output_tokens": 128, "concurrency": 1}, - "metrics": { - "successful_requests": 16, "failed_requests": 0, "duration_seconds": 4.0, - "total_input_tokens": 16384, "total_output_tokens": 2048, - "ttft_ms": {"mean": 110.9, "p50": 108.2, "p99": 121.6}, - "tpot_ms": {"mean": 7.5, "p50": 7.4, "p99": 8.1} - } -} -``` - -Set `tool` to the command name and subcommands, without options or values (for -example, `vllm bench serve`). Set `tool_version` to the exact version and -`command` to the secret-free invocation. `api` must be `chat_completions` or -`completions`. `num_requests` is the number of measured requests; -`input_tokens` and `output_tokens` are the selected per-request lengths; and -`concurrency` is the maximum number of simultaneous requests. - -Calculate all metrics from the `num_requests` benchmark requests only. Exclude -setup, health-check, and warmup requests. `successful_requests` and -`failed_requests` are request counts; all requests must succeed. -`duration_seconds` is the elapsed wall-clock time from starting the first -measured request until the last measured request completes. `total_input_tokens` -and `total_output_tokens` are actual measured token totals. `ttft_ms` is the -time-to-first-token distribution across requests. `tpot_ms` is the -time-per-output-token distribution: for each request, divide the time from the -first output token to the last by one less than the actual output token count. -For both distributions, `mean`, `p50`, and `p99` are the arithmetic mean, 50th -percentile, and 99th percentile. Do not invent missing values. - -If benchmarking fails, write a failed `final_report.json`. - -Do not write a successful `final_report.json` until the final service answers a -request using `service_model_name` from `Endpoint context:` through the dstack -service URL. +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 +335,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 +362,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/tests/_internal/cli/services/endpoints/test_agent.py b/src/tests/_internal/cli/services/endpoints/test_agent.py index c3c13ef52..91ef2e253 100644 --- a/src/tests/_internal/cli/services/endpoints/test_agent.py +++ b/src/tests/_internal/cli/services/endpoints/test_agent.py @@ -19,6 +19,7 @@ _create_agent_session_directory, _prepare_subprocess_command, _ProgressTailer, + _RecordMirror, _terminate_process, build_endpoint_agent_env, contains_redacted_value, @@ -165,8 +166,8 @@ 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.name == session.timestamp + assert {path.name for path in session.path.iterdir()} == {"agent.log", "session.json"} print_endpoint_progress("creating preset", agent_session=session) assert "creating preset" in session.log_path.read_text() assert "creating preset" in capsys.readouterr().out @@ -180,43 +181,52 @@ def test_saves_log_and_debug_files(self, tmp_path, monkeypatch, capsys): assert {path.name for path in debug_session.path.iterdir()} == { "agent.log", "endpoint.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 success_path == debug_session.path + assert json.loads((debug_session.path / "session.json").read_text()) == { + "status": "success", + "preset_id": "8f3a12c4", + } failed_session = create_endpoint_agent_session(configuration) failed_path = failed_session.finish() - assert failed_path.name.endswith("-failed") + assert failed_path == failed_session.path + assert json.loads((failed_session.path / "session.json").read_text()) == { + "status": "failed", + "preset_id": None, + } def test_avoids_existing_session_directories(self, tmp_path): timestamp = "20260714-120000-000000Z" - (tmp_path / f"{timestamp}-running").mkdir() + (tmp_path / timestamp).mkdir() path = _create_agent_session_directory(tmp_path, timestamp) - assert path.name == f"{timestamp}-1-running" + assert path.name == f"{timestamp}-1" - 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() + 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() - 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", + "preset_id": None, + } def test_reports_invalid_existing_parent(self, tmp_path, monkeypatch): monkeypatch.setenv("HOME", str(tmp_path)) @@ -391,3 +401,64 @@ 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}, + } diff --git a/src/tests/_internal/cli/services/endpoints/test_create.py b/src/tests/_internal/cli/services/endpoints/test_create.py index 99edded5d..78cf2cb3c 100644 --- a/src/tests/_internal/cli/services/endpoints/test_create.py +++ b/src/tests/_internal/cli/services/endpoints/test_create.py @@ -15,10 +15,11 @@ ) from dstack._internal.cli.services.endpoints.create import ( EndpointPresetCreateResult, - _build_prompt, + _build_constraints, _cleanup_runs, _create_endpoint_preset, _get_build_name, + _save_final_report_copy, create_endpoint_preset, ) from dstack._internal.cli.services.endpoints.store import EndpointPresetStore @@ -119,8 +120,11 @@ async def create(**kwargs): paths = list((tmp_path / ".dstack" / "agent" / "qwen").iterdir()) assert len(paths) == 1 - assert paths[0].name.endswith(f"-{preset.id}") - assert {path.name for path in paths[0].iterdir()} == {"agent.log"} + assert {path.name for path in paths[0].iterdir()} == {"agent.log", "session.json"} + assert json.loads((paths[0] / "session.json").read_text()) == { + "status": "success", + "preset_id": preset.id, + } 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 @@ -168,14 +172,18 @@ def fail_finish(self, preset_id=None): paths = list((tmp_path / ".dstack" / "agent" / "qwen").iterdir()) assert len(paths) == 1 - assert paths[0].name.endswith("-running") assert result.preset == preset assert {path.name for path in paths[0].iterdir()} == { "endpoint.dstack.yml", "agent.log", "prompt.md", + "session.json", "trace.jsonl", } + assert json.loads((paths[0] / "session.json").read_text()) == { + "status": "running", + "preset_id": None, + } assert "hf-secret" not in (paths[0] / "endpoint.dstack.yml").read_text() assert "Files remain at" in capsys.readouterr().out @@ -319,41 +327,10 @@ def test_requires_name_and_keeps_generated_prefix_bounded(self, monkeypatch): assert len(f"{build_name}-99999") <= 41 -class TestBuildPrompt: - def test_distinguishes_base_from_exact_model(self): - base_prompt = _build_prompt( - configuration=EndpointConfiguration( - name="qwen", - model={"base": "Qwen/Qwen3.5-27B"}, - context_length=8192, - ), - build_name="qwen-build", - allowed_fleets=("gpu-fleet",), - ) - exact_prompt = _build_prompt( - configuration=EndpointConfiguration( - name="qwen", - model={ - "repo": "community/Qwen3.5-27B-GPTQ-Int4", - "name": "Qwen/Qwen3.5-27B", - }, - ), - build_name="qwen-build", - allowed_fleets=("gpu-fleet",), - ) - - assert "- base_model: Qwen/Qwen3.5-27B" in base_prompt - assert "- model_repo:" not in base_prompt - assert "- context_length: 8192" in base_prompt - assert "- model_repo: community/Qwen3.5-27B-GPTQ-Int4" in exact_prompt - assert "- base_model:" not in exact_prompt - assert "- context_length:" not in exact_prompt - - class TestCleanupRuns: @pytest.mark.asyncio async def test_stops_only_recorded_build_runs(self, tmp_path, monkeypatch): - (tmp_path / "submissions.jsonl").write_text( + (tmp_path / "runs.jsonl").write_text( '{"name":"qwen-build-1"}\n{"name":"unrelated-run"}\n{"name":"qwen-build-2"}\n' ) runs = _FakeRuns() @@ -423,3 +400,85 @@ 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() 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..ce6fb9d5d --- /dev/null +++ b/src/tests/_internal/cli/services/endpoints/test_output.py @@ -0,0 +1,16 @@ +import pytest + +from dstack._internal.cli.services.endpoints.output import _format_number + +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" 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 From 53a116bcd14e5f8323b95c8b996da8585b27aa63 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Mon, 20 Jul 2026 22:02:34 -0700 Subject: [PATCH 02/25] Resumable preset sessions, flat preset store, config shorthand, lifecycle list - Preset ID minted at session start and stable through run names, the session directory, and the final preset; sessions survive Ctrl+C, hard kills, and host restarts and continue with --resume - Agent deaths without a submitted report retry in-process; the retry budget refills only when the failed attempt made progress - Durable workspace under the session directory with a /tmp symlink alias (Unix socket path limit); byte-offset persistence prevents duplicated output on resume; stale running manifests downgrade to interrupted - Flat ~/.dstack/presets// store holding preset.yaml next to session internals; lazy migration from the models--* layout; delete archives to .archive/ instead of destroying - dstack endpoint preset list shows in-flight and interrupted sessions with colored statuses, trial progress after the status, best-trial benchmark and GPU, -w watch mode, and -v details - Top-level base:/repo: configuration shorthand; nested model: deprecated unless model.name is set; --base/--repo filters on list and delete Co-Authored-By: Claude Fable 5 --- src/dstack/_internal/cli/commands/endpoint.py | 122 ++++- src/dstack/_internal/cli/models/endpoints.py | 34 +- .../_internal/cli/services/endpoints/agent.py | 405 ++++++++++++++-- .../cli/services/endpoints/create.py | 171 ++++--- .../cli/services/endpoints/output.py | 114 ++++- .../cli/services/endpoints/presets.py | 5 +- .../_internal/cli/services/endpoints/store.py | 103 +++-- .../cli/services/endpoints/verify.py | 2 + .../_internal/cli/commands/test_endpoint.py | 70 ++- .../_internal/cli/models/test_endpoints.py | 33 ++ .../cli/services/endpoints/test_agent.py | 434 ++++++++++++++++-- .../cli/services/endpoints/test_apply.py | 3 +- .../cli/services/endpoints/test_create.py | 122 ++++- .../cli/services/endpoints/test_output.py | 50 +- .../cli/services/endpoints/test_store.py | 66 ++- 15 files changed, 1466 insertions(+), 268 deletions(-) diff --git a/src/dstack/_internal/cli/commands/endpoint.py b/src/dstack/_internal/cli/commands/endpoint.py index 046276a1e..b419ec335 100644 --- a/src/dstack/_internal/cli/commands/endpoint.py +++ b/src/dstack/_internal/cli/commands/endpoint.py @@ -1,13 +1,21 @@ 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 ( + 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.output import print_endpoint_presets @@ -91,6 +99,11 @@ def _register(self) -> None: 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.set_defaults(subfunc=self._create) apply_parser = preset_subparsers.add_parser( @@ -128,10 +141,15 @@ def _register(self) -> None: help="The preset ID", ) 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" ) @@ -146,21 +164,51 @@ 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) + 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[/]" + ) result = create_endpoint_preset( api=Client.from_config(project_name=args.project), configuration=configuration, store=EndpointPresetStore(), keep_service=args.keep_service, debug=args.debug, + resume_session=resume_session, ) console.print( f"Endpoint preset [code]{result.preset.id}[/] for " @@ -190,34 +238,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) if preset is None: raise CLIError(f"Endpoint preset {args.preset!r} does not exist") + presets = [preset] message = f"Delete endpoint preset [code]{preset.id}[/] for [code]{preset.base}[/]?" else: - presets = [preset for preset in store.list() if preset.base == args.model] + 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 endpoint 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"{'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"Endpoint 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)} endpoint preset{'s' if len(presets) != 1 else ''} " + f"for [code]{args.base or args.repo}[/]" ) @@ -240,11 +289,56 @@ def _add_configuration_args(parser: argparse.ArgumentParser) -> None: 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: diff --git a/src/dstack/_internal/cli/models/endpoints.py b/src/dstack/_internal/cli/models/endpoints.py index e902a5c40..29335d883 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, @@ -94,10 +94,24 @@ class EndpointConfiguration( 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 context_length: Annotated[ Optional[PositiveInt], Field(description="The minimum required context length") ] = None @@ -144,6 +158,22 @@ def effective_max_trials(self) -> int: 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): diff --git a/src/dstack/_internal/cli/services/endpoints/agent.py b/src/dstack/_internal/cli/services/endpoints/agent.py index 5f124fd9e..f43af994f 100644 --- a/src/dstack/_internal/cli/services/endpoints/agent.py +++ b/src/dstack/_internal/cli/services/endpoints/agent.py @@ -1,16 +1,17 @@ import asyncio import json import os +import secrets import shutil import signal import subprocess import sys import tempfile -from contextlib import contextmanager, suppress +from contextlib import 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, Optional, Sequence import psutil import yaml @@ -40,6 +41,12 @@ _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 = ( @@ -118,6 +125,7 @@ class EndpointAgentSession: path: Path timestamp: str debug: bool + preset_id: str = "" _log_enabled: bool = field(default=True, init=False, repr=False) @property @@ -172,12 +180,23 @@ 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: - status = { - "status": "success" if preset_id is not None else "failed", - "preset_id": preset_id, - } - _write_private_text(self.path / _SESSION_FILENAME, json.dumps(status, indent=2) + "\n") + 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 @@ -193,16 +212,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( @@ -213,16 +300,31 @@ def create_endpoint_agent_session( if configuration.name is None: raise CLIError("Endpoint name is required to save agent output") timestamp = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S-%fZ") - parent = get_dstack_dir() / "agent" / configuration.name + 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", "") - _write_private_text( - path / _SESSION_FILENAME, - json.dumps({"status": "running", "preset_id": None}, indent=2) + "\n", - ) + manifest = { + "id": preset_id, + "status": "running", + "pid": os.getpid(), + "endpoint": 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") if debug: data = json.loads(configuration.json(exclude_none=True)) if configuration.env: @@ -239,7 +341,7 @@ def create_endpoint_agent_session( 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 _get_claude_version(auth: "ClaudeAuth") -> Optional[str]: @@ -273,17 +375,112 @@ def _get_claude_auth_status(auth: "ClaudeAuth") -> dict[str, Any]: return {"authMethod": "unknown"} -def _create_agent_session_directory(parent: Path, timestamp: str) -> Path: - index = 0 - while True: - name = timestamp if index == 0 else f"{timestamp}-{index}" - path = parent / name +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") + pid = manifest.get("pid") + if ( + status == "running" + and isinstance(pid, int) + and pid > 0 + and pid != os.getpid() + and psutil.pid_exists(pid) + ): + raise CLIError(f"Session {preset_id} appears to be running already (pid {pid})") + 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 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"): + continue + pid = manifest.get("pid") + if status == "running" and not ( + isinstance(pid, int) and pid > 0 and psutil.pid_exists(pid) + ): + status = "interrupted" + entry = dict(manifest) + entry["id"] = path.name + 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 + task_names: set[str] = set() + 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 - return path + if not isinstance(record, dict): + continue + # A trial may log several benchmark records (e.g. a re-run); records + # sharing a task name are one trial. + task = record.get("task") + task_name = task.get("name") if isinstance(task, dict) else None + if isinstance(task_name, str) and task_name: + task_names.add(task_name) + else: + count += 1 + benchmark = record.get("benchmark") + if not isinstance(benchmark, dict): + continue + metrics = benchmark.get("metrics") or {} + workload = benchmark.get("workload") or {} + duration = metrics.get("duration_seconds") + tokens = metrics.get("total_output_tokens") + if not isinstance(duration, (int, float)) or duration <= 0: + continue + if not isinstance(tokens, (int, float)): + continue + tok_s = tokens / duration + if best is None or tok_s > best["tok_s"]: + resources = record.get("resources") or {} + gpu = resources.get("gpu") if isinstance(resources, dict) else None + gpu_text = None + if isinstance(gpu, dict) and gpu.get("name"): + gpu_text = str(gpu["name"]) + if gpu.get("memory"): + gpu_text += f":{gpu['memory']}" + if gpu.get("count"): + gpu_text += f":{gpu['count']}" + best = { + "tok_s": tok_s, + "concurrency": workload.get("concurrency"), + "gpu": gpu_text, + } + return {"count": count + len(task_names), "best": best} def get_claude_auth() -> ClaudeAuth: @@ -354,28 +551,97 @@ 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)) + 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: + 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) + 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 _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( @@ -422,24 +688,13 @@ async def run_endpoint_agent( if proc is not None and proc.returncode is None: await _terminate_process(proc) raise - 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() 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: @@ -615,7 +870,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", @@ -641,6 +898,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 @@ -724,7 +983,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" @@ -784,6 +1052,27 @@ def _terminate_windows_process_tree(pid: int) -> None: psutil.wait_procs(alive, timeout=3) +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 + with suppress(OSError): + _write_private_text(self._path, json.dumps(self._data) + "\n") + + class _ProgressTailer: def __init__( self, @@ -791,11 +1080,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: @@ -809,6 +1102,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: @@ -821,11 +1116,21 @@ 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]) -> None: + 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 = 0 + 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: @@ -845,6 +1150,8 @@ def flush(self) -> None: return chunk = data[: end + 1].decode("utf-8", errors="replace") self._offset += end + 1 + if self._offset_store is not None and self._offset_key: + self._offset_store.set(self._offset_key, self._offset) try: if not self._target.exists(): _write_private_text(self._target, "") diff --git a/src/dstack/_internal/cli/services/endpoints/create.py b/src/dstack/_internal/cli/services/endpoints/create.py index 7b353ff80..89b9182ae 100644 --- a/src/dstack/_internal/cli/services/endpoints/create.py +++ b/src/dstack/_internal/cli/services/endpoints/create.py @@ -1,7 +1,7 @@ import asyncio +import dataclasses import json import os -import secrets import uuid from contextlib import suppress from dataclasses import dataclass @@ -17,15 +17,17 @@ from dstack._internal.cli.services.endpoints.agent import ( EndpointAgentSession, EndpointAgentWorkspace, + attach_agent_workspace, build_endpoint_agent_env, contains_redacted_value, + create_agent_workspace, create_endpoint_agent_session, - endpoint_agent_workspace, get_claude_auth, get_redacted_values, get_sensitive_inherited_env_values, print_endpoint_progress, redact, + remove_agent_workspace, run_endpoint_agent, ) from dstack._internal.cli.services.endpoints.presets import endpoint_preset_to_data @@ -60,8 +62,9 @@ def create_endpoint_preset( keep_service: bool = False, build_name: Optional[str] = None, debug: bool = False, + resume_session: Optional[EndpointAgentSession] = 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( @@ -73,12 +76,18 @@ def create_endpoint_preset( keep_service=keep_service, build_name=build_name, agent_session=agent_session, + resume=resume_session is not None, ) ) + except KeyboardInterrupt: + _suspend_agent_session(agent_session) + 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 @@ -91,13 +100,30 @@ async def _create_endpoint_preset( keep_service: bool = False, build_name: Optional[str] = None, agent_session: EndpointAgentSession, + resume: bool = False, ) -> 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 resume: + 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: tuple[str, ...] = () + else: + 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, agent_session.preset_id) + agent_session.update_manifest(status="running", pid=os.getpid(), claude_model=auth.model) endpoint_env = configuration.env.as_dict() token = getattr(api.client, "_token", None) @@ -115,61 +141,68 @@ async def _create_endpoint_preset( 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: - env = build_endpoint_agent_env( - api=api, - endpoint_env=endpoint_env, - auth=auth, - workspace=workspace, - token=token, - ) + env = build_endpoint_agent_env( + api=api, + endpoint_env=endpoint_env, + auth=auth, + workspace=workspace, + token=token, + ) + prompt = get_endpoint_agent_system_prompt() + if not resume: constraints_text = _build_constraints( configuration=configuration, build_name=build_name, allowed_fleets=allowed_fleets, ) workspace.constraints_path.write_text(constraints_text, encoding="utf-8") - prompt = get_endpoint_agent_system_prompt() if agent_session.debug: agent_session.write_prompt(prompt) agent_session.write_constraints(constraints_text) agent_session.write_agent_info(auth) - try: - process_output = await run_endpoint_agent( - prompt=prompt, - env=env, + try: + process_output = await run_endpoint_agent( + prompt=prompt, + env=env, + workspace=workspace, + 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, + 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, + ) + if contains_redacted_value(endpoint_preset_to_data(preset), redacted_values): + raise CLIError("Generated endpoint preset contains a secret value") + preset_path = store.save(preset) + print_endpoint_progress( + f"Saved endpoint preset {preset.id} for {preset.base} at {preset_path}.", + agent_session=agent_session, + ) + creation_succeeded = True + except (KeyboardInterrupt, asyncio.CancelledError): + interrupted = True + raise + finally: + if agent_session.debug: + _save_final_report_copy( workspace=workspace, - auth=auth, - redacted_values=redacted_values, agent_session=agent_session, - ) - report = load_endpoint_agent_report( - output=process_output, - workspace=workspace, redacted_values=redacted_values, ) - run = api.client.runs.get(api.project, report.run_name) - preset = build_verified_endpoint_preset( - run=run, - endpoint_configuration=source_configuration, - report=report, - ) - if contains_redacted_value(endpoint_preset_to_data(preset), redacted_values): - raise CLIError("Generated endpoint preset contains a secret value") - preset_path = store.save(preset) - print_endpoint_progress( - f"Saved endpoint preset {preset.id} for {preset.base} at {preset_path}.", - agent_session=agent_session, - ) - creation_succeeded = True - finally: - if agent_session.debug: - _save_final_report_copy( - workspace=workspace, - agent_session=agent_session, - redacted_values=redacted_values, - ) + if not interrupted: keep_final_service = keep_service and creation_succeeded try: await _cleanup_runs( @@ -221,25 +254,51 @@ 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: +def _suspend_agent_session(session: EndpointAgentSession) -> None: + try: + session.finish("interrupted") + except OSError as e: + warn(f"Could not record the interrupted session state: {e}") + 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 endpoint preset create -f " + f"--resume {session.preset_id}[/], or stop the runs with [code]dstack stop[/]." + ) + + +def _get_build_name(endpoint_name: Optional[str], suffix: 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("-") + # Leave room for the preset id and numeric submission suffix while retaining + # a recognizable prefix. + prefix = endpoint_name[:26].rstrip("-") return f"{prefix}-{suffix}" +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 _get_allowed_fleets(api: Client, configuration: EndpointConfiguration) -> tuple[str, ...]: if configuration.fleets is not None: return tuple( diff --git a/src/dstack/_internal/cli/services/endpoints/output.py b/src/dstack/_internal/cli/services/endpoints/output.py index d040978c3..b6ed587e9 100644 --- a/src/dstack/_internal/cli/services/endpoints/output.py +++ b/src/dstack/_internal/cli/services/endpoints/output.py @@ -1,4 +1,6 @@ from collections import defaultdict +from datetime import datetime +from typing import Any, Optional from rich.table import Table @@ -9,43 +11,111 @@ 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": ("done", "grey"), + "running": ("clauding", "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 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("RESOURCES" if verbose else "GPU", style="secondary") + if verbose: + table.add_column("CONTEXT", justify="right", 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: + repo_to_base[preset.model] = preset.base + sessions_by_model: dict[str, list[dict[str, Any]]] = defaultdict(list) + for session in sessions or []: + 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": f"[bold]{base}[/]"}) + for preset in presets_by_base.get(base, []): _add_preset(table, preset, verbose=verbose) + for session in sessions_by_model.get(base, []): + _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]) -> 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 = _format_status(str(session.get("status", ""))) + trials = session.get("trials") + max_trials = session.get("max_trials") + if isinstance(trials, dict) and (trials.get("count") or isinstance(max_trials, int)): + progress = str(trials.get("count") or 0) + if isinstance(max_trials, int): + progress += f"/{max_trials}" + # The trial progress stays outside the status markup to render in the + # default color. + status += f" ({progress})" + 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", "")), + "GPU": gpu, + "RESOURCES": gpu, + "BENCHMARK": benchmark, + "STATUS": status, + "SUBMITTED": created, }, ) - if preset.model != preset.base: + + +def _add_preset(table: Table, preset: EndpointPreset, *, verbose: bool) -> None: + groups = preset.service.replica_groups + column = "RESOURCES" if verbose else "GPU" + row = { + "ID": preset.id, + column: _format_resources(groups[0].resources, verbose=verbose), + "STATUS": _format_status("ready"), + "BENCHMARK": format_endpoint_benchmark(preset, verbose=verbose), + "SUBMITTED": pretty_date(preset.created_at), + } + if verbose: + row["CONTEXT"] = format_endpoint_context_length(preset) + add_row_from_dict(table, row) + if verbose and preset.model != preset.base: add_row_from_dict( table, - {"MODEL": f" repo={preset.model}"}, + {"BASE": f" repo={preset.model}"}, style="secondary", ) if len(groups) > 1: @@ -53,7 +123,7 @@ def _add_preset(table: Table, preset: EndpointPreset, *, verbose: bool) -> None: add_row_from_dict( table, { - "MODEL": f" group={group.name}", + "BASE": f" group={group.name}", column: _format_resources(group.resources, verbose=verbose), }, style="secondary", @@ -72,7 +142,7 @@ def format_endpoint_benchmark(preset: EndpointPreset, *, verbose: bool = False) 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)}", ] diff --git a/src/dstack/_internal/cli/services/endpoints/presets.py b/src/dstack/_internal/cli/services/endpoints/presets.py index 727133fa3..9f0b44e52 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,7 @@ def build_endpoint_preset( model: str, context_length: int, benchmark: EndpointBenchmark, + preset_id: Optional[str] = None, ) -> EndpointPreset: service = service.copy(deep=True) service.name = None @@ -39,7 +40,7 @@ def build_endpoint_preset( set_service_gpu_vendors_from_validations(service, [validation]) return EndpointPreset( 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(), diff --git a/src/dstack/_internal/cli/services/endpoints/store.py b/src/dstack/_internal/cli/services/endpoints/store.py index bcb7c24db..efe8e791c 100644 --- a/src/dstack/_internal/cli/services/endpoints/store.py +++ b/src/dstack/_internal/cli/services/endpoints/store.py @@ -1,8 +1,10 @@ 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 @@ -10,38 +12,48 @@ from dstack._internal.cli.models.endpoint_presets import EndpointPreset from dstack._internal.cli.models.endpoints import EndpointConfiguration from dstack._internal.cli.services.endpoints.presets import endpoint_preset_to_data +from dstack._internal.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") 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", ) @@ -62,27 +74,31 @@ 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: @@ -91,23 +107,10 @@ def _load(self, path: Path) -> EndpointPreset: except (OSError, ValidationError, yaml.YAMLError) as e: raise CLIError(f"Invalid endpoint preset file {path}: {e}") from e - def _path(self, base: str, preset_id: str) -> Path: - if not preset_id or any(char in preset_id for char in "/\\"): - raise CLIError("Endpoint preset ID must not contain path separators") - return self._directory(base) / f"{preset_id}.yaml" - - def _find_preset_paths(self, preset_id: str) -> List[Path]: - if not preset_id or any(char in preset_id for char in "/\\"): - raise CLIError("Endpoint preset ID must not contain path separators") - return [ - path - for directory in self.root.glob("models--*") - if (path := directory / f"{preset_id}.yaml").is_file() - ] - def _directory(self, base: str) -> Path: - directory = "models--" + base.replace("/", "--").replace("\\", "--") - return self.root / directory +def _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 endpoint preset ID: {preset_id!r}") def load_endpoint_configuration(path: str) -> tuple[str, EndpointConfiguration]: @@ -129,8 +132,16 @@ def _parse_endpoint_configuration(stream: TextIO) -> EndpointConfiguration: data = yaml.safe_load(stream) if not isinstance(data, dict): raise ConfigurationError("Endpoint configuration must be a YAML object") - return EndpointConfiguration.parse_obj(data) + 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 + 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 diff --git a/src/dstack/_internal/cli/services/endpoints/verify.py b/src/dstack/_internal/cli/services/endpoints/verify.py index d5f2872ec..3953dfd21 100644 --- a/src/dstack/_internal/cli/services/endpoints/verify.py +++ b/src/dstack/_internal/cli/services/endpoints/verify.py @@ -61,6 +61,7 @@ def build_verified_endpoint_preset( run: Run, endpoint_configuration: EndpointConfiguration, report: AgentFinalReport, + preset_id: 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") @@ -107,6 +108,7 @@ def build_verified_endpoint_preset( model=report.model, context_length=report.context_length, benchmark=benchmark, + preset_id=preset_id, ) diff --git a/src/tests/_internal/cli/commands/test_endpoint.py b/src/tests/_internal/cli/commands/test_endpoint.py index ac843c910..0b25f9a08 100644 --- a/src/tests/_internal/cli/commands/test_endpoint.py +++ b/src/tests/_internal/cli/commands/test_endpoint.py @@ -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,7 +67,7 @@ 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): @@ -118,13 +118,15 @@ def test_lists_and_deletes_preset_without_api_client(self, tmp_path, capsys): 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 @@ -132,11 +134,12 @@ def test_lists_and_deletes_preset_without_api_client(self, tmp_path, capsys): 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()) + assert "hardware=" in joined_verbose + assert "api=" in joined_verbose + assert "n=16" in joined_verbose + assert "con=1" in joined_verbose + assert "1K->128" in joined_verbose with patch("dstack.api.Client.from_config") as from_config: assert ( @@ -198,7 +201,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 +212,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"], + ["endpoint", "preset", "delete", flag, getattr(preset, attribute), "-y"], home_dir=tmp_path, ) == 0 @@ -216,6 +221,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( + ["endpoint", "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 = ["endpoint", "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( 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 91ef2e253..154827252 100644 --- a/src/tests/_internal/cli/services/endpoints/test_agent.py +++ b/src/tests/_internal/cli/services/endpoints/test_agent.py @@ -4,6 +4,7 @@ import shutil import subprocess import sys +from pathlib import Path from types import SimpleNamespace import psutil @@ -16,18 +17,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 @@ -121,8 +125,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() @@ -137,8 +142,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" @@ -152,6 +158,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)) @@ -165,9 +180,16 @@ 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 == session.timestamp + 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"} + 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 @@ -187,28 +209,14 @@ def test_saves_log_and_debug_files(self, tmp_path, monkeypatch, capsys): 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") + success_path = debug_session.finish("success") assert success_path == debug_session.path - assert json.loads((debug_session.path / "session.json").read_text()) == { - "status": "success", - "preset_id": "8f3a12c4", - } + assert json.loads((debug_session.path / "session.json").read_text())["status"] == "success" failed_session = create_endpoint_agent_session(configuration) - failed_path = failed_session.finish() + failed_path = failed_session.finish("failed") assert failed_path == failed_session.path - assert json.loads((failed_session.path / "session.json").read_text()) == { - "status": "failed", - "preset_id": None, - } - - def test_avoids_existing_session_directories(self, tmp_path): - timestamp = "20260714-120000-000000Z" - (tmp_path / timestamp).mkdir() - - path = _create_agent_session_directory(tmp_path, timestamp) - - assert path.name == f"{timestamp}-1" + 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" @@ -220,20 +228,16 @@ def test_finish_writes_status_in_place(self, tmp_path): debug=False, ) - path = session.finish() + path = session.finish("failed") assert path == session_dir - assert json.loads((session_dir / "session.json").read_text()) == { - "status": "failed", - "preset_id": None, - } + 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( @@ -462,3 +466,365 @@ def test_writes_model_params_and_auth(self, tmp_path, monkeypatch): "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_distinct_task_names_not_benchmark_records(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) + + assert summary["count"] == 3 + assert summary["best"] == {"tok_s": 2300.0, "concurrency": 8, "gpu": "A40:48GB:1"} diff --git a/src/tests/_internal/cli/services/endpoints/test_apply.py b/src/tests/_internal/cli/services/endpoints/test_apply.py index c9962ffd3..3f629632e 100644 --- a/src/tests/_internal/cli/services/endpoints/test_apply.py +++ b/src/tests/_internal/cli/services/endpoints/test_apply.py @@ -139,8 +139,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]context=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 78cf2cb3c..975f5c68a 100644 --- a/src/tests/_internal/cli/services/endpoints/test_create.py +++ b/src/tests/_internal/cli/services/endpoints/test_create.py @@ -11,7 +11,9 @@ EndpointAgentProcessOutput, EndpointAgentSession, EndpointAgentWorkspace, + create_agent_workspace, print_endpoint_progress, + remove_agent_workspace, ) from dstack._internal.cli.services.endpoints.create import ( EndpointPresetCreateResult, @@ -77,7 +79,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, @@ -118,13 +120,16 @@ 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 {path.name for path in paths[0].iterdir()} == {"agent.log", "session.json"} - assert json.loads((paths[0] / "session.json").read_text()) == { - "status": "success", - "preset_id": preset.id, - } + 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 @@ -170,7 +175,11 @@ 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 result.preset == preset assert {path.name for path in paths[0].iterdir()} == { @@ -180,10 +189,7 @@ def fail_finish(self, preset_id=None): "session.json", "trace.jsonl", } - assert json.loads((paths[0] / "session.json").read_text()) == { - "status": "running", - "preset_id": None, - } + assert json.loads((paths[0] / "session.json").read_text())["status"] == "running" assert "hf-secret" not in (paths[0] / "endpoint.dstack.yml").read_text() assert "Files remain at" in capsys.readouterr().out @@ -237,7 +243,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 @@ -264,8 +270,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"), @@ -302,6 +307,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, ) @@ -314,16 +320,13 @@ async def run_agent(**kwargs): class TestBuildName: - def test_requires_name_and_keeps_generated_prefix_bounded(self, monkeypatch): + def test_requires_name_and_keeps_generated_prefix_bounded(self): 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", - ) + _get_build_name(None, "a1b2c3d4") - build_name = _get_build_name("qwen-endpoint-with-a-name-that-is-forty-one") + build_name = _get_build_name("qwen-endpoint-with-a-name-that-is-forty-one", "a1b2c3d4") + assert build_name.endswith("-a1b2c3d4") assert len(f"{build_name}-99999") <= 41 @@ -369,6 +372,7 @@ def _agent_session(tmp_path, *, debug: bool = False) -> EndpointAgentSession: path=path, timestamp="20260714-120000Z", debug=debug, + preset_id="ab12cd34", ) @@ -482,3 +486,79 @@ def test_missing_report_is_no_op(self, tmp_path): ) 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) diff --git a/src/tests/_internal/cli/services/endpoints/test_output.py b/src/tests/_internal/cli/services/endpoints/test_output.py index ce6fb9d5d..7e7125495 100644 --- a/src/tests/_internal/cli/services/endpoints/test_output.py +++ b/src/tests/_internal/cli/services/endpoints/test_output.py @@ -1,6 +1,7 @@ import pytest +from rich.table import Table -from dstack._internal.cli.services.endpoints.output import _format_number +from dstack._internal.cli.services.endpoints.output import _add_session, _format_number pytestmark = pytest.mark.windows @@ -14,3 +15,50 @@ 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[/] (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[/] (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[/] (2)" diff --git a/src/tests/_internal/cli/services/endpoints/test_store.py b/src/tests/_internal/cli/services/endpoints/test_store.py index 6088ec8cd..eba675101 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,6 +13,7 @@ EndpointBenchmarkMetrics, EndpointBenchmarkWorkload, ) +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.models.envs import EnvSentinel @@ -68,7 +71,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 +92,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 +151,35 @@ 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: endpoint\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: endpoint\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 From 6bf004f987e8d4ce81f6e2cbf4fc4f1d6b5e1685 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Mon, 20 Jul 2026 22:24:53 -0700 Subject: [PATCH 03/25] Tie agent life to CLI life; sort preset list newest-first MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - A watchdog process holds a pipe from the CLI and kills the agent's process tree the moment the pipe snaps — covering deaths the CLI cannot react to (SIGKILL included), which previously left an orphaned agent able to keep submitting GPU runs. Graceful stops disarm the watchdog. - POSIX kills the agent's process group; Windows kills the process tree - Presets and sessions sort newest-first within each base-model group, matching dstack ps Co-Authored-By: Claude Fable 5 --- .../_internal/cli/services/endpoints/agent.py | 77 +++++++++++++++++++ .../cli/services/endpoints/output.py | 11 ++- .../cli/services/endpoints/test_agent.py | 55 +++++++++++++ .../cli/services/endpoints/test_output.py | 43 +++++++++++ 4 files changed, 184 insertions(+), 2 deletions(-) diff --git a/src/dstack/_internal/cli/services/endpoints/agent.py b/src/dstack/_internal/cli/services/endpoints/agent.py index f43af994f..c2cae73f4 100644 --- a/src/dstack/_internal/cli/services/endpoints/agent.py +++ b/src/dstack/_internal/cli/services/endpoints/agent.py @@ -643,6 +643,7 @@ async def _run_claude_process( agent_session: EndpointAgentSession, ) -> tuple[EndpointAgentProcessOutput, int]: proc: Optional[asyncio.subprocess.Process] = None + watchdog: Optional[subprocess.Popen] = None try: proc = await asyncio.create_subprocess_exec( *command, @@ -654,6 +655,7 @@ async def _run_claude_process( start_new_session=not IS_WINDOWS, limit=_CLAUDE_STREAM_LIMIT, ) + watchdog = _start_agent_watchdog(proc.pid) assert proc.stdin is not None assert proc.stdout is not None assert proc.stderr is not None @@ -688,6 +690,9 @@ async def _run_claude_process( if proc is not None and proc.returncode is None: await _terminate_process(proc) raise + finally: + if watchdog is not None: + await asyncio.to_thread(_stop_agent_watchdog, watchdog) output = stdout_output if output.report_data is None: @@ -1036,6 +1041,78 @@ async def _terminate_process(proc: asyncio.subprocess.Process) -> None: await proc.wait() +# The watchdog reads its stdin until EOF. The CLI holds the write end and never +# writes while the agent runs, so EOF means the CLI is gone — any death, +# including SIGKILL, which the CLI cannot react to itself. A disarm byte is +# written when the CLI stops the agent on its own, so the watchdog exits +# without touching anything (and cannot kill a reused pid). +_AGENT_WATCHDOG_SCRIPT = """ +import os +import signal +import sys +import time + +import psutil + +agent_pid = int(sys.argv[1]) +if sys.stdin.buffer.read(): + sys.exit(0) # disarmed: the CLI stopped the agent itself + + +def _kill(hard): + if os.name == "posix": + try: + os.killpg(agent_pid, signal.SIGKILL if hard else signal.SIGTERM) + except OSError: + pass + return + try: + root = psutil.Process(agent_pid) + processes = [*root.children(recursive=True), root] + except psutil.NoSuchProcess: + return + for process in processes: + try: + process.kill() if hard else process.terminate() + except psutil.NoSuchProcess: + pass + + +if psutil.pid_exists(agent_pid): + _kill(hard=False) + deadline = time.monotonic() + 15 + while time.monotonic() < deadline and psutil.pid_exists(agent_pid): + time.sleep(0.5) + if psutil.pid_exists(agent_pid): + _kill(hard=True) +""" + + +def _start_agent_watchdog(agent_pid: int) -> subprocess.Popen: + """Tie the agent's life to this process: if this process dies in a way it + cannot react to, the watchdog kills the agent's process tree.""" + return subprocess.Popen( + [sys.executable, "-c", _AGENT_WATCHDOG_SCRIPT, str(agent_pid)], + stdin=subprocess.PIPE, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=not IS_WINDOWS, + ) + + +def _stop_agent_watchdog(watchdog: subprocess.Popen) -> None: + if watchdog.stdin is not None: + with suppress(OSError, ValueError): + watchdog.stdin.write(b"x") + watchdog.stdin.flush() + with suppress(OSError, ValueError): + watchdog.stdin.close() + try: + watchdog.wait(timeout=10) + except subprocess.TimeoutExpired: + watchdog.kill() + + def _terminate_windows_process_tree(pid: int) -> None: try: root = psutil.Process(pid) diff --git a/src/dstack/_internal/cli/services/endpoints/output.py b/src/dstack/_internal/cli/services/endpoints/output.py index b6ed587e9..db24f8073 100644 --- a/src/dstack/_internal/cli/services/endpoints/output.py +++ b/src/dstack/_internal/cli/services/endpoints/output.py @@ -50,9 +50,16 @@ def print_endpoint_presets( for base in sorted({*presets_by_base, *sessions_by_model}, key=str.lower): add_row_from_dict(table, {"BASE": f"[bold]{base}[/]"}) - for preset in presets_by_base.get(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) - for session in sessions_by_model.get(base, []): + 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() diff --git a/src/tests/_internal/cli/services/endpoints/test_agent.py b/src/tests/_internal/cli/services/endpoints/test_agent.py index 154827252..7611ed3df 100644 --- a/src/tests/_internal/cli/services/endpoints/test_agent.py +++ b/src/tests/_internal/cli/services/endpoints/test_agent.py @@ -2,8 +2,10 @@ import json import os import shutil +import signal import subprocess import sys +from contextlib import suppress from pathlib import Path from types import SimpleNamespace @@ -20,6 +22,8 @@ _prepare_subprocess_command, _ProgressTailer, _RecordMirror, + _start_agent_watchdog, + _stop_agent_watchdog, _summarize_session_trials, _terminate_process, attach_agent_workspace, @@ -407,6 +411,57 @@ async def test_terminates_windows_process_tree(self): assert not psutil.pid_exists(child_pid) +class TestAgentWatchdog: + @pytest.mark.skipif(IS_WINDOWS, reason="exercises POSIX signals and sessions") + def test_kills_agent_when_cli_dies_hard(self): + # A fake CLI spawns a fake agent in its own session, arms the real + # watchdog, and hangs. SIGKILL of the CLI must take the agent down. + cli_script = ( + "import subprocess, sys, time\n" + "from dstack._internal.cli.services.endpoints.agent import _start_agent_watchdog\n" + "agent = subprocess.Popen(\n" + " [sys.executable, '-c', 'import time; time.sleep(300)'], start_new_session=True\n" + ")\n" + "_start_agent_watchdog(agent.pid)\n" + "print(agent.pid, flush=True)\n" + "time.sleep(300)\n" + ) + cli = subprocess.Popen( + [sys.executable, "-c", cli_script], stdout=subprocess.PIPE, text=True + ) + try: + assert cli.stdout is not None + agent_pid = int(cli.stdout.readline()) + assert psutil.pid_exists(agent_pid) + + cli.kill() + cli.wait(timeout=10) + + try: + psutil.Process(agent_pid).wait(timeout=20) + except psutil.NoSuchProcess: + pass + assert not psutil.pid_exists(agent_pid) + finally: + with suppress(OSError): + cli.kill() + + @pytest.mark.skipif(IS_WINDOWS, reason="exercises POSIX signals and sessions") + def test_disarmed_watchdog_leaves_agent_untouched(self): + agent = subprocess.Popen( + [sys.executable, "-c", "import time; time.sleep(300)"], start_new_session=True + ) + try: + watchdog = _start_agent_watchdog(agent.pid) + _stop_agent_watchdog(watchdog) + + assert watchdog.poll() is not None + assert agent.poll() is None + finally: + with suppress(OSError): + os.killpg(agent.pid, signal.SIGKILL) + + class TestRecordMirror: def test_mirrors_complete_lines_redacted(self, tmp_path): source = tmp_path / "runs.jsonl" diff --git a/src/tests/_internal/cli/services/endpoints/test_output.py b/src/tests/_internal/cli/services/endpoints/test_output.py index 7e7125495..bc00817a7 100644 --- a/src/tests/_internal/cli/services/endpoints/test_output.py +++ b/src/tests/_internal/cli/services/endpoints/test_output.py @@ -62,3 +62,46 @@ def test_counts_without_max_trials(self): row = _session_row({"id": "ab12cd34", "status": "interrupted", "trials": {"count": 2}}) assert row["STATUS"] == "[bold gold1]interrupted[/] (2)" + + +class TestGroupOrdering: + def test_sorts_presets_and_sessions_newest_first(self, monkeypatch): + from datetime import timedelta + from io import StringIO + + from rich.console import Console + from rich.theme import Theme + + from dstack._internal.cli.services.endpoints import output as output_module + from tests._internal.cli.endpoint_presets import get_endpoint_preset + + 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") From 91b0c86099b2627c5b40ebfa6617ae215862cf13 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Tue, 21 Jul 2026 11:17:45 -0700 Subject: [PATCH 04/25] Redact known secrets from the agent report instead of failing the session A final report echoing a known secret (e.g. the bearer token in benchmark.command) was rejected wholesale, discarding an otherwise verified session at its very last step. Scrub known secret values from the report structure before validation; the bearer check still rejects unknown leaked tokens. Co-Authored-By: Claude Fable 5 --- .../cli/services/endpoints/verify.py | 14 +++++++ .../cli/services/endpoints/test_verify.py | 39 +++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/src/dstack/_internal/cli/services/endpoints/verify.py b/src/dstack/_internal/cli/services/endpoints/verify.py index 3953dfd21..9886c1b76 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: diff --git a/src/tests/_internal/cli/services/endpoints/test_verify.py b/src/tests/_internal/cli/services/endpoints/test_verify.py index 29452b4b8..aa55f2ae9 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,37 @@ 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",)) From 289c18437f3ed52bf95e4a252f24d8b4b2d9fe66 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Tue, 21 Jul 2026 11:59:12 -0700 Subject: [PATCH 05/25] Rename to `dstack preset`, add apply --id, fix report validation, polish list - Configuration `type: preset`; flat `dstack preset {list,get,create,apply,delete}` command tree; user-facing strings follow (internal names are a follow-up) - `apply` selects among repeatable `--id` candidates (given order, capacity-aware) replacing the machine-local `preset:` configuration field - The bearer-token report guard only rejects credential-shaped values; prose such as "auth via bearer header from env" failed two live sessions - Preset rows show `success (n/m)` from the completed creation session; trial count is the record count again (one long-lived task commonly hosts several trials, so distinct task names undercount) - List polish: dimmed base group rows and trial progress, green `success` Co-Authored-By: Claude Fable 5 --- src/dstack/_internal/cli/commands/endpoint.py | 63 ++++++++++--------- .../_internal/cli/models/endpoint_presets.py | 5 ++ src/dstack/_internal/cli/models/endpoints.py | 13 +--- .../_internal/cli/services/endpoints/agent.py | 24 +++---- .../_internal/cli/services/endpoints/apply.py | 39 +++++++----- .../cli/services/endpoints/create.py | 10 +-- .../cli/services/endpoints/output.py | 53 +++++++++++----- .../cli/services/endpoints/presets.py | 2 +- .../_internal/cli/services/endpoints/store.py | 10 +-- .../cli/services/endpoints/verify.py | 4 +- .../_internal/cli/commands/test_endpoint.py | 43 +++++++------ .../cli/services/endpoints/test_agent.py | 12 ++-- .../cli/services/endpoints/test_apply.py | 35 +++++++---- .../cli/services/endpoints/test_create.py | 6 +- .../cli/services/endpoints/test_output.py | 50 +++++++++++---- .../cli/services/endpoints/test_store.py | 4 +- .../cli/services/endpoints/test_verify.py | 16 +++++ 17 files changed, 235 insertions(+), 154 deletions(-) diff --git a/src/dstack/_internal/cli/commands/endpoint.py b/src/dstack/_internal/cli/commands/endpoint.py index b419ec335..93c6f158e 100644 --- a/src/dstack/_internal/cli/commands/endpoint.py +++ b/src/dstack/_internal/cli/commands/endpoint.py @@ -33,8 +33,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( @@ -43,20 +43,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) @@ -64,7 +57,7 @@ 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") @@ -78,7 +71,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) @@ -108,12 +101,18 @@ def _register(self) -> None: 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. Can be repeated", + ) apply_parser.add_argument( "-y", "--yes", action="store_true", help="Do not ask for confirmation" ) @@ -130,7 +129,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) @@ -211,7 +210,7 @@ def _create(self, args: argparse.Namespace) -> None: resume_session=resume_session, ) 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: @@ -220,7 +219,7 @@ def _create(self, args: argparse.Namespace) -> None: def _get(self, args: argparse.Namespace) -> None: preset = EndpointPresetStore().get(args.preset) if preset is None: - raise CLIError(f"Endpoint preset {args.preset!r} does not exist") + raise CLIError(f"Preset {args.preset!r} does not exist") print(preset.json()) def _apply(self, args: argparse.Namespace) -> None: @@ -230,7 +229,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(), @@ -241,17 +240,17 @@ def _delete(self, args: argparse.Namespace) -> None: if args.preset is not None: preset = store.get(args.preset) if preset is None: - raise CLIError(f"Endpoint preset {args.preset!r} does not exist") + raise CLIError(f"Preset {args.preset!r} does not exist") presets = [preset] - message = f"Delete endpoint preset [code]{preset.id}[/] for [code]{preset.base}[/]?" + message = f"Delete preset [code]{preset.id}[/] for [code]{preset.base}[/]?" else: target = args.base or args.repo presets = _filter_presets(store.list(), base=args.base, repo=args.repo) if not presets: kind = "base model" if args.base else "model repo" - raise CLIError(f"No endpoint presets found for {kind} {target!r}") + raise CLIError(f"No presets found for {kind} {target!r}") message = ( - f"Delete {len(presets)} endpoint preset" + f"Delete {len(presets)} preset" f"{'s' if len(presets) != 1 else ''} for [code]{target}[/]?" ) if not args.yes and not confirm_ask(message): @@ -261,11 +260,11 @@ def _delete(self, args: argparse.Namespace) -> None: store.delete(preset.id) if args.preset is not None: console.print( - f"Endpoint preset [code]{presets[0].id}[/] for [code]{presets[0].base}[/] deleted" + f"Preset [code]{presets[0].id}[/] for [code]{presets[0].base}[/] deleted" ) else: console.print( - f"Deleted {len(presets)} endpoint preset{'s' if len(presets) != 1 else ''} " + f"Deleted {len(presets)} preset{'s' if len(presets) != 1 else ''} " f"for [code]{args.base or args.repo}[/]" ) @@ -277,13 +276,13 @@ 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", ) @@ -345,9 +344,11 @@ def _apply_name(configuration: EndpointConfiguration, name: str | None) -> None: if name is not None: configuration.name = name if configuration.name is None: - raise CLIError("Endpoint name is required. Set `name` in the configuration or use --name") + raise CLIError( + "The service name is required. Set `name` in the configuration or use --name" + ) if not is_valid_dstack_resource_name(configuration.name): - raise CLIError("Endpoint name must match '^[a-z][a-z0-9-]{1,40}$'") + raise CLIError("The name must match '^[a-z][a-z0-9-]{1,40}$'") def _get_effective_configuration( diff --git a/src/dstack/_internal/cli/models/endpoint_presets.py b/src/dstack/_internal/cli/models/endpoint_presets.py index aea33f499..7f1058d6a 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: diff --git a/src/dstack/_internal/cli/models/endpoints.py b/src/dstack/_internal/cli/models/endpoints.py index 29335d883..f2e00c80c 100644 --- a/src/dstack/_internal/cli/models/endpoints.py +++ b/src/dstack/_internal/cli/models/endpoints.py @@ -84,10 +84,10 @@ 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, @@ -115,9 +115,6 @@ class EndpointConfiguration( context_length: Annotated[ Optional[PositiveInt], Field(description="The minimum required context length") ] = None - preset: Annotated[ - Optional[str], Field(description="The preset ID to use when applying the endpoint") - ] = None max_trials: Annotated[ Optional[PositiveInt], Field( @@ -180,12 +177,6 @@ def parse_model(cls, value: Any) -> Any: return {"repo": _validate_model(value, field="model")} return value - @validator("preset") - def validate_preset(cls, value: Optional[str]) -> Optional[str]: - if value is not None and not value.strip(): - raise ValueError("Endpoint preset must be a non-empty string") - return value - class EndpointPresetConstraints(CoreModel): """The effective constraints for endpoint preset creation, saved as `constraints.json` diff --git a/src/dstack/_internal/cli/services/endpoints/agent.py b/src/dstack/_internal/cli/services/endpoints/agent.py index c2cae73f4..a987fe347 100644 --- a/src/dstack/_internal/cli/services/endpoints/agent.py +++ b/src/dstack/_internal/cli/services/endpoints/agent.py @@ -298,7 +298,7 @@ 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_presets_dir() path: Optional[Path] = None @@ -332,7 +332,7 @@ def create_endpoint_agent_session( else: data.pop("env", None) _write_private_text( - path / "endpoint.dstack.yml", + path / "preset.dstack.yml", yaml.safe_dump(data, sort_keys=False), ) _write_private_text(path / "trace.jsonl", "") @@ -414,7 +414,7 @@ def list_agent_sessions() -> list[dict[str, Any]]: 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"): + if status not in ("running", "interrupted", "success"): continue pid = manifest.get("pid") if status == "running" and not ( @@ -436,7 +436,6 @@ def _summarize_session_trials(path: Path) -> Optional[dict[str, Any]]: except OSError: lines = [] count = 0 - task_names: set[str] = set() best: Optional[dict[str, Any]] = None for line in lines: try: @@ -445,14 +444,9 @@ def _summarize_session_trials(path: Path) -> Optional[dict[str, Any]]: continue if not isinstance(record, dict): continue - # A trial may log several benchmark records (e.g. a re-run); records - # sharing a task name are one trial. - task = record.get("task") - task_name = task.get("name") if isinstance(task, dict) else None - if isinstance(task_name, str) and task_name: - task_names.add(task_name) - else: - count += 1 + # 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 @@ -480,7 +474,7 @@ def _summarize_session_trials(path: Path) -> Optional[dict[str, Any]]: "concurrency": workload.get("concurrency"), "gpu": gpu_text, } - return {"count": count + len(task_names), "best": best} + return {"count": count, "best": best} def get_claude_auth() -> ClaudeAuth: @@ -792,7 +786,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: @@ -859,7 +853,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) diff --git a/src/dstack/_internal/cli/services/endpoints/apply.py b/src/dstack/_internal/cli/services/endpoints/apply.py index 350a47dd9..23dc25d4b 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 @@ -32,21 +32,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 +67,32 @@ 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_id = {preset.id: preset for preset in presets} + missing = [preset_id for preset_id in preset_ids if preset_id not in presets_by_id] + 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 ids in turn. + return [presets_by_id[preset_id] for preset_id in preset_ids] + + 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: diff --git a/src/dstack/_internal/cli/services/endpoints/create.py b/src/dstack/_internal/cli/services/endpoints/create.py index 89b9182ae..892cb79ff 100644 --- a/src/dstack/_internal/cli/services/endpoints/create.py +++ b/src/dstack/_internal/cli/services/endpoints/create.py @@ -185,10 +185,10 @@ async def _create_endpoint_preset( preset_id=agent_session.preset_id or None, ) if contains_redacted_value(endpoint_preset_to_data(preset), redacted_values): - raise CLIError("Generated endpoint preset contains a secret value") + raise CLIError("Generated 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}.", + f"Saved preset {preset.id} for {preset.base} at {preset_path}.", agent_session=agent_session, ) creation_succeeded = True @@ -274,14 +274,16 @@ def _suspend_agent_session(session: EndpointAgentSession) -> None: "Its runs may still be active and accruing cost." ) console.print( - f"Resume with [code]dstack endpoint preset create -f " + f"Resume with [code]dstack preset create -f " f"--resume {session.preset_id}[/], or stop the runs with [code]dstack stop[/]." ) def _get_build_name(endpoint_name: Optional[str], suffix: str) -> str: if endpoint_name is None: - raise CLIError("Endpoint name is required. Set `name` in the configuration or use --name") + raise CLIError( + "The service name is required. Set `name` in the configuration or use --name" + ) # Leave room for the preset id and numeric submission suffix while retaining # a recognizable prefix. prefix = endpoint_name[:26].rstrip("-") diff --git a/src/dstack/_internal/cli/services/endpoints/output.py b/src/dstack/_internal/cli/services/endpoints/output.py index db24f8073..a57023dca 100644 --- a/src/dstack/_internal/cli/services/endpoints/output.py +++ b/src/dstack/_internal/cli/services/endpoints/output.py @@ -12,7 +12,7 @@ from dstack._internal.utils.common import pretty_date, pretty_resources _STATUS_DISPLAY = { - "ready": ("done", "grey"), + "ready": ("success", "green"), "running": ("clauding", "bold deep_sky_blue1"), "interrupted": ("interrupted", "bold gold1"), "failed": ("failed", "indian_red1"), @@ -24,6 +24,21 @@ def _format_status(status: str) -> str: return f"[{style}]{text}[/]" if style else text +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, @@ -44,17 +59,22 @@ def print_endpoint_presets( presets_by_base[preset.base].append(preset) repo_to_base[preset.model] = preset.base sessions_by_model: dict[str, list[dict[str, Any]]] = defaultdict(list) + creations_by_id: dict[str, dict[str, Any]] = {} for session in sessions or []: + if str(session.get("status")) == "success": + # A completed creation session decorates its preset row. + creations_by_id[str(session.get("id"))] = session + continue model = str(session.get("model") or "unknown") sessions_by_model[repo_to_base.get(model, model)].append(session) for base in sorted({*presets_by_base, *sessions_by_model}, key=str.lower): - add_row_from_dict(table, {"BASE": f"[bold]{base}[/]"}) + add_row_from_dict(table, {"BASE": f"[secondary]{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) + _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 ""), @@ -65,7 +85,7 @@ def print_endpoint_presets( console.print() -def _add_session(table: Table, session: dict[str, Any]) -> None: +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): @@ -75,16 +95,9 @@ def _add_session(table: Table, session: dict[str, Any]) -> None: created = created_at benchmark = "" gpu = "" - status = _format_status(str(session.get("status", ""))) + status = _format_status(str(session.get("status", ""))) + _format_trial_progress(session) trials = session.get("trials") - max_trials = session.get("max_trials") - if isinstance(trials, dict) and (trials.get("count") or isinstance(max_trials, int)): - progress = str(trials.get("count") or 0) - if isinstance(max_trials, int): - progress += f"/{max_trials}" - # The trial progress stays outside the status markup to render in the - # default color. - status += f" ({progress})" + if isinstance(trials, dict): best = trials.get("best") if isinstance(best, dict): parts = ["best trial:"] @@ -106,13 +119,19 @@ def _add_session(table: Table, session: dict[str, Any]) -> None: ) -def _add_preset(table: Table, preset: EndpointPreset, *, verbose: bool) -> None: +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, column: _format_resources(groups[0].resources, verbose=verbose), - "STATUS": _format_status("ready"), + "STATUS": _format_status("ready") + _format_trial_progress(creation), "BENCHMARK": format_endpoint_benchmark(preset, verbose=verbose), "SUBMITTED": pretty_date(preset.created_at), } @@ -122,7 +141,7 @@ def _add_preset(table: Table, preset: EndpointPreset, *, verbose: bool) -> None: if verbose and preset.model != preset.base: add_row_from_dict( table, - {"BASE": f" repo={preset.model}"}, + {"BASE": f" repo={preset.model}"}, style="secondary", ) if len(groups) > 1: @@ -130,7 +149,7 @@ def _add_preset(table: Table, preset: EndpointPreset, *, verbose: bool) -> None: add_row_from_dict( table, { - "BASE": f" group={group.name}", + "BASE": f" group={group.name}", column: _format_resources(group.resources, verbose=verbose), }, style="secondary", diff --git a/src/dstack/_internal/cli/services/endpoints/presets.py b/src/dstack/_internal/cli/services/endpoints/presets.py index 9f0b44e52..922d21a7e 100644 --- a/src/dstack/_internal/cli/services/endpoints/presets.py +++ b/src/dstack/_internal/cli/services/endpoints/presets.py @@ -116,7 +116,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/store.py b/src/dstack/_internal/cli/services/endpoints/store.py index efe8e791c..f28289c4a 100644 --- a/src/dstack/_internal/cli/services/endpoints/store.py +++ b/src/dstack/_internal/cli/services/endpoints/store.py @@ -42,7 +42,7 @@ def get(self, preset_id: str) -> EndpointPreset | None: return None 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: @@ -105,12 +105,12 @@ def _load(self, path: Path) -> EndpointPreset: 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 _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 endpoint preset ID: {preset_id!r}") + raise CLIError(f"Invalid preset ID: {preset_id!r}") def load_endpoint_configuration(path: str) -> tuple[str, EndpointConfiguration]: @@ -131,12 +131,12 @@ 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") + 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" diff --git a/src/dstack/_internal/cli/services/endpoints/verify.py b/src/dstack/_internal/cli/services/endpoints/verify.py index 9886c1b76..33a9a9754 100644 --- a/src/dstack/_internal/cli/services/endpoints/verify.py +++ b/src/dstack/_internal/cli/services/endpoints/verify.py @@ -85,14 +85,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 ( diff --git a/src/tests/_internal/cli/commands/test_endpoint.py b/src/tests/_internal/cli/commands/test_endpoint.py index 0b25f9a08..43d8eb0d1 100644 --- a/src/tests/_internal/cli/commands/test_endpoint.py +++ b/src/tests/_internal/cli/commands/test_endpoint.py @@ -73,7 +73,7 @@ def test_prints_submitted_column(self, monkeypatch): 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 ( @@ -84,7 +84,7 @@ def test_handles_keyboard_interrupt(self, tmp_path, capsys): ), ): exit_code = run_dstack_cli( - ["endpoint", "preset", "create", "-f", str(configuration_path)], + ["preset", "create", "-f", str(configuration_path)], home_dir=tmp_path, ) @@ -112,7 +112,7 @@ 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) @@ -132,7 +132,7 @@ def test_lists_and_deletes_preset_without_api_client(self, tmp_path, capsys): assert "108ms" in output assert "A6000:48GB:1" not in output - assert run_dstack_cli(["endpoint", "preset", "list", "-v"], home_dir=tmp_path) == 0 + assert run_dstack_cli(["preset", "list", "-v"], home_dir=tmp_path) == 0 verbose_output = capsys.readouterr().out joined_verbose = "".join(verbose_output.split()) assert "hardware=" in joined_verbose @@ -145,7 +145,6 @@ def test_lists_and_deletes_preset_without_api_client(self, tmp_path, capsys): assert ( run_dstack_cli( [ - "endpoint", "preset", "delete", preset.id, @@ -167,7 +166,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 @@ -183,8 +182,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): @@ -212,7 +211,7 @@ def test_deletes_presets_by_model_without_api_client(self, tmp_path, flag_attrib with patch("dstack.api.Client.from_config") as from_config: assert ( run_dstack_cli( - ["endpoint", "preset", "delete", flag, getattr(preset, attribute), "-y"], + ["preset", "delete", flag, getattr(preset, attribute), "-y"], home_dir=tmp_path, ) == 0 @@ -232,7 +231,7 @@ def test_delete_by_model_keeps_other_presets(self, tmp_path): assert ( run_dstack_cli( - ["endpoint", "preset", "delete", "--base", preset.base, "-y"], + ["preset", "delete", "--base", preset.base, "-y"], home_dir=tmp_path, ) == 0 @@ -250,7 +249,7 @@ def test_lists_presets_filtered_by_model(self, tmp_path, capsys, flag_attribute) preset.copy(update={"id": "01234567", "base": "meta/Llama-4", "model": "meta/Llama-4"}) ) - args = ["endpoint", "preset", "list", "--json", flag, getattr(preset, attribute)] + args = ["preset", "list", "--json", flag, getattr(preset, attribute)] assert run_dstack_cli(args, home_dir=tmp_path) == 0 output = json.loads(capsys.readouterr().out) @@ -270,7 +269,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 @@ -296,7 +295,6 @@ def test_merges_profile_configuration_and_cli_args(self, tmp_path): ): exit_code = run_dstack_cli( [ - "endpoint", "preset", "create", "-f", @@ -326,17 +324,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 ( @@ -344,7 +348,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", @@ -361,5 +364,5 @@ 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 diff --git a/src/tests/_internal/cli/services/endpoints/test_agent.py b/src/tests/_internal/cli/services/endpoints/test_agent.py index 7611ed3df..5d1ba4809 100644 --- a/src/tests/_internal/cli/services/endpoints/test_agent.py +++ b/src/tests/_internal/cli/services/endpoints/test_agent.py @@ -202,17 +202,17 @@ 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() + 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" @@ -860,7 +860,7 @@ def test_refusals(self, tmp_path, monkeypatch): class TestSummarizeSessionTrials: - def test_counts_distinct_task_names_not_benchmark_records(self, tmp_path): + 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}}, @@ -881,5 +881,7 @@ def test_counts_distinct_task_names_not_benchmark_records(self, tmp_path): summary = _summarize_session_trials(path) - assert summary["count"] == 3 + # 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"} diff --git a/src/tests/_internal/cli/services/endpoints/test_apply.py b/src/tests/_internal/cli/services/endpoints/test_apply.py index 3f629632e..e13b8c8ba 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])), diff --git a/src/tests/_internal/cli/services/endpoints/test_create.py b/src/tests/_internal/cli/services/endpoints/test_create.py index 975f5c68a..7a5ec002e 100644 --- a/src/tests/_internal/cli/services/endpoints/test_create.py +++ b/src/tests/_internal/cli/services/endpoints/test_create.py @@ -183,14 +183,14 @@ def fail_finish(self, preset_id=None): assert len(paths) == 1 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 json.loads((paths[0] / "session.json").read_text())["status"] == "running" - assert "hf-secret" not in (paths[0] / "endpoint.dstack.yml").read_text() + 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): @@ -321,7 +321,7 @@ async def run_agent(**kwargs): class TestBuildName: def test_requires_name_and_keeps_generated_prefix_bounded(self): - with pytest.raises(CLIError, match="Endpoint name is required"): + with pytest.raises(CLIError, match="The service name is required"): _get_build_name(None, "a1b2c3d4") build_name = _get_build_name("qwen-endpoint-with-a-name-that-is-forty-one", "a1b2c3d4") diff --git a/src/tests/_internal/cli/services/endpoints/test_output.py b/src/tests/_internal/cli/services/endpoints/test_output.py index bc00817a7..311d7b3ab 100644 --- a/src/tests/_internal/cli/services/endpoints/test_output.py +++ b/src/tests/_internal/cli/services/endpoints/test_output.py @@ -1,7 +1,14 @@ +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 @@ -41,7 +48,7 @@ def test_shows_progress_after_status_and_best_benchmark(self): } ) - assert row["STATUS"] == "[bold deep_sky_blue1]clauding[/] (2/3)" + 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" @@ -50,7 +57,7 @@ def test_shows_zero_progress_without_benchmark(self): {"id": "ab12cd34", "status": "running", "max_trials": 3, "trials": {"count": 0}} ) - assert row["STATUS"] == "[bold deep_sky_blue1]clauding[/] (0/3)" + assert row["STATUS"] == "[bold deep_sky_blue1]clauding[/] [secondary](0/3)[/]" assert row["BENCHMARK"] == "" def test_omits_progress_without_trials_data(self): @@ -61,20 +68,11 @@ def test_omits_progress_without_trials_data(self): def test_counts_without_max_trials(self): row = _session_row({"id": "ab12cd34", "status": "interrupted", "trials": {"count": 2}}) - assert row["STATUS"] == "[bold gold1]interrupted[/] (2)" + assert row["STATUS"] == "[bold gold1]interrupted[/] [secondary](2)[/]" class TestGroupOrdering: def test_sorts_presets_and_sessions_newest_first(self, monkeypatch): - from datetime import timedelta - from io import StringIO - - from rich.console import Console - from rich.theme import Theme - - from dstack._internal.cli.services.endpoints import output as output_module - from tests._internal.cli.endpoint_presets import get_endpoint_preset - buffer = StringIO() monkeypatch.setattr( output_module, @@ -105,3 +103,31 @@ def test_sorts_presets_and_sessions_newest_first(self, monkeypatch): 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 diff --git a/src/tests/_internal/cli/services/endpoints/test_store.py b/src/tests/_internal/cli/services/endpoints/test_store.py index eba675101..5cbc7e7c3 100644 --- a/src/tests/_internal/cli/services/endpoints/test_store.py +++ b/src/tests/_internal/cli/services/endpoints/test_store.py @@ -156,7 +156,7 @@ def test_preserves_literal_env_values(self, tmp_path: Path): class TestParseEndpointConfiguration: @pytest.mark.parametrize("key", ["base", "repo"]) def test_warns_on_nested_model_without_name(self, key: str): - stream = StringIO(f"type: endpoint\nmodel:\n {key}: Qwen/Qwen3.5-27B\n") + 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) @@ -176,7 +176,7 @@ def test_warns_on_nested_model_without_name(self, key: str): ], ) def test_does_not_warn_on_preferred_syntax(self, body: str): - stream = StringIO(f"type: endpoint\n{body}") + stream = StringIO(f"type: preset\n{body}") with patch.object(store_module, "warn") as warn: configuration = store_module._parse_endpoint_configuration(stream) diff --git a/src/tests/_internal/cli/services/endpoints/test_verify.py b/src/tests/_internal/cli/services/endpoints/test_verify.py index aa55f2ae9..ee598694b 100644 --- a/src/tests/_internal/cli/services/endpoints/test_verify.py +++ b/src/tests/_internal/cli/services/endpoints/test_verify.py @@ -119,3 +119,19 @@ def test_still_rejects_unknown_bearer_token(self, tmp_path): 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 From 41ac512de42a8eec2e60c4fa52e88222b8af570c Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Tue, 21 Jul 2026 12:05:51 -0700 Subject: [PATCH 06/25] Unify the benchmark format across list, list -v, and apply - BENCHMARK renders identically with and without -v; the verbose-only detail dump is gone (full data remains in `preset get --json`) - -v adds exactly two things: a `ctx=` benchmark prefix (replacing the CONTEXT column) and a dimmed indented `repo=` on the preset's own row - `apply` shows the selected preset with the same `ctx= con= tok/s TTFT` string Co-Authored-By: Claude Fable 5 --- .../_internal/cli/services/endpoints/apply.py | 5 +- .../cli/services/endpoints/output.py | 58 ++----------------- .../_internal/cli/commands/test_endpoint.py | 18 +++--- .../cli/services/endpoints/test_apply.py | 2 +- 4 files changed, 16 insertions(+), 67 deletions(-) diff --git a/src/dstack/_internal/cli/services/endpoints/apply.py b/src/dstack/_internal/cli/services/endpoints/apply.py index 23dc25d4b..baef1a21b 100644 --- a/src/dstack/_internal/cli/services/endpoints/apply.py +++ b/src/dstack/_internal/cli/services/endpoints/apply.py @@ -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 @@ -162,7 +161,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/output.py b/src/dstack/_internal/cli/services/endpoints/output.py index a57023dca..71a4c897f 100644 --- a/src/dstack/_internal/cli/services/endpoints/output.py +++ b/src/dstack/_internal/cli/services/endpoints/output.py @@ -6,7 +6,6 @@ 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 @@ -48,8 +47,6 @@ def print_endpoint_presets( table.add_column("BASE", no_wrap=True) table.add_column("ID", no_wrap=True) table.add_column("RESOURCES" if verbose else "GPU", style="secondary") - if verbose: - table.add_column("CONTEXT", justify="right", 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") @@ -69,7 +66,7 @@ def print_endpoint_presets( 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": f"[secondary]{base}[/]"}) + 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 @@ -135,15 +132,9 @@ def _add_preset( "BENCHMARK": format_endpoint_benchmark(preset, verbose=verbose), "SUBMITTED": pretty_date(preset.created_at), } - if verbose: - row["CONTEXT"] = format_endpoint_context_length(preset) - add_row_from_dict(table, row) if verbose and preset.model != preset.base: - add_row_from_dict( - table, - {"BASE": f" repo={preset.model}"}, - style="secondary", - ) + 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( @@ -156,16 +147,10 @@ def _add_preset( ) -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"con={workload.concurrency}", @@ -173,38 +158,10 @@ def format_endpoint_benchmark(preset: EndpointPreset, *, verbose: bool = False) 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: @@ -225,11 +182,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/tests/_internal/cli/commands/test_endpoint.py b/src/tests/_internal/cli/commands/test_endpoint.py index 43d8eb0d1..28a339c0b 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): @@ -103,7 +103,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"}), ), @@ -115,8 +115,10 @@ def test_lists_and_deletes_preset_without_api_client(self, tmp_path, capsys): 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 "8f3a12c4" in output # The repo row is shown only in verbose mode. @@ -132,14 +134,12 @@ def test_lists_and_deletes_preset_without_api_client(self, tmp_path, capsys): assert "108ms" in output assert "A6000:48GB:1" not in output - assert run_dstack_cli(["preset", "list", "-v"], home_dir=tmp_path) == 0 - verbose_output = capsys.readouterr().out joined_verbose = "".join(verbose_output.split()) - assert "hardware=" in joined_verbose - assert "api=" in joined_verbose - assert "n=16" in joined_verbose + # 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 "1K->128" in joined_verbose + assert "hardware=" not in joined_verbose with patch("dstack.api.Client.from_config") as from_config: assert ( diff --git a/src/tests/_internal/cli/services/endpoints/test_apply.py b/src/tests/_internal/cli/services/endpoints/test_apply.py index e13b8c8ba..fab87e900 100644 --- a/src/tests/_internal/cli/services/endpoints/test_apply.py +++ b/src/tests/_internal/cli/services/endpoints/test_apply.py @@ -150,7 +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, con=1 42.1 tok/s TTFT 108ms[/])", + "Preset": "8f3a12c4 ([secondary]ctx=32K con=1 42.1 tok/s TTFT 108ms[/])", }, ) From c333889dce93c351f4839426e92ef67643ff7f16 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Tue, 21 Jul 2026 13:25:16 -0700 Subject: [PATCH 07/25] Add the preset `prompt`: user instructions for the creation agent - `prompt:` in the configuration, inline or as a file `path:` relative to the configuration file; validated and capped - The agent contract gains the text via template directives in system_prompt.md (``); without a prompt the rendered contract stays byte-identical. The default no-patching limits gain an escape clause the prompt can invoke explicitly - The resolved prompt is pinned per session and kept across resumes, with a warning when the configuration changes meanwhile - The list shows `verifying` once a running session's trial budget is spent and the final service is being deployed and verified Co-Authored-By: Claude Fable 5 --- src/dstack/_internal/cli/commands/endpoint.py | 5 +- src/dstack/_internal/cli/models/endpoints.py | 32 ++++++++ .../_internal/cli/services/endpoints/agent.py | 11 +++ .../cli/services/endpoints/create.py | 14 +++- .../cli/services/endpoints/output.py | 19 ++++- .../cli/services/endpoints/prompt.py | 31 ++++++- .../endpoints/resources/system_prompt.md | 9 ++- .../_internal/cli/services/endpoints/store.py | 28 ++++++- .../cli/services/endpoints/test_create.py | 80 +++++++++++++++++++ .../cli/services/endpoints/test_output.py | 16 ++++ .../cli/services/endpoints/test_prompt.py | 46 +++++++++++ .../cli/services/endpoints/test_store.py | 37 ++++++++- 12 files changed, 320 insertions(+), 8 deletions(-) create mode 100644 src/tests/_internal/cli/services/endpoints/test_prompt.py diff --git a/src/dstack/_internal/cli/commands/endpoint.py b/src/dstack/_internal/cli/commands/endpoint.py index 93c6f158e..a6a6d870b 100644 --- a/src/dstack/_internal/cli/commands/endpoint.py +++ b/src/dstack/_internal/cli/commands/endpoint.py @@ -22,6 +22,7 @@ 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 @@ -191,8 +192,9 @@ def _list(self, args: argparse.Namespace) -> None: time.sleep(5) def _create(self, args: argparse.Namespace) -> None: - _, configuration = load_endpoint_configuration(args.configuration_file) + configuration_path, configuration = load_endpoint_configuration(args.configuration_file) configuration = _get_effective_configuration(configuration, args) + user_prompt = resolve_endpoint_prompt(configuration, configuration_path) resume_session = None if getattr(args, "resume", None): resume_session = load_resumable_agent_session(args.resume) @@ -208,6 +210,7 @@ def _create(self, args: argparse.Namespace) -> None: keep_service=args.keep_service, debug=args.debug, resume_session=resume_session, + user_prompt=user_prompt, ) console.print( f"Preset [code]{result.preset.id}[/] for " diff --git a/src/dstack/_internal/cli/models/endpoints.py b/src/dstack/_internal/cli/models/endpoints.py index f2e00c80c..03762c41b 100644 --- a/src/dstack/_internal/cli/models/endpoints.py +++ b/src/dstack/_internal/cli/models/endpoints.py @@ -69,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 @@ -112,6 +127,14 @@ class EndpointConfiguration( 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 @@ -177,6 +200,15 @@ def parse_model(cls, value: Any) -> Any: return {"repo": _validate_model(value, field="model")} return value + @validator("prompt") + def validate_prompt(cls, value: Any) -> Any: + if isinstance(value, str): + if not value.strip(): + raise ValueError("Prompt must be a non-empty string") + if len(value) > MAX_PROMPT_LENGTH: + raise ValueError(f"Prompt must be at most {MAX_PROMPT_LENGTH} characters") + return value + class EndpointPresetConstraints(CoreModel): """The effective constraints for endpoint preset creation, saved as `constraints.json` diff --git a/src/dstack/_internal/cli/services/endpoints/agent.py b/src/dstack/_internal/cli/services/endpoints/agent.py index a987fe347..1e2f3deb7 100644 --- a/src/dstack/_internal/cli/services/endpoints/agent.py +++ b/src/dstack/_internal/cli/services/endpoints/agent.py @@ -36,6 +36,7 @@ _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" @@ -147,6 +148,16 @@ def trials_path(self) -> Path: 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) diff --git a/src/dstack/_internal/cli/services/endpoints/create.py b/src/dstack/_internal/cli/services/endpoints/create.py index 892cb79ff..56d527682 100644 --- a/src/dstack/_internal/cli/services/endpoints/create.py +++ b/src/dstack/_internal/cli/services/endpoints/create.py @@ -63,6 +63,7 @@ def create_endpoint_preset( build_name: Optional[str] = None, debug: bool = False, resume_session: Optional[EndpointAgentSession] = None, + user_prompt: Optional[str] = None, ) -> EndpointPresetCreateResult: agent_session = resume_session or create_endpoint_agent_session(configuration, debug=debug) try: @@ -77,6 +78,7 @@ def create_endpoint_preset( build_name=build_name, agent_session=agent_session, resume=resume_session is not None, + user_prompt=user_prompt, ) ) except KeyboardInterrupt: @@ -101,10 +103,18 @@ async def _create_endpoint_preset( build_name: Optional[str] = None, agent_session: EndpointAgentSession, resume: bool = False, + user_prompt: Optional[str] = None, ) -> EndpointPresetCreateResult: source_configuration = source_configuration or configuration initial_resume_session_id: Optional[str] = None if 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() @@ -150,8 +160,10 @@ async def _create_endpoint_preset( workspace=workspace, token=token, ) - prompt = get_endpoint_agent_system_prompt() + prompt = get_endpoint_agent_system_prompt(user_prompt=user_prompt) if not resume: + if user_prompt: + agent_session.write_user_prompt(user_prompt) constraints_text = _build_constraints( configuration=configuration, build_name=build_name, diff --git a/src/dstack/_internal/cli/services/endpoints/output.py b/src/dstack/_internal/cli/services/endpoints/output.py index 71a4c897f..f2e2393d2 100644 --- a/src/dstack/_internal/cli/services/endpoints/output.py +++ b/src/dstack/_internal/cli/services/endpoints/output.py @@ -13,6 +13,7 @@ _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"), } @@ -23,6 +24,17 @@ def _format_status(status: str) -> str: 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.""" @@ -92,7 +104,12 @@ def _add_session(table: Table, session: dict[str, Any], base_label: Optional[str created = created_at benchmark = "" gpu = "" - status = _format_status(str(session.get("status", ""))) + _format_trial_progress(session) + 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") diff --git a/src/dstack/_internal/cli/services/endpoints/prompt.py b/src/dstack/_internal/cli/services/endpoints/prompt.py index 4d3fbd628..0123047ee 100644 --- a/src/dstack/_internal/cli/services/endpoints/prompt.py +++ b/src/dstack/_internal/cli/services/endpoints/prompt.py @@ -1,9 +1,36 @@ +import re from pathlib import Path +from typing import Optional + +from dstack._internal.core.errors import CLIError _SYSTEM_PROMPT_PATH = Path(__file__).resolve().parent / "resources" / "system_prompt.md" +# `` emits CONTENT (with `{NAME}` interpolated) when the +# variable NAME is set, and nothing otherwise. The conditional text lives in +# the document; this module only applies the rule. +_DIRECTIVE_PATTERN = re.compile(r"", re.DOTALL) + # TODO: reintroduce a `# Resume` section in system_prompt.md once session resume # (seeded from `runs.jsonl` and `trials.jsonl`) is designed. -def get_endpoint_agent_system_prompt() -> str: - return _SYSTEM_PROMPT_PATH.read_text(encoding="utf-8").strip() +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 2f0f2c08e..ab62a35b4 100644 --- a/src/dstack/_internal/cli/services/endpoints/resources/system_prompt.md +++ b/src/dstack/_internal/cli/services/endpoints/resources/system_prompt.md @@ -35,8 +35,15 @@ 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. +patching drivers, or patching serving framework source code. + ## CLI And Skills All trials and the final verification are done using `dstack`. This includes diff --git a/src/dstack/_internal/cli/services/endpoints/store.py b/src/dstack/_internal/cli/services/endpoints/store.py index f28289c4a..24e234e2d 100644 --- a/src/dstack/_internal/cli/services/endpoints/store.py +++ b/src/dstack/_internal/cli/services/endpoints/store.py @@ -10,7 +10,11 @@ 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 @@ -145,3 +149,25 @@ def _parse_endpoint_configuration(stream: TextIO) -> EndpointConfiguration: 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/tests/_internal/cli/services/endpoints/test_create.py b/src/tests/_internal/cli/services/endpoints/test_create.py index 7a5ec002e..590ba5d7b 100644 --- a/src/tests/_internal/cli/services/endpoints/test_create.py +++ b/src/tests/_internal/cli/services/endpoints/test_create.py @@ -562,3 +562,83 @@ async def run_agent(**kwargs): 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) diff --git a/src/tests/_internal/cli/services/endpoints/test_output.py b/src/tests/_internal/cli/services/endpoints/test_output.py index 311d7b3ab..6bb188736 100644 --- a/src/tests/_internal/cli/services/endpoints/test_output.py +++ b/src/tests/_internal/cli/services/endpoints/test_output.py @@ -131,3 +131,19 @@ def test_completed_creation_decorates_preset_row_without_extra_session_row(self, 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 5cbc7e7c3..1ceb731f7 100644 --- a/src/tests/_internal/cli/services/endpoints/test_store.py +++ b/src/tests/_internal/cli/services/endpoints/test_store.py @@ -13,9 +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, @@ -183,3 +184,37 @@ def test_does_not_warn_on_preferred_syntax(self, body: str): 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, + ) From 15807f905d7706d4322f17f1a466c11a719f5433 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Tue, 21 Jul 2026 14:51:45 -0700 Subject: [PATCH 08/25] Preset names as mutable pointers; apply-style create plan and confirmation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - The configuration `name` names the preset, claimed at session start: creating with a used name detaches it from the current holder (preset or session of any status) after confirmation, docker-tag style; the old preset survives nameless. `get`, `delete`, and `apply --id` accept names wherever they accept ids; the list gains a NAME column - Name-less creation is allowed; run names derive from a model slug - `create` shows an apply-style plan before starting the agent — Project, User, the effective fleets, and their offers (agent-free, best-effort) — and always asks for confirmation; `-y` skips - `print_offers` extracted from `print_run_plan` for reuse Co-Authored-By: Claude Fable 5 --- src/dstack/_internal/cli/commands/endpoint.py | 84 +++++++++++--- .../_internal/cli/models/endpoint_presets.py | 2 + .../_internal/cli/services/endpoints/agent.py | 26 +++++ .../_internal/cli/services/endpoints/apply.py | 16 ++- .../cli/services/endpoints/create.py | 70 ++++++++++-- .../cli/services/endpoints/output.py | 3 + .../cli/services/endpoints/presets.py | 3 + .../_internal/cli/services/endpoints/store.py | 15 +++ .../cli/services/endpoints/verify.py | 2 + src/dstack/_internal/cli/utils/run.py | 26 ++++- .../_internal/cli/commands/test_endpoint.py | 105 +++++++++++++++++- .../cli/services/endpoints/test_create.py | 9 +- .../cli/services/endpoints/test_store.py | 17 +++ 13 files changed, 342 insertions(+), 36 deletions(-) diff --git a/src/dstack/_internal/cli/commands/endpoint.py b/src/dstack/_internal/cli/commands/endpoint.py index a6a6d870b..0fc51a5b2 100644 --- a/src/dstack/_internal/cli/commands/endpoint.py +++ b/src/dstack/_internal/cli/commands/endpoint.py @@ -13,11 +13,15 @@ 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, 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 ( + create_endpoint_preset, + plan_endpoint_preset, +) from dstack._internal.cli.services.endpoints.output import print_endpoint_presets from dstack._internal.cli.services.endpoints.store import ( EndpointPresetStore, @@ -61,7 +65,7 @@ def _register(self) -> None: 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", @@ -98,6 +102,9 @@ def _register(self) -> None: 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) apply_parser = preset_subparsers.add_parser( @@ -112,7 +119,7 @@ def _register(self) -> None: action="append", dest="preset_ids", metavar="ID", - help="Deploy the best available preset among these IDs. Can be repeated", + 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" @@ -138,7 +145,7 @@ def _register(self) -> None: "preset", nargs="?", metavar="ID", - help="The preset ID", + help="The preset ID or name", ) delete_target.add_argument( "--base", @@ -193,8 +200,9 @@ def _list(self, args: argparse.Namespace) -> None: def _create(self, args: argparse.Namespace) -> None: configuration_path, configuration = load_endpoint_configuration(args.configuration_file) - configuration = _get_effective_configuration(configuration, args) + 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) @@ -203,14 +211,22 @@ def _create(self, args: argparse.Namespace) -> None: "[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"Preset [code]{result.preset.id}[/] for " @@ -220,7 +236,8 @@ def _create(self, args: argparse.Namespace) -> None: console.print(f"Final service [code]{result.final_run_name}[/] kept running") 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"Preset {args.preset!r} does not exist") print(preset.json()) @@ -241,7 +258,7 @@ def _apply(self, args: argparse.Namespace) -> None: def _delete(self, args: argparse.Namespace) -> None: store = EndpointPresetStore() 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"Preset {args.preset!r} does not exist") presets = [preset] @@ -343,22 +360,61 @@ def _session_matches_model( 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( - "The service 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("The name must match '^[a-z][a-z0-9-]{1,40}$'") +def _confirm_preset_creation( + store: EndpointPresetStore, name: str | None, *, assume_yes: bool +) -> bool: + """One apply-style confirmation; detaches 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" Detach it and create 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.detach_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) diff --git a/src/dstack/_internal/cli/models/endpoint_presets.py b/src/dstack/_internal/cli/models/endpoint_presets.py index 7f1058d6a..75aa778e8 100644 --- a/src/dstack/_internal/cli/models/endpoint_presets.py +++ b/src/dstack/_internal/cli/models/endpoint_presets.py @@ -104,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/services/endpoints/agent.py b/src/dstack/_internal/cli/services/endpoints/agent.py index 1e2f3deb7..d86853868 100644 --- a/src/dstack/_internal/cli/services/endpoints/agent.py +++ b/src/dstack/_internal/cli/services/endpoints/agent.py @@ -329,6 +329,7 @@ def create_endpoint_agent_session( "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, @@ -415,6 +416,30 @@ def load_resumable_agent_session(preset_id: str) -> EndpointAgentSession: 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 = [] @@ -434,6 +459,7 @@ def list_agent_sessions() -> list[dict[str, Any]]: 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) diff --git a/src/dstack/_internal/cli/services/endpoints/apply.py b/src/dstack/_internal/cli/services/endpoints/apply.py index baef1a21b..8aed115b7 100644 --- a/src/dstack/_internal/cli/services/endpoints/apply.py +++ b/src/dstack/_internal/cli/services/endpoints/apply.py @@ -73,14 +73,22 @@ def _get_candidate_presets( ) -> list[EndpointPreset]: if not preset_ids: return presets - presets_by_id = {preset.id: preset for preset in presets} - missing = [preset_id for preset_id in preset_ids if preset_id not in presets_by_id] + 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 ids in turn. - return [presets_by_id[preset_id] for preset_id in preset_ids] + # 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( diff --git a/src/dstack/_internal/cli/services/endpoints/create.py b/src/dstack/_internal/cli/services/endpoints/create.py index 56d527682..d9db24f2d 100644 --- a/src/dstack/_internal/cli/services/endpoints/create.py +++ b/src/dstack/_internal/cli/services/endpoints/create.py @@ -2,12 +2,15 @@ import dataclasses import json import os +import re import uuid from contextlib import suppress from dataclasses import dataclass from pathlib import Path from typing import Optional, Sequence +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 ( @@ -19,6 +22,7 @@ EndpointAgentWorkspace, attach_agent_workspace, build_endpoint_agent_env, + claimed_session_name, contains_redacted_value, create_agent_workspace, create_endpoint_agent_session, @@ -38,9 +42,12 @@ load_endpoint_agent_report, ) from dstack._internal.cli.utils.common import console, warn +from dstack._internal.cli.utils.run import print_offers from dstack._internal.core.errors import CLIError, ConfigurationError +from dstack._internal.core.models.configurations import TaskConfiguration from dstack._internal.core.models.envs import 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 @@ -64,6 +71,7 @@ def create_endpoint_preset( debug: bool = False, resume_session: Optional[EndpointAgentSession] = None, user_prompt: Optional[str] = None, + allowed_fleets: Optional[tuple[str, ...]] = None, ) -> EndpointPresetCreateResult: agent_session = resume_session or create_endpoint_agent_session(configuration, debug=debug) try: @@ -79,6 +87,7 @@ def create_endpoint_preset( agent_session=agent_session, resume=resume_session is not None, user_prompt=user_prompt, + allowed_fleets=allowed_fleets, ) ) except KeyboardInterrupt: @@ -104,6 +113,7 @@ async def _create_endpoint_preset( agent_session: EndpointAgentSession, resume: bool = False, user_prompt: Optional[str] = None, + allowed_fleets: Optional[tuple[str, ...]] = None, ) -> EndpointPresetCreateResult: source_configuration = source_configuration or configuration initial_resume_session_id: Optional[str] = None @@ -125,14 +135,17 @@ async def _create_endpoint_preset( 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: tuple[str, ...] = () + allowed_fleets = () else: - allowed_fleets = _get_allowed_fleets(api, configuration) + 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, agent_session.preset_id) + build_name = build_name or _get_build_name( + configuration.name, configuration.model.api_model_name, agent_session.preset_id + ) agent_session.update_manifest(status="running", pid=os.getpid(), claude_model=auth.model) endpoint_env = configuration.env.as_dict() @@ -195,6 +208,7 @@ async def _create_endpoint_preset( 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") @@ -291,17 +305,23 @@ def _suspend_agent_session(session: EndpointAgentSession) -> None: ) -def _get_build_name(endpoint_name: Optional[str], suffix: str) -> str: - if endpoint_name is None: - raise CLIError( - "The service name is required. Set `name` in the configuration or use --name" - ) +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 = endpoint_name[:26].rstrip("-") + 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")) @@ -313,6 +333,38 @@ def _load_build_name(workspace: EndpointAgentWorkspace) -> str: 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( diff --git a/src/dstack/_internal/cli/services/endpoints/output.py b/src/dstack/_internal/cli/services/endpoints/output.py index f2e2393d2..c0e23fde5 100644 --- a/src/dstack/_internal/cli/services/endpoints/output.py +++ b/src/dstack/_internal/cli/services/endpoints/output.py @@ -58,6 +58,7 @@ def print_endpoint_presets( table = Table(box=None) 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) @@ -124,6 +125,7 @@ def _add_session(table: Table, session: dict[str, Any], base_label: Optional[str table, { "ID": str(session.get("id", "")), + "NAME": str(session.get("name") or ""), "GPU": gpu, "RESOURCES": gpu, "BENCHMARK": benchmark, @@ -144,6 +146,7 @@ def _add_preset( 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), diff --git a/src/dstack/_internal/cli/services/endpoints/presets.py b/src/dstack/_internal/cli/services/endpoints/presets.py index 922d21a7e..553730346 100644 --- a/src/dstack/_internal/cli/services/endpoints/presets.py +++ b/src/dstack/_internal/cli/services/endpoints/presets.py @@ -27,6 +27,7 @@ def build_endpoint_preset( context_length: int, benchmark: EndpointBenchmark, preset_id: Optional[str] = None, + name: Optional[str] = None, ) -> EndpointPreset: service = service.copy(deep=True) service.name = None @@ -39,6 +40,7 @@ def build_endpoint_preset( ) set_service_gpu_vendors_from_validations(service, [validation]) return EndpointPreset( + name=name, base=base_model, id=preset_id or make_endpoint_preset_id(service, context_length=context_length), model=model, @@ -68,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(), diff --git a/src/dstack/_internal/cli/services/endpoints/store.py b/src/dstack/_internal/cli/services/endpoints/store.py index 24e234e2d..766cdebd6 100644 --- a/src/dstack/_internal/cli/services/endpoints/store.py +++ b/src/dstack/_internal/cli/services/endpoints/store.py @@ -74,6 +74,21 @@ 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 detach_name(self, name: str) -> EndpointPreset | None: + """Detaches `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: diff --git a/src/dstack/_internal/cli/services/endpoints/verify.py b/src/dstack/_internal/cli/services/endpoints/verify.py index 33a9a9754..f1917b1c4 100644 --- a/src/dstack/_internal/cli/services/endpoints/verify.py +++ b/src/dstack/_internal/cli/services/endpoints/verify.py @@ -76,6 +76,7 @@ def build_verified_endpoint_preset( 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") @@ -116,6 +117,7 @@ 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, 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 28a339c0b..e6216e0b5 100644 --- a/src/tests/_internal/cli/commands/test_endpoint.py +++ b/src/tests/_internal/cli/commands/test_endpoint.py @@ -78,13 +78,17 @@ def test_handles_keyboard_interrupt(self, tmp_path, capsys): 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( - ["preset", "create", "-f", str(configuration_path)], + ["preset", "create", "-y", "-f", str(configuration_path)], home_dir=tmp_path, ) @@ -288,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, @@ -297,6 +305,7 @@ def test_merges_profile_configuration_and_cli_args(self, tmp_path): [ "preset", "create", + "-y", "-f", str(configuration_path), "--name", @@ -366,3 +375,97 @@ def test_apply_passes_selected_profile_and_preset_ids( assert apply.call_args.kwargs["profile_name"] == "gpu" 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/services/endpoints/test_create.py b/src/tests/_internal/cli/services/endpoints/test_create.py index 590ba5d7b..d01ac3cd4 100644 --- a/src/tests/_internal/cli/services/endpoints/test_create.py +++ b/src/tests/_internal/cli/services/endpoints/test_create.py @@ -320,11 +320,12 @@ async def run_agent(**kwargs): class TestBuildName: - def test_requires_name_and_keeps_generated_prefix_bounded(self): - with pytest.raises(CLIError, match="The service name is required"): - _get_build_name(None, "a1b2c3d4") + def test_derives_slug_for_nameless_and_keeps_prefix_bounded(self): + assert _get_build_name(None, "Qwen/Qwen3.5-27B", "a1b2c3d4") == "qwen3-5-27b-a1b2c3d4" - build_name = _get_build_name("qwen-endpoint-with-a-name-that-is-forty-one", "a1b2c3d4") + build_name = _get_build_name( + "qwen-endpoint-with-a-name-that-is-forty-one", "Qwen/Qwen3.5-27B", "a1b2c3d4" + ) assert build_name.endswith("-a1b2c3d4") assert len(f"{build_name}-99999") <= 41 diff --git a/src/tests/_internal/cli/services/endpoints/test_store.py b/src/tests/_internal/cli/services/endpoints/test_store.py index 1ceb731f7..27c88ddca 100644 --- a/src/tests/_internal/cli/services/endpoints/test_store.py +++ b/src/tests/_internal/cli/services/endpoints/test_store.py @@ -218,3 +218,20 @@ def test_rejects_missing_and_empty_prompt_files(self, tmp_path: Path): 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 + + detached = store.detach_name("qwen") + + assert detached.id == named.id + assert store.find_by_name("qwen") is None + assert store.get(named.id).name is None From 1c84e6f3082a4f239a30f5e3c047f2c5a05ef0c1 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Tue, 21 Jul 2026 16:35:14 -0700 Subject: [PATCH 09/25] Detachable preset creation sessions; replace the agent watchdog The agent's streams go to workspace files instead of a pipe into the CLI, so the agent survives CLI death and a later attach continues parsing from the persisted offsets. The session manifest tracks the agent's own pid, and list liveness keys on it, so a CLI-less agent shows as running (clauding/verifying) rather than interrupted. - Ctrl+C during create/attach asks to stop or detach; detaching (or any hard CLI death) leaves the agent working and visible - `dstack preset attach ` follows a detached session to completion and finalizes it; `stop ` terminates the agent and offers to stop its active runs - Stopping a session offers to stop its non-terminal runs; the cost warning only shows when the user keeps them for a later resume - Preset name conflicts prompt to reassign the name, leaving the old preset name-less - The watchdog is removed: CLI death now means detach, which is safe because detached sessions stay visible and attachable Co-Authored-By: Claude Fable 5 --- src/dstack/_internal/cli/commands/endpoint.py | 68 ++- .../_internal/cli/services/endpoints/agent.py | 472 ++++++++++++------ .../cli/services/endpoints/create.py | 251 ++++++++-- .../_internal/cli/services/endpoints/store.py | 4 +- .../cli/services/endpoints/test_agent.py | 126 +++-- .../cli/services/endpoints/test_create.py | 152 +++++- .../cli/services/endpoints/test_store.py | 4 +- 7 files changed, 817 insertions(+), 260 deletions(-) diff --git a/src/dstack/_internal/cli/commands/endpoint.py b/src/dstack/_internal/cli/commands/endpoint.py index 0fc51a5b2..c68d01568 100644 --- a/src/dstack/_internal/cli/commands/endpoint.py +++ b/src/dstack/_internal/cli/commands/endpoint.py @@ -14,13 +14,16 @@ 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 ( + 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 ( @@ -107,6 +110,30 @@ def _register(self) -> None: ) 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 a preset", @@ -235,6 +262,31 @@ def _create(self, args: argparse.Namespace) -> None: 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: store = EndpointPresetStore() preset = store.get(args.preset) or store.find_by_name(args.preset) @@ -373,10 +425,20 @@ def _apply_name(configuration: EndpointConfiguration, name: str | None, *, requi 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; detaches the name from any holder on yes.""" + """One apply-style confirmation; reassigns the name from any holder on yes.""" preset_holder = None session_holders = [] if name is not None: @@ -393,7 +455,7 @@ def _confirm_preset_creation( if holders: message = ( f"The name [code]{name}[/] is already used by {', '.join(holders)}." - f" Detach it and create a new preset?" + f" Reassign it to a new preset?" ) elif name is not None: message = f"Create the preset [code]{name}[/]?" @@ -402,7 +464,7 @@ def _confirm_preset_creation( if not assume_yes and not confirm_ask(message): return False if preset_holder is not None and name is not None: - store.detach_name(name) + store.release_name(name) for session in session_holders: session.update_manifest(name=None) return True diff --git a/src/dstack/_internal/cli/services/endpoints/agent.py b/src/dstack/_internal/cli/services/endpoints/agent.py index d86853868..75df539db 100644 --- a/src/dstack/_internal/cli/services/endpoints/agent.py +++ b/src/dstack/_internal/cli/services/endpoints/agent.py @@ -7,11 +7,12 @@ import subprocess import sys import tempfile -from contextlib import 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, Optional, Sequence +from typing import Any, AsyncIterator, Callable, Optional, Protocol, Sequence import psutil import yaml @@ -41,7 +42,6 @@ _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. " @@ -120,6 +120,14 @@ def constraints_path(self) -> Path: 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: @@ -337,16 +345,16 @@ def create_endpoint_agent_session( "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 / "preset.dstack.yml", - yaml.safe_dump(data, sort_keys=False), - ) _write_private_text(path / "trace.jsonl", "") except OSError as e: if path is not None: @@ -398,15 +406,11 @@ def load_resumable_agent_session(preset_id: str) -> EndpointAgentSession: 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") - pid = manifest.get("pid") - if ( - status == "running" - and isinstance(pid, int) - and pid > 0 - and pid != os.getpid() - and psutil.pid_exists(pid) - ): - raise CLIError(f"Session {preset_id} appears to be running already (pid {pid})") + 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" @@ -416,6 +420,53 @@ def load_resumable_agent_session(preset_id: str) -> EndpointAgentSession: 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: @@ -452,10 +503,7 @@ def list_agent_sessions() -> list[dict[str, Any]]: status = manifest.get("status") if status not in ("running", "interrupted", "success"): continue - pid = manifest.get("pid") - if status == "running" and not ( - isinstance(pid, int) and pid > 0 and psutil.pid_exists(pid) - ): + if status == "running" and not session_process_alive(manifest): status = "interrupted" entry = dict(manifest) entry["id"] = path.name @@ -584,33 +632,9 @@ async def run_endpoint_agent( agent_session: EndpointAgentSession, initial_resume_session_id: Optional[str] = None, ) -> EndpointAgentProcessOutput: - 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: + 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) @@ -653,15 +677,6 @@ async def run_endpoint_agent( agent_session=agent_session, ) await asyncio.sleep(delay) - 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 _run_claude_process( @@ -674,62 +689,54 @@ async def _run_claude_process( agent_session: EndpointAgentSession, ) -> tuple[EndpointAgentProcessOutput, int]: proc: Optional[asyncio.subprocess.Process] = None - watchdog: Optional[subprocess.Popen] = 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) ) - watchdog = _start_agent_watchdog(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: - if watchdog is not None: - await asyncio.to_thread(_stop_agent_watchdog, watchdog) - 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, returncode @@ -992,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], @@ -1072,76 +1194,38 @@ async def _terminate_process(proc: asyncio.subprocess.Process) -> None: await proc.wait() -# The watchdog reads its stdin until EOF. The CLI holds the write end and never -# writes while the agent runs, so EOF means the CLI is gone — any death, -# including SIGKILL, which the CLI cannot react to itself. A disarm byte is -# written when the CLI stops the agent on its own, so the watchdog exits -# without touching anything (and cannot kill a reused pid). -_AGENT_WATCHDOG_SCRIPT = """ -import os -import signal -import sys -import time - -import psutil - -agent_pid = int(sys.argv[1]) -if sys.stdin.buffer.read(): - sys.exit(0) # disarmed: the CLI stopped the agent itself - - -def _kill(hard): - if os.name == "posix": - try: - os.killpg(agent_pid, signal.SIGKILL if hard else signal.SIGTERM) - except OSError: - pass +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 - try: - root = psutil.Process(agent_pid) - processes = [*root.children(recursive=True), root] - except psutil.NoSuchProcess: + if IS_WINDOWS: + _terminate_windows_process_tree(agent_pid) return - for process in processes: - try: - process.kill() if hard else process.terminate() - except psutil.NoSuchProcess: - pass - - -if psutil.pid_exists(agent_pid): - _kill(hard=False) - deadline = time.monotonic() + 15 - while time.monotonic() < deadline and psutil.pid_exists(agent_pid): - time.sleep(0.5) - if psutil.pid_exists(agent_pid): - _kill(hard=True) -""" + 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 _start_agent_watchdog(agent_pid: int) -> subprocess.Popen: - """Tie the agent's life to this process: if this process dies in a way it - cannot react to, the watchdog kills the agent's process tree.""" - return subprocess.Popen( - [sys.executable, "-c", _AGENT_WATCHDOG_SCRIPT, str(agent_pid)], - stdin=subprocess.PIPE, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - start_new_session=not IS_WINDOWS, - ) +def _pid_running(pid: int) -> bool: + try: + return psutil.Process(pid).status() != psutil.STATUS_ZOMBIE + except psutil.Error: + return False -def _stop_agent_watchdog(watchdog: subprocess.Popen) -> None: - if watchdog.stdin is not None: - with suppress(OSError, ValueError): - watchdog.stdin.write(b"x") - watchdog.stdin.flush() - with suppress(OSError, ValueError): - watchdog.stdin.close() +def _process_started_at(pid: int) -> Optional[float]: try: - watchdog.wait(timeout=10) - except subprocess.TimeoutExpired: - watchdog.kill() + return psutil.Process(pid).create_time() + except psutil.Error: + return None def _terminate_windows_process_tree(pid: int) -> None: @@ -1160,6 +1244,64 @@ 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.""" diff --git a/src/dstack/_internal/cli/services/endpoints/create.py b/src/dstack/_internal/cli/services/endpoints/create.py index d9db24f2d..763c2a7bf 100644 --- a/src/dstack/_internal/cli/services/endpoints/create.py +++ b/src/dstack/_internal/cli/services/endpoints/create.py @@ -9,6 +9,7 @@ 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 @@ -21,6 +22,7 @@ EndpointAgentSession, EndpointAgentWorkspace, attach_agent_workspace, + attach_endpoint_agent, build_endpoint_agent_env, claimed_session_name, contains_redacted_value, @@ -29,10 +31,13 @@ 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 get_endpoint_agent_system_prompt @@ -41,11 +46,11 @@ 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.configurations import TaskConfiguration -from dstack._internal.core.models.envs import EnvSentinel +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 @@ -61,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, @@ -91,7 +217,7 @@ def create_endpoint_preset( ) ) except KeyboardInterrupt: - _suspend_agent_session(agent_session) + _stop_or_detach_agent_session(agent_session, api) raise except BaseException: _finish_agent_session(agent_session, "failed") @@ -112,12 +238,18 @@ async def _create_endpoint_preset( 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 initial_resume_session_id: Optional[str] = None - if resume: + 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: @@ -146,7 +278,10 @@ async def _create_endpoint_preset( build_name = build_name or _get_build_name( configuration.name, configuration.model.api_model_name, agent_session.preset_id ) - agent_session.update_manifest(status="running", pid=os.getpid(), claude_model=auth.model) + 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) @@ -155,26 +290,28 @@ 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 - env = build_endpoint_agent_env( - api=api, - endpoint_env=endpoint_env, - auth=auth, - workspace=workspace, - token=token, - ) + if auth is not None: + env = build_endpoint_agent_env( + api=api, + endpoint_env=endpoint_env, + auth=auth, + workspace=workspace, + token=token, + ) prompt = get_endpoint_agent_system_prompt(user_prompt=user_prompt) - if not resume: + if not resume and not attach: if user_prompt: agent_session.write_user_prompt(user_prompt) constraints_text = _build_constraints( @@ -186,17 +323,28 @@ async def _create_endpoint_preset( if agent_session.debug: agent_session.write_prompt(prompt) agent_session.write_constraints(constraints_text) - agent_session.write_agent_info(auth) + if auth is not None: + agent_session.write_agent_info(auth) try: - process_output = await run_endpoint_agent( - prompt=prompt, - env=env, - workspace=workspace, - auth=auth, - redacted_values=redacted_values, - agent_session=agent_session, - initial_resume_session_id=initial_resume_session_id, - ) + 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, + workspace=workspace, + 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, workspace=workspace, @@ -290,19 +438,56 @@ def _finish_agent_session( console.print(f"Agent log saved to [code]{path / 'agent.log'}[/]") -def _suspend_agent_session(session: EndpointAgentSession) -> None: +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}") - 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[/]." - ) + 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: diff --git a/src/dstack/_internal/cli/services/endpoints/store.py b/src/dstack/_internal/cli/services/endpoints/store.py index 766cdebd6..561d88a47 100644 --- a/src/dstack/_internal/cli/services/endpoints/store.py +++ b/src/dstack/_internal/cli/services/endpoints/store.py @@ -80,8 +80,8 @@ def find_by_name(self, name: str) -> EndpointPreset | None: return preset return None - def detach_name(self, name: str) -> EndpointPreset | None: - """Detaches `name` from the preset holding it, keeping the preset.""" + 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 diff --git a/src/tests/_internal/cli/services/endpoints/test_agent.py b/src/tests/_internal/cli/services/endpoints/test_agent.py index 5d1ba4809..3b569fd9d 100644 --- a/src/tests/_internal/cli/services/endpoints/test_agent.py +++ b/src/tests/_internal/cli/services/endpoints/test_agent.py @@ -22,8 +22,6 @@ _prepare_subprocess_command, _ProgressTailer, _RecordMirror, - _start_agent_watchdog, - _stop_agent_watchdog, _summarize_session_trials, _terminate_process, attach_agent_workspace, @@ -187,7 +185,11 @@ def test_saves_log_and_debug_files(self, tmp_path, monkeypatch, capsys): 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"} + 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" @@ -411,57 +413,6 @@ async def test_terminates_windows_process_tree(self): assert not psutil.pid_exists(child_pid) -class TestAgentWatchdog: - @pytest.mark.skipif(IS_WINDOWS, reason="exercises POSIX signals and sessions") - def test_kills_agent_when_cli_dies_hard(self): - # A fake CLI spawns a fake agent in its own session, arms the real - # watchdog, and hangs. SIGKILL of the CLI must take the agent down. - cli_script = ( - "import subprocess, sys, time\n" - "from dstack._internal.cli.services.endpoints.agent import _start_agent_watchdog\n" - "agent = subprocess.Popen(\n" - " [sys.executable, '-c', 'import time; time.sleep(300)'], start_new_session=True\n" - ")\n" - "_start_agent_watchdog(agent.pid)\n" - "print(agent.pid, flush=True)\n" - "time.sleep(300)\n" - ) - cli = subprocess.Popen( - [sys.executable, "-c", cli_script], stdout=subprocess.PIPE, text=True - ) - try: - assert cli.stdout is not None - agent_pid = int(cli.stdout.readline()) - assert psutil.pid_exists(agent_pid) - - cli.kill() - cli.wait(timeout=10) - - try: - psutil.Process(agent_pid).wait(timeout=20) - except psutil.NoSuchProcess: - pass - assert not psutil.pid_exists(agent_pid) - finally: - with suppress(OSError): - cli.kill() - - @pytest.mark.skipif(IS_WINDOWS, reason="exercises POSIX signals and sessions") - def test_disarmed_watchdog_leaves_agent_untouched(self): - agent = subprocess.Popen( - [sys.executable, "-c", "import time; time.sleep(300)"], start_new_session=True - ) - try: - watchdog = _start_agent_watchdog(agent.pid) - _stop_agent_watchdog(watchdog) - - assert watchdog.poll() is not None - assert agent.poll() is None - finally: - with suppress(OSError): - os.killpg(agent.pid, signal.SIGKILL) - - class TestRecordMirror: def test_mirrors_complete_lines_redacted(self, tmp_path): source = tmp_path / "runs.jsonl" @@ -885,3 +836,70 @@ def test_counts_records_even_when_trials_share_a_task(self, tmp_path): # 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) diff --git a/src/tests/_internal/cli/services/endpoints/test_create.py b/src/tests/_internal/cli/services/endpoints/test_create.py index d01ac3cd4..d8a64ecfc 100644 --- a/src/tests/_internal/cli/services/endpoints/test_create.py +++ b/src/tests/_internal/cli/services/endpoints/test_create.py @@ -21,7 +21,10 @@ _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 @@ -126,7 +129,11 @@ async def create(**kwargs): if path.is_dir() and not path.name.startswith(".") ] assert len(paths) == 1 - assert {path.name for path in paths[0].iterdir()} == {"agent.log", "session.json"} + 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 @@ -643,3 +650,146 @@ async def run_agent(**kwargs): 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_store.py b/src/tests/_internal/cli/services/endpoints/test_store.py index 27c88ddca..f205c6964 100644 --- a/src/tests/_internal/cli/services/endpoints/test_store.py +++ b/src/tests/_internal/cli/services/endpoints/test_store.py @@ -230,8 +230,8 @@ def test_finds_and_detaches_names(self, tmp_path: Path): assert store.find_by_name("qwen").id == named.id assert store.find_by_name("other") is None - detached = store.detach_name("qwen") + released = store.release_name("qwen") - assert detached.id == named.id + assert released.id == named.id assert store.find_by_name("qwen") is None assert store.get(named.id).name is None From ec01c27e3505293c41e698ab29a2dcd0b09fb640 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Tue, 21 Jul 2026 16:48:12 -0700 Subject: [PATCH 10/25] Fix offset store clobbering across shared readers and mirrors The stream readers and the record mirrors each open a separate _OffsetStore on the same file with disjoint keys; writing the whole in-memory view dropped the other store's offsets, causing duplicated mirror records on the next attach or resume. Update only the written key, re-reading the file first. Co-Authored-By: Claude Fable 5 --- .../_internal/cli/services/endpoints/agent.py | 11 +++++++++- .../cli/services/endpoints/test_agent.py | 21 +++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/src/dstack/_internal/cli/services/endpoints/agent.py b/src/dstack/_internal/cli/services/endpoints/agent.py index 75df539db..0c97fdddf 100644 --- a/src/dstack/_internal/cli/services/endpoints/agent.py +++ b/src/dstack/_internal/cli/services/endpoints/agent.py @@ -1319,8 +1319,17 @@ def get(self, key: str) -> int: 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(self._data) + "\n") + _write_private_text(self._path, json.dumps(data) + "\n") class _ProgressTailer: diff --git a/src/tests/_internal/cli/services/endpoints/test_agent.py b/src/tests/_internal/cli/services/endpoints/test_agent.py index 3b569fd9d..a8fabd89b 100644 --- a/src/tests/_internal/cli/services/endpoints/test_agent.py +++ b/src/tests/_internal/cli/services/endpoints/test_agent.py @@ -903,3 +903,24 @@ def test_detach_keeps_the_agent_and_stop_terminates_it(self, tmp_path, monkeypat 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 From 57f2f201bb22cd8c8dd574409a5230800909a472 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Wed, 22 Jul 2026 00:02:28 -0700 Subject: [PATCH 11/25] Align preset list statuses with dstack ps; move NAME to last column Statuses use the same palette as `dstack ps` (verified=grey, clauding= sea_green3, verifying/interrupted/failed matching), and NAME becomes a trailing dimmed column so the data columns stay contiguous. Co-Authored-By: Claude Fable 5 --- src/dstack/_internal/cli/services/endpoints/output.py | 6 +++--- src/tests/_internal/cli/services/endpoints/test_output.py | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/dstack/_internal/cli/services/endpoints/output.py b/src/dstack/_internal/cli/services/endpoints/output.py index c0e23fde5..460248917 100644 --- a/src/dstack/_internal/cli/services/endpoints/output.py +++ b/src/dstack/_internal/cli/services/endpoints/output.py @@ -11,8 +11,8 @@ from dstack._internal.utils.common import pretty_date, pretty_resources _STATUS_DISPLAY = { - "ready": ("success", "green"), - "running": ("clauding", "bold deep_sky_blue1"), + "ready": ("verified", "grey"), + "running": ("clauding", "bold sea_green3"), "verifying": ("verifying", "bold deep_sky_blue1"), "interrupted": ("interrupted", "bold gold1"), "failed": ("failed", "indian_red1"), @@ -58,11 +58,11 @@ def print_endpoint_presets( table = Table(box=None) 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") + table.add_column("NAME", no_wrap=True, style="secondary") presets_by_base: dict[str, list[EndpointPreset]] = defaultdict(list) repo_to_base: dict[str, str] = {} for preset in presets: diff --git a/src/tests/_internal/cli/services/endpoints/test_output.py b/src/tests/_internal/cli/services/endpoints/test_output.py index 6bb188736..eccad45f4 100644 --- a/src/tests/_internal/cli/services/endpoints/test_output.py +++ b/src/tests/_internal/cli/services/endpoints/test_output.py @@ -48,7 +48,7 @@ def test_shows_progress_after_status_and_best_benchmark(self): } ) - assert row["STATUS"] == "[bold deep_sky_blue1]clauding[/] [secondary](2/3)[/]" + assert row["STATUS"] == "[bold sea_green3]clauding[/] [secondary](2/3)[/]" assert row["BENCHMARK"] == "best trial: con=8 2339 tok/s" assert row["GPU"] == "A40:48GB:1" @@ -57,7 +57,7 @@ def test_shows_zero_progress_without_benchmark(self): {"id": "ab12cd34", "status": "running", "max_trials": 3, "trials": {"count": 0}} ) - assert row["STATUS"] == "[bold deep_sky_blue1]clauding[/] [secondary](0/3)[/]" + assert row["STATUS"] == "[bold sea_green3]clauding[/] [secondary](0/3)[/]" assert row["BENCHMARK"] == "" def test_omits_progress_without_trials_data(self): @@ -129,7 +129,7 @@ def test_completed_creation_decorates_preset_row_without_extra_session_row(self, output_module.print_endpoint_presets([preset], sessions=sessions) text = buffer.getvalue() - assert "success (3/4)" in text + assert "verified (3/4)" in text assert text.count(preset.id) == 1 @@ -146,4 +146,4 @@ def test_running_session_with_remaining_trials_stays_clauding(self): {"id": "ab12cd34", "status": "running", "max_trials": 2, "trials": {"count": 1}} ) - assert row["STATUS"].startswith("[bold deep_sky_blue1]clauding[/]") + assert row["STATUS"].startswith("[bold sea_green3]clauding[/]") From dbd70cc34a6d8b0bfd8d663922c003cdf7ee6cb0 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Wed, 22 Jul 2026 00:02:28 -0700 Subject: [PATCH 12/25] Require `model` on the final service in the agent prompt MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `model` is what makes the run a served, readiness-checked model endpoint — i.e. a preset. Tell the agent a service that never passes the readiness probe is a real failure, not something to work around by removing `model`. Co-Authored-By: Claude Fable 5 --- .../cli/services/endpoints/resources/system_prompt.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 ab62a35b4..f6b610d1b 100644 --- a/src/dstack/_internal/cli/services/endpoints/resources/system_prompt.md +++ b/src/dstack/_internal/cli/services/endpoints/resources/system_prompt.md @@ -297,7 +297,10 @@ 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`). +`constraints.json` (see `# Constraints`). `model` is required: it also enables +`dstack`'s default readiness probe — an independent check that the model can +serve requests. If the service never passes the probe, treat that as a real +failure of the configuration, not something to work around by removing `model`. Before the final benchmark, verify the model through the service: send real requests using the client-facing model name and check that the model works From be6198fc7cf8b6877f7999de0112297845958a5f Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Wed, 22 Jul 2026 09:43:50 -0700 Subject: [PATCH 13/25] Add `dstack preset logs [-f]`; reconcile detached preset creations on read MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace `dstack preset attach` with `dstack preset logs [-f]`: dump a preset creation's log at any status; `-f` re-owns and finalizes a detached or running creation, and extra followers tail read-only instead of being rejected. Reconcile-on-read: `list`, `get`, and `apply` finalize a completed-but-orphaned creation from its on-disk report, so the preset is saved regardless of how the CLI went away (detached, killed, or crashed). Use advisory file locks (flock/msvcrt) for creation ownership instead of pid claims — no steal race, no stale locks, kernel auto-release on exit. Ctrl+C during create reports Detached/interrupted via the interrupt handler; stop auto-stops runs with a spinner and never re-saves a finished creation. Co-Authored-By: Claude Opus 4.8 --- src/dstack/_internal/cli/commands/endpoint.py | 107 +++--- .../_internal/cli/services/endpoints/agent.py | 209 +++++++++-- .../cli/services/endpoints/create.py | 348 +++++++++++++----- .../_internal/cli/commands/test_endpoint.py | 4 +- .../cli/services/endpoints/test_agent.py | 8 +- .../cli/services/endpoints/test_create.py | 250 ++++++++++++- 6 files changed, 725 insertions(+), 201 deletions(-) diff --git a/src/dstack/_internal/cli/commands/endpoint.py b/src/dstack/_internal/cli/commands/endpoint.py index d80f33b00..35731d4c1 100644 --- a/src/dstack/_internal/cli/commands/endpoint.py +++ b/src/dstack/_internal/cli/commands/endpoint.py @@ -1,6 +1,7 @@ import argparse import os import time +from contextlib import suppress from pathlib import Path from argcomplete import FilesCompleter # type: ignore[attr-defined] @@ -20,9 +21,10 @@ ) from dstack._internal.cli.services.endpoints.apply import apply_endpoint_preset from dstack._internal.cli.services.endpoints.create import ( - attach_endpoint_preset, create_endpoint_preset, plan_endpoint_preset, + reconcile_detached_sessions, + show_endpoint_session_logs, stop_endpoint_session, ) from dstack._internal.cli.services.endpoints.output import print_endpoint_presets @@ -106,29 +108,35 @@ def _register(self) -> None: create_parser.add_argument( "--resume", metavar="ID", - help="Resume an interrupted preset creation session by its preset ID", + help="Resume an interrupted preset creation by its preset ID", ) create_parser.add_argument( "-y", "--yes", action="store_true", help="Do not ask for confirmation" ) create_parser.set_defaults(subfunc=self._create) - attach_parser = preset_subparsers.add_parser( - "attach", - help="Attach to a detached preset creation session", + logs_parser = preset_subparsers.add_parser( + "logs", + help="Show a preset's creation log", formatter_class=self._parser.formatter_class, ) - attach_parser.add_argument("preset", metavar="ID", help="The preset ID or name") - attach_parser.add_argument( + logs_parser.add_argument("preset", metavar="ID", help="The preset ID or name") + logs_parser.add_argument( + "-f", + "--follow", + action="store_true", + help="Follow to completion and save the preset", + ) + logs_parser.add_argument( "--keep-service", action="store_true", - help="Leave the verified service running", + help="Leave the verified service running (with -f)", ) - attach_parser.set_defaults(subfunc=self._attach) + logs_parser.set_defaults(subfunc=self._logs) stop_parser = preset_subparsers.add_parser( "stop", - help="Stop a running preset creation session", + help="Stop a running preset creation", formatter_class=self._parser.formatter_class, ) stop_parser.add_argument("preset", metavar="ID", help="The preset ID or name") @@ -200,15 +208,24 @@ def _command(self, args: argparse.Namespace) -> None: console.print("\nOperation interrupted by user. Exiting...") exit(0) + def _reconcile(self) -> None: + """Finalize any detached/orphaned session whose agent already completed, + so a saved preset never depends on a foreground CLI surviving. Fully + best-effort — it must never make a read command fail.""" + with suppress(Exception): + reconcile_detached_sessions(EndpointPresetStore()) + def _list(self, args: argparse.Namespace) -> None: base = getattr(args, "base", None) repo = getattr(args, "repo", None) if getattr(args, "json", False): + self._reconcile() presets = _filter_presets(EndpointPresetStore().list(), base=base, repo=repo) print(EndpointPresetListOutput(presets=presets).json()) return verbose = getattr(args, "verbose", False) while True: + self._reconcile() presets = EndpointPresetStore().list() sessions = list_agent_sessions() if base or repo: @@ -239,7 +256,7 @@ def _create(self, args: argparse.Namespace) -> None: if getattr(args, "max_trials", None) is not None: console.print( "[warning]--max-trials is ignored when resuming: " - "session constraints are fixed at creation[/]" + "the constraints are fixed at creation[/]" ) api = Client.from_config(project_name=args.project) allowed_fleets = None @@ -248,16 +265,19 @@ def _create(self, args: argparse.Namespace) -> None: if not _confirm_preset_creation(store, configuration.name, assume_yes=args.yes): console.print("\nExiting...") return - result = create_endpoint_preset( - api=api, - configuration=configuration, - store=store, - keep_service=args.keep_service, - debug=args.debug, - resume_session=resume_session, - user_prompt=user_prompt, - allowed_fleets=allowed_fleets, - ) + try: + result = create_endpoint_preset( + api=api, + configuration=configuration, + store=store, + keep_service=args.keep_service, + debug=args.debug, + resume_session=resume_session, + user_prompt=user_prompt, + allowed_fleets=allowed_fleets, + ) + except KeyboardInterrupt: + return # the interrupt handler already reported detach / stop console.print( f"Preset [code]{result.preset.id}[/] for " f"[code]{result.preset.base}[/] saved to [code]{result.path}[/]" @@ -265,32 +285,34 @@ def _create(self, args: argparse.Namespace) -> None: 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 _logs(self, args: argparse.Namespace) -> None: + try: + result = show_endpoint_session_logs( + project=args.project, + store=EndpointPresetStore(), + preset_id=_resolve_session_ref(args.preset), + follow=args.follow, + keep_service=args.keep_service, + ) + except KeyboardInterrupt: + return # a log viewer: Ctrl+C just stops watching, quietly + if result is not None: + 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}[/]?" - ): + if not args.yes and not confirm_ask(f"Stop creating preset [code]{preset_id}[/]?"): console.print("\nExiting...") return - stop_endpoint_session( - Client.from_config(project_name=args.project), preset_id, assume_yes=args.yes - ) + stop_endpoint_session(Client.from_config(project_name=args.project), preset_id) def _get(self, args: argparse.Namespace) -> None: + self._reconcile() store = EndpointPresetStore() preset = store.get(args.preset) or store.find_by_name(args.preset) if preset is None: @@ -298,6 +320,7 @@ def _get(self, args: argparse.Namespace) -> None: print(preset.json()) def _apply(self, args: argparse.Namespace) -> None: + self._reconcile() configuration_path, configuration = load_endpoint_configuration(args.configuration_file) configuration = _get_effective_configuration(configuration, args) apply_endpoint_preset( @@ -367,7 +390,7 @@ def _add_list_args(parser: argparse.ArgumentParser) -> None: "-w", "--watch", action="store_true", - help="Watch presets and sessions in realtime", + help="Watch presets in realtime", ) parser.add_argument( "--json", @@ -454,7 +477,7 @@ def _confirm_preset_creation( 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) + holders.extend(f"preset [code]{session.preset_id}[/]" for session in session_holders) if holders: message = ( f"The name [code]{name}[/] is already used by {', '.join(holders)}." diff --git a/src/dstack/_internal/cli/services/endpoints/agent.py b/src/dstack/_internal/cli/services/endpoints/agent.py index 0c97fdddf..4369f9de2 100644 --- a/src/dstack/_internal/cli/services/endpoints/agent.py +++ b/src/dstack/_internal/cli/services/endpoints/agent.py @@ -12,7 +12,7 @@ from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path -from typing import Any, AsyncIterator, Callable, Optional, Protocol, Sequence +from typing import Any, AsyncIterator, Callable, Iterator, Optional, Protocol, Sequence import psutil import yaml @@ -87,6 +87,11 @@ ) +class SessionBusyError(CLIError): + """Another live process owns the session — it is following or finalizing it. + Callers that only want to view can fall back to a read-only follow.""" + + @dataclass(frozen=True) class EndpointAgentWorkspace: path: Path @@ -135,6 +140,10 @@ class EndpointAgentSession: timestamp: str debug: bool preset_id: str = "" + # Whether progress lines echo to this process's console (a live attach), on + # top of always being recorded to agent.log. Background reconcile sets it + # False so finalizing a detached session stays silent on the read command. + echo: bool = field(default=True, repr=False) _log_enabled: bool = field(default=True, init=False, repr=False) @property @@ -197,7 +206,8 @@ def append_log(self, line: str) -> None: f.flush() except OSError as e: self._log_enabled = False - console.print(f"[warning]Could not write agent log {self.log_path}: {e}[/]") + if self.echo: + console.print(f"[warning]Could not write agent log {self.log_path}: {e}[/]") def read_manifest(self) -> dict[str, Any]: try: @@ -258,12 +268,12 @@ def attach_agent_workspace(session: EndpointAgentSession) -> EndpointAgentWorksp 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") + raise CLIError("The preset creation has no workspace to resume") real, alias = Path(real_value), Path(alias_value) if not real.is_dir(): raise CLIError( - f"The session workspace no longer exists: {real}. " - "Stop any leftover runs manually and start a new session." + f"The preset creation workspace no longer exists: {real}. " + "Stop any leftover runs manually and create a new preset." ) if alias != real: _ensure_workspace_alias(alias, real) @@ -336,6 +346,7 @@ def create_endpoint_agent_session( "id": preset_id, "status": "running", "pid": os.getpid(), + "pid_started_at": _process_started_at(os.getpid()), "endpoint": configuration.name, "name": configuration.name, "model": getattr(configuration.model, "base", None) @@ -400,21 +411,19 @@ def load_resumable_agent_session(preset_id: str) -> EndpointAgentSession: 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}") + raise CLIError(f"Unknown preset: {preset_id}") status = manifest.get("status") if status == "success": - raise CLIError(f"Session {preset_id} completed successfully; nothing to resume") + raise CLIError(f"Preset {preset_id} is already created; nothing to resume") if status == "failed": - raise CLIError(f"Session {preset_id} failed and cannot be resumed") + raise CLIError(f"Preset {preset_id} creation 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}[/]" + f"Preset {preset_id} is still being created;" + f" follow it with [code]dstack preset logs -f {preset_id}[/]" ) if not manifest.get("claude_session_id"): - raise CLIError( - f"Session {preset_id} stopped before the agent started; start a new session" - ) + raise CLIError(f"Preset {preset_id} creation stopped before it started; create a new one") session.debug = bool(manifest.get("debug")) session.timestamp = str(manifest.get("created_at") or "") return session @@ -437,7 +446,12 @@ def session_process_alive(manifest: dict[str, Any]) -> bool: 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) + if not isinstance(pid, int) or pid <= 0 or pid == os.getpid(): + return False + # Guard the CLI pid with its start time too: after an ungraceful CLI death + # the OS can recycle the pid, and a bare pid_exists() would read a dead + # session as still owned (falsely blocking reconcile / follow). + return _pid_alive(pid, manifest.get("pid_started_at")) def load_attachable_agent_session(preset_id: str) -> EndpointAgentSession: @@ -445,21 +459,26 @@ def load_attachable_agent_session(preset_id: str) -> EndpointAgentSession: 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}") + raise CLIError(f"Unknown preset: {preset_id}") status = manifest.get("status") if status == "success": - raise CLIError(f"Session {preset_id} completed successfully; nothing to attach") + raise CLIError(f"Preset {preset_id} is already created") if status == "failed": - raise CLIError(f"Session {preset_id} failed and cannot be attached") + raise CLIError(f"Preset {preset_id} creation failed") if status == "interrupted": raise CLIError( - f"Session {preset_id} is interrupted; resume it with" - f" [code]dstack preset create -f --resume {preset_id}[/]" + f"Preset {preset_id} creation was 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});" + if ( + isinstance(pid, int) + and pid > 0 + and pid != os.getpid() + and _pid_alive(pid, manifest.get("pid_started_at")) + ): + raise SessionBusyError( + f"Preset {preset_id} is already being followed by another CLI (pid {pid});" f" stop or detach it there with Ctrl+C" ) session.debug = bool(manifest.get("debug")) @@ -467,6 +486,109 @@ def load_attachable_agent_session(preset_id: str) -> EndpointAgentSession: return session +def load_agent_session(preset_id: str) -> EndpointAgentSession: + """Loads a session of any status for read-only inspection (its log).""" + path = get_presets_dir() / preset_id + session = EndpointAgentSession(path=path, timestamp="", debug=False, preset_id=preset_id) + if not path.is_dir() or not session.read_manifest(): + raise CLIError(f"Unknown preset: {preset_id}") + return session + + +def print_session_log(session: EndpointAgentSession) -> None: + """Prints the session's redacted progress log verbatim (no markup).""" + try: + content = session.log_path.read_text(encoding="utf-8") + except OSError: + content = "" + if content.strip(): + console.print(Text(content.rstrip("\n")), soft_wrap=True) + else: + console.print(f"No log output yet for session [code]{session.preset_id}[/].") + + +def mark_session_owner( + session: EndpointAgentSession, + *, + project: Optional[str] = None, + keep_service: Optional[bool] = None, + claude_model: Optional[str] = None, +) -> None: + """Records this process as the session's owner (pid + start time) and, when + given, the finalize context a later detached reconcile needs (project and + keep-service intent). `None` fields are left untouched.""" + fields: dict[str, Any] = { + "status": "running", + "pid": os.getpid(), + "pid_started_at": _process_started_at(os.getpid()), + } + if project is not None: + fields["project"] = project + if keep_service is not None: + fields["keep_service"] = keep_service + if claude_model is not None: + fields["claude_model"] = claude_model + session.update_manifest(**fields) + + +def session_report_exists(manifest: dict[str, Any]) -> bool: + """Whether the agent left a final report on disk — the durable completion + signal a detached session is finalized from.""" + workspace = manifest.get("workspace") + if not isinstance(workspace, str) or not workspace: + return False + return (Path(workspace) / "w" / _FINAL_REPORT_FILENAME).is_file() + + +def try_claim_session(session: EndpointAgentSession) -> Optional[int]: + """Takes an exclusive, kernel-held lock for the duration of a session's + finalization, so concurrent readers can't both finalize it. Returns an open + file descriptor to release via `release_session_claim`, or None if another + process holds it. The kernel drops the lock if the holder dies, so there are + no stale locks to reason about.""" + try: + fd = os.open(session.path / ".reconcile.lock", os.O_CREAT | os.O_RDWR, 0o600) + except OSError: + return None + if _try_lock_fd(fd): + return fd + with suppress(OSError): + os.close(fd) + return None + + +def release_session_claim(fd: Optional[int]) -> None: + if fd is not None: + # Closing the descriptor releases the kernel lock. + with suppress(OSError): + os.close(fd) + + +def _try_lock_fd(fd: int) -> bool: + """Non-blocking exclusive lock on an open fd; True if acquired, False if + another process holds it.""" + if IS_WINDOWS: + import msvcrt + + try: + # A 1-byte range lock at offset 0 (allowed past EOF on Windows). + msvcrt.locking( # pyright: ignore[reportAttributeAccessIssue] + fd, + msvcrt.LK_NBLCK, # pyright: ignore[reportAttributeAccessIssue] + 1, + ) + return True + except OSError: + return False + import fcntl + + try: + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + return True + except OSError: + return False + + def claimed_session_name(manifest: dict[str, Any]) -> Optional[str]: """The name this session holds; pre-claim manifests fall back to the endpoint.""" if "name" in manifest: @@ -476,29 +598,29 @@ def claimed_session_name(manifest: dict[str, Any]) -> Optional[str]: 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.""" +def iter_agent_sessions() -> Iterator[EndpointAgentSession]: + """Yields a handle for every session directory under the presets dir.""" root = get_presets_dir() - claims = [] if not root.is_dir(): - return claims + return 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 + if path.is_dir() and not path.name.startswith((".", "models--")): + yield EndpointAgentSession(path=path, timestamp="", debug=False, preset_id=path.name) + + +def find_session_name_claims(name: str) -> list[EndpointAgentSession]: + """Sessions of any status holding `name`, including failed ones.""" + return [ + session + for session in iter_agent_sessions() + if claimed_session_name(session.read_manifest()) == name + ] def 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) + for session in iter_agent_sessions(): + path = session.path manifest = session.read_manifest() status = manifest.get("status") if status not in ("running", "interrupted", "success"): @@ -669,7 +791,7 @@ async def run_endpoint_agent( if session_id is not None: resume_session_id = session_id attempt_prompt = _RESUME_PROMPT - action = "resuming the session" + action = "resuming" else: action = "retrying" print_endpoint_progress( @@ -744,6 +866,8 @@ def print_endpoint_progress(message: str, *, agent_session: EndpointAgentSession timestamp = datetime.now(timezone.utc).astimezone().strftime("%Y-%m-%d %H:%M:%S") message = message.rstrip("\r\n") agent_session.append_log(f"[{timestamp}] {message}") + if not agent_session.echo: + return console.print( Text(f"[{timestamp}]", style="log.time"), Text(message, style="log.message"), @@ -1021,6 +1145,7 @@ async def _session_tailers( redacted_values=redacted_values, offset_store=offset_store, offset_key="runs", + echo=agent_session.echo, ), _RecordMirror( source=workspace.trials_path, @@ -1028,6 +1153,7 @@ async def _session_tailers( redacted_values=redacted_values, offset_store=offset_store, offset_key="trials", + echo=agent_session.echo, ), ] tailer_tasks = [ @@ -1383,6 +1509,7 @@ def __init__( redacted_values: Sequence[str], offset_store: Optional[_OffsetStore] = None, offset_key: str = "", + echo: bool = True, ) -> None: self._source = source self._target = target @@ -1391,6 +1518,7 @@ def __init__( self._offset_key = offset_key self._offset = offset_store.get(offset_key) if offset_store and offset_key else 0 self._enabled = True + self._echo = echo async def run(self) -> None: while True: @@ -1419,7 +1547,8 @@ def flush(self) -> None: f.flush() except OSError as e: self._enabled = False - console.print(f"[warning]Could not mirror {self._target.name}: {e}[/]") + if self._echo: + console.print(f"[warning]Could not mirror {self._target.name}: {e}[/]") def _parse_progress(line: str) -> Optional[str]: diff --git a/src/dstack/_internal/cli/services/endpoints/create.py b/src/dstack/_internal/cli/services/endpoints/create.py index f3b42a7ed..263b57c76 100644 --- a/src/dstack/_internal/cli/services/endpoints/create.py +++ b/src/dstack/_internal/cli/services/endpoints/create.py @@ -3,14 +3,16 @@ import json import os import re +import time import uuid from contextlib import suppress from dataclasses import dataclass from pathlib import Path -from typing import Optional, Sequence +from typing import Any, Optional, Sequence import yaml from rich.table import Table +from rich.text import Text from dstack._internal.cli.models.endpoint_agent import AgentFinalReport from dstack._internal.cli.models.endpoint_presets import EndpointPreset @@ -21,6 +23,7 @@ from dstack._internal.cli.services.endpoints.agent import ( EndpointAgentSession, EndpointAgentWorkspace, + SessionBusyError, attach_agent_workspace, attach_endpoint_agent, build_endpoint_agent_env, @@ -31,13 +34,20 @@ get_claude_auth, get_redacted_values, get_sensitive_inherited_env_values, + iter_agent_sessions, + load_agent_session, load_attachable_agent_session, + mark_session_owner, print_endpoint_progress, + print_session_log, redact, + release_session_claim, remove_agent_workspace, run_endpoint_agent, session_process_alive, + session_report_exists, terminate_agent_process, + try_claim_session, ) from dstack._internal.cli.services.endpoints.presets import endpoint_preset_to_data from dstack._internal.cli.services.endpoints.prompt import get_endpoint_agent_system_prompt @@ -75,72 +85,217 @@ def __init__(self, error: Optional[str]) -> None: self.error = error -def attach_endpoint_preset( +def follow_endpoint_preset( *, api: Client, store: EndpointPresetStore, preset_id: str, keep_service: bool = False, + wait_for_run_stop: bool = True, + echo: bool = True, + claim: bool = False, ) -> EndpointPresetCreateResult: + """Re-owns a detached session: follows its agent to completion, then + verifies and saves the preset (the finalize role, which must run CLI-side + for secret-scrubbing and server-verified preset building). + + `claim=True` takes the exclusive finalize lock so a concurrent `logs -f` and + reconcile can't both finalize the same session. `wait_for_run_stop=False` + and `echo=False` make it non-blocking and silent for reconcile.""" agent_session = load_attachable_agent_session(preset_id) + agent_session.echo = echo + lock = try_claim_session(agent_session) if claim else None + if claim and lock is None: + raise SessionBusyError(f"Preset {preset_id} is being finalized by another process") + try: + configuration = _load_session_configuration(agent_session) + try: + result = asyncio.run( + _create_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, + wait_for_run_stop=wait_for_run_stop, + ) + ) + except KeyboardInterrupt: + # `logs -f` is a viewer: Ctrl+C detaches and leaves the agent + # running (reconcile finalizes it later), never stops it. + _detach_agent_session(agent_session) + raise + except AgentExitedWithoutReport as e: + _stop_active_session_runs(api, agent_session) + _suspend_agent_session(agent_session) + raise CLIError(str(e)) from e + except CLIError: + # Definitive: a failure/invalid report, an unverifiable service, or a + # leaked secret — the preset genuinely cannot be built, so fail it. + _finish_agent_session(agent_session, "failed") + remove_agent_workspace(agent_session) + raise + # A transient error (network / OS) propagates untouched: the completed + # session and its report stay intact for a later follow or reconcile. + _finish_agent_session(agent_session, "success") + remove_agent_workspace(agent_session) + return result + finally: + release_session_claim(lock) + + +def _load_session_configuration(agent_session: EndpointAgentSession) -> EndpointConfiguration: 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" + f"Preset {agent_session.preset_id} has no saved configuration and cannot be" + f" followed; resume it with [code]--resume {agent_session.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( + return 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, + raise CLIError(f"Could not read the preset configuration: {e}") from e + + +def show_endpoint_session_logs( + *, + project: Optional[str], + store: EndpointPresetStore, + preset_id: str, + follow: bool, + keep_service: bool, +) -> Optional[EndpointPresetCreateResult]: + """`logs`: dump a session's log (any status). With `follow`, a still-live + session is re-owned, followed to completion, and its preset saved; a + finished session just prints its log. Returns the saved preset, if any.""" + session = load_agent_session(preset_id) + status = session.read_manifest().get("status") + if not follow or status in ("success", "failed", "interrupted"): + print_session_log(session) + if follow and status == "interrupted": + console.print( + f"\nPreset [code]{preset_id}[/] creation was interrupted; resume it with" + f" [code]dstack preset create -f --resume {preset_id}[/]." ) + return None + # Following a live session: print the log so far, then stream new progress + # (a future --since could bound this). The client is built only here, so a + # read-only dump never needs a server or authentication. + print_session_log(session) + try: + return follow_endpoint_preset( + api=Client.from_config(project_name=project), + store=store, + preset_id=preset_id, + keep_service=keep_service, + claim=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 + except SessionBusyError: + # Another CLI already owns the finalize; follow read-only instead of + # refusing, so any number of viewers can watch the same preset at once. + _follow_session_log_readonly(session) + return None -def stop_endpoint_session(api: Client, preset_id: str, *, assume_yes: bool = False) -> None: +def _follow_session_log_readonly(session: EndpointAgentSession) -> None: + """Read-only follow: another CLI owns the finalize, so just stream the log it + writes until the preset reaches a terminal state.""" + try: + offset = session.log_path.stat().st_size + except OSError: + offset = 0 + while True: + try: + with session.log_path.open("r", encoding="utf-8", errors="replace") as f: + f.seek(offset) + chunk = f.read() + offset = f.tell() + except OSError: + chunk = "" + if chunk: + console.print(Text(chunk.rstrip("\n")), soft_wrap=True) + if session.read_manifest().get("status") in ("success", "failed", "interrupted"): + return + time.sleep(1) + + +def reconcile_detached_sessions(store: EndpointPresetStore) -> None: + """Finalizes sessions whose agent completed while no CLI was attached + (graceful detach, or an ungraceful CLI death). This is what makes the saved + preset independent of a foreground process: any read command runs it, and + the work materializes from the on-disk report. + + Best-effort and parallel-safe — finalize takes an exclusive claim, and every + error is swallowed so the calling read command never fails. + """ + for session in iter_agent_sessions(): + if _is_reconcilable(session.read_manifest()): + _reconcile_session(session, store) + + +def _is_reconcilable(manifest: dict[str, Any]) -> bool: + # An orphaned session (no live owner) whose agent left a completion report. + # A session interrupted mid-work has no report and stays resumable; one + # stopped *after* the agent finished is finalized by `stop` itself, not here. + # Sessions created before finalize context was persisted lack `project` and + # are skipped — they finalize interactively via `logs -f`. + return ( + manifest.get("status") == "running" + and bool(manifest.get("project")) + and session_report_exists(manifest) + and not session_process_alive(manifest) + ) + + +def _reconcile_session(session: EndpointAgentSession, store: EndpointPresetStore) -> None: + manifest = session.read_manifest() + try: + api = Client.from_config(project_name=str(manifest.get("project") or "")) + except Exception: # noqa: BLE001 — offline/misconfigured must not break the read command + return + # follow_endpoint_preset takes the finalize claim (so a concurrent `logs -f` + # or reconcile can't double-finalize), records the terminal status itself, + # and leaves the session intact on a transient error. Every outcome is silent + # here — the result shows in the list that follows. + with suppress(Exception): + follow_endpoint_preset( + api=api, + store=store, + preset_id=session.preset_id, + keep_service=bool(manifest.get("keep_service")), + wait_for_run_stop=False, + echo=False, + claim=True, + ) + + +def stop_endpoint_session(api: Client, preset_id: str) -> 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) + manifest = session.read_manifest() + # If the agent already finished, there is nothing to stop. Leave the session + # for a read to save (reconcile-on-read) rather than discarding it as + # interrupted — but `stop` itself never saves; it only stops. + if session_report_exists(manifest) and not session_process_alive(manifest): + console.print(f"Preset [code]{preset_id}[/] has already finished.") + return + terminate_agent_process(manifest) + _stop_active_session_runs(api, session) + _suspend_agent_session(session) -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.""" +def _stop_active_session_runs(api: Client, session: EndpointAgentSession) -> None: + """Stops the session's non-terminal runs (with a spinner), like `dstack + stop`. Keeping a trial instance warm for resume is the detach path, not this.""" try: lines = session.runs_path.read_text(encoding="utf-8").splitlines() except OSError: - return False + return names = [] for line in lines: try: @@ -158,15 +313,9 @@ def _stop_active_session_runs( 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 + return with console.status("Stopping runs..."): api.client.runs.stop(api.project, active, abort=False) - return True def _resolve_endpoint_env_best_effort( @@ -239,6 +388,7 @@ async def _create_endpoint_preset( agent_session: EndpointAgentSession, resume: bool = False, attach: bool = False, + wait_for_run_stop: bool = True, user_prompt: Optional[str] = None, allowed_fleets: Optional[tuple[str, ...]] = None, ) -> EndpointPresetCreateResult: @@ -254,7 +404,7 @@ async def _create_endpoint_preset( 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" + "The configuration prompt is ignored when resuming: the preset keeps its original prompt" ) user_prompt = pinned_prompt auth = get_claude_auth() @@ -272,16 +422,20 @@ async def _create_endpoint_preset( 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") + raise CLIError("The project has no fleets. Create one before creating a preset") 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()) + # Record ownership + the finalize context (project, keep-service) so a later + # detached reconcile can complete this session from disk alone. + mark_session_owner( + agent_session, + project=api.project, + keep_service=keep_service, + claude_model=auth.model if auth is not None else None, + ) endpoint_env = configuration.env.as_dict() token = getattr(api.client, "_token", None) @@ -361,10 +515,6 @@ async def _create_endpoint_preset( 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 @@ -386,6 +536,7 @@ async def _create_endpoint_preset( final_run_name=report.run_name if report is not None else None, keep_final_service=keep_final_service, agent_session=agent_session, + wait_for_stop=wait_for_run_stop, ) except Exception as e: cleanup_error = str(e) @@ -397,10 +548,15 @@ async def _create_endpoint_preset( workspace=workspace, final_run_name=report.run_name if report is not None else None, agent_session=agent_session, + wait_for_stop=wait_for_run_stop, ) if cleanup_error is not None: - raise CLIError(f"Failed to clean up preset creation runs: {cleanup_error}") + # The preset is already saved by this point; a failed cleanup only means + # trial runs may still be running. Warn rather than fail the (successful) + # session — otherwise a transient blip would discard completed work. + if agent_session.echo: + warn(f"Failed to stop preset creation runs: {cleanup_error}") assert preset is not None assert preset_path is not None assert report is not None @@ -431,63 +587,54 @@ def _finish_agent_session( status: str, ) -> None: try: - path = session.finish(status) + 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'}[/]") + if session.echo: + warn(f"Could not finalize agent output. Files remain at {session.path}: {e}") + + +def _detach_agent_session(session: EndpointAgentSession) -> None: + """Releases ownership but leaves the agent running — it stays visible and + reconcilable in `dstack preset`. Silent: `logs -f` calls this on Ctrl+C, and + a viewer that just stops watching shouldn't announce anything.""" + session.update_manifest(pid=None) def _stop_or_detach_agent_session( session: EndpointAgentSession, api: Optional[Client] = None ) -> None: - """Apply-style interrupt: stop the session, or detach and leave the agent + """`create` interrupt: stop the session, or detach and leave the agent working — it stays visible as a running session in `dstack preset`.""" manifest = session.read_manifest() agent_alive = session_process_alive({**manifest, "pid": None}) stop = True if agent_alive: try: - stop = confirm_ask(f"Stop the preset creation session [code]{session.preset_id}[/]?") + stop = confirm_ask(f"Stop creating preset [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) + if not stop: + _detach_agent_session(session) + console.print( + f"\nDetached. Follow with [code]dstack preset logs -f {session.preset_id}[/]" + f" or stop with [code]dstack preset stop {session.preset_id}[/]." + ) return - 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}[/]." - ) + terminate_agent_process(manifest) + if api is not None: + _stop_active_session_runs(api, session) + _suspend_agent_session(session) -def _suspend_agent_session( - session: EndpointAgentSession, *, runs_left_active: bool = True -) -> None: +def _suspend_agent_session(session: EndpointAgentSession) -> 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}[/]." - ) + warn(f"Could not record the interrupted preset state: {e}") + console.print(f"\nPreset [code]{session.preset_id}[/] creation interrupted.") + console.print( + f"Resume it with [code]dstack preset create -f --resume {session.preset_id}[/]." + ) def _get_build_name(name: Optional[str], model_name: str, suffix: str) -> str: @@ -511,10 +658,10 @@ 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 + raise CLIError(f"The preset creation constraints cannot be read: {e}") from e prefix = data.get("run_name_prefix") if isinstance(data, dict) else None if not isinstance(prefix, str) or not prefix: - raise CLIError("The session constraints do not contain a run name prefix") + raise CLIError("The preset creation constraints do not contain a run name prefix") return prefix @@ -523,7 +670,7 @@ def plan_endpoint_preset(*, api: Client, configuration: EndpointConfiguration) - 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") + raise CLIError("The project has no fleets. Create one before creating a preset") _print_fleet_offers(api, allowed_fleets) return allowed_fleets @@ -617,6 +764,7 @@ async def _cleanup_runs( final_run_name: Optional[str], agent_session: EndpointAgentSession, keep_final_service: bool = False, + wait_for_stop: bool = True, ) -> None: run_names = _load_submitted_run_names(workspace.runs_path) if final_run_name is not None: @@ -634,6 +782,10 @@ async def _cleanup_runs( if not active_names: return api.client.runs.stop(api.project, active_names, abort=False) + if not wait_for_stop: + # Background reconcile issues the stop but must not block the read + # command it runs inside; the runs terminate on the server regardless. + return deadline = asyncio.get_running_loop().time() + _RUN_STOP_TIMEOUT_SECONDS pending = set(active_names) while pending: diff --git a/src/tests/_internal/cli/commands/test_endpoint.py b/src/tests/_internal/cli/commands/test_endpoint.py index e6216e0b5..0fac79b71 100644 --- a/src/tests/_internal/cli/commands/test_endpoint.py +++ b/src/tests/_internal/cli/commands/test_endpoint.py @@ -93,7 +93,9 @@ def test_handles_keyboard_interrupt(self, tmp_path, capsys): ) assert exit_code == 0 - assert "Operation interrupted by user" in capsys.readouterr().out + # create reports its own Ctrl+C outcome (Detached / interrupted) via the + # interrupt handler, so the command layer stays silent — no generic line. + assert "Operation interrupted by user" not in capsys.readouterr().out def test_lists_and_deletes_preset_without_api_client(self, tmp_path, capsys): preset = get_endpoint_preset() diff --git a/src/tests/_internal/cli/services/endpoints/test_agent.py b/src/tests/_internal/cli/services/endpoints/test_agent.py index a8fabd89b..a4e7c3b4e 100644 --- a/src/tests/_internal/cli/services/endpoints/test_agent.py +++ b/src/tests/_internal/cli/services/endpoints/test_agent.py @@ -549,7 +549,7 @@ async def test_resumes_after_connection_error(self, tmp_path, monkeypatch, capsy 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 + assert "resuming" in capsys.readouterr().out @pytest.mark.asyncio async def test_resumes_on_any_unreported_death(self, tmp_path, monkeypatch): @@ -775,7 +775,7 @@ def test_treats_dead_running_session_as_resumable(self, tmp_path, monkeypatch): 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"): + with pytest.raises(CLIError, match="Unknown preset"): monkeypatch.setenv("HOME", str(tmp_path)) monkeypatch.setenv("USERPROFILE", str(tmp_path)) load_resumable_agent_session("00000000") @@ -802,11 +802,11 @@ def test_refusals(self, tmp_path, monkeypatch): "dstack._internal.cli.services.endpoints.agent.psutil.pid_exists", lambda pid: True, ) - with pytest.raises(CLIError, match="running already"): + with pytest.raises(CLIError, match="still being created"): 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"): + with pytest.raises(CLIError, match="stopped before it started"): load_resumable_agent_session("aa000004") diff --git a/src/tests/_internal/cli/services/endpoints/test_create.py b/src/tests/_internal/cli/services/endpoints/test_create.py index d8a64ecfc..d49b1b76b 100644 --- a/src/tests/_internal/cli/services/endpoints/test_create.py +++ b/src/tests/_internal/cli/services/endpoints/test_create.py @@ -1,5 +1,6 @@ import asyncio import json +import os import uuid from types import SimpleNamespace @@ -12,11 +13,18 @@ EndpointAgentSession, EndpointAgentWorkspace, create_agent_workspace, + load_agent_session, + mark_session_owner, print_endpoint_progress, + print_session_log, + release_session_claim, remove_agent_workspace, + session_process_alive, + try_claim_session, ) from dstack._internal.cli.services.endpoints.create import ( EndpointPresetCreateResult, + SessionBusyError, _build_constraints, _cleanup_runs, _create_endpoint_preset, @@ -24,8 +32,9 @@ _print_fleet_offers, _save_final_report_copy, _stop_active_session_runs, - attach_endpoint_preset, create_endpoint_preset, + follow_endpoint_preset, + reconcile_detached_sessions, ) from dstack._internal.cli.services.endpoints.store import EndpointPresetStore from dstack._internal.core.errors import CLIError @@ -138,8 +147,6 @@ async def create(**kwargs): 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 def test_debug_finalization_error_does_not_mask_success(self, tmp_path, monkeypatch, capsys): monkeypatch.setenv("HOME", str(tmp_path)) @@ -238,7 +245,7 @@ async def test_checks_active_fleets_before_claude_auth(self, tmp_path, monkeypat lambda: pytest.fail("Claude auth must not be checked without an active fleet"), ) - with pytest.raises(CLIError, match="no active fleets"): + with pytest.raises(CLIError, match="no fleets"): await _create_endpoint_preset( api=api, configuration=EndpointConfiguration( @@ -527,7 +534,7 @@ async def create(**kwargs): manifest = json.loads((sessions[0] / "session.json").read_text()) assert manifest["status"] == "interrupted" output = capsys.readouterr().out - assert "Resume with" in output + assert "--resume" in output assert sessions[0].name in output @pytest.mark.asyncio @@ -671,7 +678,45 @@ def test_no_offers_shows_the_shared_warning_without_failing(self, capsys): assert "No matching instance offers available" in out -class TestAttachEndpointPreset: +class TestSessionLog: + def _session(self, tmp_path, preset_id: str, status: str, log: str) -> EndpointAgentSession: + session_dir = tmp_path / preset_id + session_dir.mkdir() + (session_dir / "agent.log").write_text(log) + session = EndpointAgentSession( + path=session_dir, timestamp="t", debug=False, preset_id=preset_id + ) + session.update_manifest(status=status) + return session + + def test_load_agent_session_reads_any_status(self, tmp_path, monkeypatch): + self._session(tmp_path, "dead0000", "failed", "[t] boom\n") + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.agent.get_presets_dir", + lambda: tmp_path, + ) + # A failed session is off-limits to follow/resume, but its log is readable. + session = load_agent_session("dead0000") + assert session.preset_id == "dead0000" + with pytest.raises(CLIError, match="Unknown preset"): + load_agent_session("nope0000") + + def test_print_session_log_dumps_log_verbatim(self, tmp_path, monkeypatch, capsys): + session = self._session( + tmp_path, "abcd0000", "success", "[t] trial 1 done\n[t] saved preset\n" + ) + print_session_log(session) + out = capsys.readouterr().out + assert "trial 1 done" in out + assert "saved preset" in out + + def test_print_session_log_notes_empty_log(self, tmp_path, capsys): + session = self._session(tmp_path, "empty000", "running", "") + print_session_log(session) + assert "No log output yet" in capsys.readouterr().out + + +class TestFollowEndpointPreset: def _detached_session(self, tmp_path, configuration_yaml: str) -> EndpointAgentSession: session_dir = tmp_path / "ab12cd34" session_dir.mkdir() @@ -708,7 +753,7 @@ async def fake_attach(**kwargs): fake_attach, ) - result = attach_endpoint_preset( + result = follow_endpoint_preset( api=creation_context.api, store=creation_context.store, preset_id="ab12cd34", @@ -738,7 +783,7 @@ async def fake_attach(**kwargs): ) with pytest.raises(CLIError, match="agent died"): - attach_endpoint_preset( + follow_endpoint_preset( api=creation_context.api, store=creation_context.store, preset_id="ab12cd34", @@ -746,6 +791,28 @@ async def fake_attach(**kwargs): assert session.read_manifest()["status"] == "interrupted" + def test_backs_off_when_claim_is_held(self, creation_context, monkeypatch, tmp_path): + session = self._detached_session( + tmp_path, "type: preset\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\n" + ) + # Another live holder owns the finalize lock (reconcile or logs -f). + held = try_claim_session(session) + assert held is not None + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.create.load_attachable_agent_session", + lambda preset_id: session, + ) + # claim=True must refuse rather than double-finalize; the session is untouched. + with pytest.raises(SessionBusyError): + follow_endpoint_preset( + api=creation_context.api, + store=creation_context.store, + preset_id="ab12cd34", + claim=True, + ) + assert session.read_manifest()["status"] == "running" + release_session_claim(held) + class TestStopActiveSessionRuns: def _session(self, tmp_path) -> EndpointAgentSession: @@ -778,18 +845,169 @@ def test_stops_only_non_terminal_runs(self, tmp_path): {"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 + # No per-run prompt: active runs are stopped automatically (like dstack stop). + _stop_active_session_runs(api, self._session(tmp_path)) 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 + def test_stops_nothing_when_all_terminal(self, tmp_path): 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 - ) + api = self._api({"qwen-build-1": RunStatus.DONE, "qwen-build-2": RunStatus.DONE}, stopped) - assert _stop_active_session_runs(api, self._session(tmp_path), assume_yes=False) is False + _stop_active_session_runs(api, self._session(tmp_path)) assert stopped == [] + + +class TestReconcileDetachedSessions: + def _session_dir( + self, + tmp_path, + preset_id="dead0001", + *, + status="running", + project="main", + with_report=True, + keep_service=False, + owner_alive=False, + ): + session_dir = tmp_path / preset_id + (session_dir / "workspace" / "w").mkdir(parents=True) + manifest = { + "id": preset_id, + "status": status, + "keep_service": keep_service, + # A dead pid unless a live owner is requested below. + "pid": 987654321, + "pid_started_at": 0.0, + "workspace": str(session_dir / "workspace"), + } + if project is not None: + manifest["project"] = project + if owner_alive: + # A live pid with no recorded start time reads as an active owner. + manifest["agent_pid"] = os.getpid() + (session_dir / "session.json").write_text(json.dumps(manifest)) + if with_report: + (session_dir / "workspace" / "w" / "final_report.json").write_text("{}") + return session_dir + + def _patch(self, monkeypatch, tmp_path, follow): + # reconcile iterates via agent.iter_agent_sessions -> agent.get_presets_dir. + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.agent.get_presets_dir", + lambda: tmp_path, + ) + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.create.Client", + SimpleNamespace(from_config=lambda project_name=None: SimpleNamespace()), + ) + monkeypatch.setattr( + "dstack._internal.cli.services.endpoints.create.follow_endpoint_preset", + follow, + ) + + def _recording_follow(self, calls): + def follow(**kwargs): + calls.append(kwargs) + return SimpleNamespace( + preset=SimpleNamespace(id=kwargs["preset_id"], base="Qwen/Qwen3.5-27B") + ) + + return follow + + def test_finalizes_eligible_detached_session(self, tmp_path, monkeypatch): + self._session_dir(tmp_path, keep_service=True) + calls: list = [] + self._patch(monkeypatch, tmp_path, self._recording_follow(calls)) + reconcile_detached_sessions(EndpointPresetStore(tmp_path / "store")) + assert len(calls) == 1 + assert calls[0]["preset_id"] == "dead0001" + # Honors persisted keep-service; non-interactive, non-blocking, silent, + # and under the finalize claim (parallel-safe). + assert calls[0]["keep_service"] is True + assert calls[0]["wait_for_run_stop"] is False + assert calls[0]["echo"] is False + assert calls[0]["claim"] is True + + @pytest.mark.parametrize( + "kwargs", + [ + {"status": "success"}, + {"status": "failed"}, + {"with_report": False}, + {"project": None}, + {"owner_alive": True}, + # `interrupted` is `stop`'s job now, not reconcile's — with or without a report. + {"status": "interrupted"}, + {"status": "interrupted", "with_report": False}, + ], + ) + def test_skips_ineligible(self, tmp_path, monkeypatch, kwargs): + self._session_dir(tmp_path, **kwargs) + calls: list = [] + self._patch(monkeypatch, tmp_path, self._recording_follow(calls)) + reconcile_detached_sessions(EndpointPresetStore(tmp_path / "store")) + assert calls == [] + + def test_never_raises_when_finalize_fails(self, tmp_path, monkeypatch): + self._session_dir(tmp_path) + + def boom(**kwargs): + raise RuntimeError("finalize blew up") + + self._patch(monkeypatch, tmp_path, boom) + # A read command must never fail because reconcile did. + reconcile_detached_sessions(EndpointPresetStore(tmp_path / "store")) + + +class TestSessionClaim: + def _session(self, tmp_path): + (tmp_path / "sess").mkdir() + return EndpointAgentSession( + path=tmp_path / "sess", timestamp="", debug=False, preset_id="sess" + ) + + def test_claim_is_exclusive_and_releasable(self, tmp_path): + session = self._session(tmp_path) + first = try_claim_session(session) + assert first is not None + assert try_claim_session(session) is None # the kernel lock is held + release_session_claim(first) + again = try_claim_session(session) + assert again is not None + release_session_claim(again) + + def test_claim_acquires_when_lock_file_is_unheld(self, tmp_path): + session = self._session(tmp_path) + # A leftover lock file from a crashed run holds no kernel lock: the file's + # presence must not block a new claim (no stale-lock reasoning needed). + (session.path / ".reconcile.lock").write_text("stale") + fd = try_claim_session(session) + assert fd is not None + release_session_claim(fd) + + +class TestSessionProcessAlive: + def test_recycled_pid_with_stale_start_time_is_not_alive(self): + # A live pid whose recorded start time does not match — the pid was recycled. + assert session_process_alive({"agent_pid": os.getpid(), "agent_started_at": 0.0}) is False + + def test_dead_pids_are_not_alive(self): + assert session_process_alive({"pid": 987654321, "pid_started_at": 0.0}) is False + assert session_process_alive({}) is False + + +class TestMarkSessionOwner: + def test_persists_finalize_context(self, tmp_path): + (tmp_path / "s").mkdir() + session = EndpointAgentSession( + path=tmp_path / "s", timestamp="", debug=False, preset_id="s" + ) + session.update_manifest(status="running") + mark_session_owner(session, project="main", keep_service=True) + manifest = session.read_manifest() + assert manifest["project"] == "main" + assert manifest["keep_service"] is True + assert manifest["pid"] == os.getpid() + assert "pid_started_at" in manifest From c911c2a6e8c00e58b00f3ccdb6dd56b94fa3e774 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Wed, 22 Jul 2026 11:31:49 -0700 Subject: [PATCH 14/25] =?UTF-8?q?Rename=20endpoint=E2=86=92preset=20in=20t?= =?UTF-8?q?he=20CLI;=20require=20max=5Ftrials=20on=20create?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rename the preset feature's internal identifiers, modules, and directories from endpoint/Endpoint to preset/Preset. CLI-only — no server, core, or API code is involved (that design was already removed). Notable moves: cli/services/endpoints -> cli/services/presets, cli/commands/endpoint.py -> preset.py, and the models split into configurations.py (PresetConfiguration) and presets.py (Preset). Also drops the redundant `endpoint` session-manifest key, which held the same value as `name`. Require max_trials for preset creation instead of silently defaulting to 3: create now errors unless it is set in the configuration or passed via --max-trials. The model field stays optional so apply and partial configs are unaffected. Trim the saved-preset line to `Preset saved`. Co-Authored-By: Claude Opus 4.8 --- .../cli/commands/{endpoint.py => preset.py} | 102 ++++---- src/dstack/_internal/cli/main.py | 4 +- .../{endpoints.py => configurations.py} | 33 ++- .../{endpoint_agent.py => preset_agent.py} | 8 +- .../{endpoint_trials.py => preset_trials.py} | 6 +- .../{endpoint_presets.py => presets.py} | 40 ++-- .../{endpoints => presets}/__init__.py | 0 .../services/{endpoints => presets}/agent.py | 124 +++++----- .../services/{endpoints => presets}/apply.py | 42 ++-- .../services/{endpoints => presets}/create.py | 170 +++++++------- .../services/{endpoints => presets}/output.py | 16 +- .../{endpoints => presets}/presets.py | 32 +-- .../services/{endpoints => presets}/prompt.py | 2 +- .../resources/system_prompt.md | 0 .../services/{endpoints => presets}/store.py | 44 ++-- .../services/{endpoints => presets}/verify.py | 64 ++--- .../{test_endpoint.py => test_preset.py} | 136 ++++++----- ...st_endpoints.py => test_configurations.py} | 48 ++-- ...ndpoint_presets.py => preset_factories.py} | 40 ++-- .../{endpoints => presets}/__init__.py | 0 .../{endpoints => presets}/test_agent.py | 116 ++++----- .../{endpoints => presets}/test_apply.py | 46 ++-- .../{endpoints => presets}/test_create.py | 221 +++++++++--------- .../{endpoints => presets}/test_output.py | 14 +- .../{endpoints => presets}/test_prompt.py | 18 +- .../{endpoints => presets}/test_store.py | 106 +++++---- .../{endpoints => presets}/test_trials.py | 12 +- .../{endpoints => presets}/test_verify.py | 52 ++--- 28 files changed, 752 insertions(+), 744 deletions(-) rename src/dstack/_internal/cli/commands/{endpoint.py => preset.py} (85%) rename src/dstack/_internal/cli/models/{endpoints.py => configurations.py} (88%) rename src/dstack/_internal/cli/models/{endpoint_agent.py => preset_agent.py} (95%) rename src/dstack/_internal/cli/models/{endpoint_trials.py => preset_trials.py} (77%) rename src/dstack/_internal/cli/models/{endpoint_presets.py => presets.py} (87%) rename src/dstack/_internal/cli/services/{endpoints => presets}/__init__.py (100%) rename src/dstack/_internal/cli/services/{endpoints => presets}/agent.py (93%) rename src/dstack/_internal/cli/services/{endpoints => presets}/apply.py (83%) rename src/dstack/_internal/cli/services/{endpoints => presets}/create.py (86%) rename src/dstack/_internal/cli/services/{endpoints => presets}/output.py (94%) rename src/dstack/_internal/cli/services/{endpoints => presets}/presets.py (91%) rename src/dstack/_internal/cli/services/{endpoints => presets}/prompt.py (94%) rename src/dstack/_internal/cli/services/{endpoints => presets}/resources/system_prompt.md (100%) rename src/dstack/_internal/cli/services/{endpoints => presets}/store.py (82%) rename src/dstack/_internal/cli/services/{endpoints => presets}/verify.py (75%) rename src/tests/_internal/cli/commands/{test_endpoint.py => test_preset.py} (77%) rename src/tests/_internal/cli/models/{test_endpoints.py => test_configurations.py} (53%) rename src/tests/_internal/cli/{endpoint_presets.py => preset_factories.py} (81%) rename src/tests/_internal/cli/services/{endpoints => presets}/__init__.py (100%) rename src/tests/_internal/cli/services/{endpoints => presets}/test_agent.py (89%) rename src/tests/_internal/cli/services/{endpoints => presets}/test_apply.py (77%) rename src/tests/_internal/cli/services/{endpoints => presets}/test_create.py (82%) rename src/tests/_internal/cli/services/{endpoints => presets}/test_output.py (91%) rename src/tests/_internal/cli/services/{endpoints => presets}/test_prompt.py (70%) rename src/tests/_internal/cli/services/{endpoints => presets}/test_store.py (67%) rename src/tests/_internal/cli/services/{endpoints => presets}/test_trials.py (83%) rename src/tests/_internal/cli/services/{endpoints => presets}/test_verify.py (74%) diff --git a/src/dstack/_internal/cli/commands/endpoint.py b/src/dstack/_internal/cli/commands/preset.py similarity index 85% rename from src/dstack/_internal/cli/commands/endpoint.py rename to src/dstack/_internal/cli/commands/preset.py index 35731d4c1..5b4620e4c 100644 --- a/src/dstack/_internal/cli/commands/endpoint.py +++ b/src/dstack/_internal/cli/commands/preset.py @@ -7,31 +7,31 @@ from argcomplete import FilesCompleter # type: ignore[attr-defined] from dstack._internal.cli.commands import BaseCommand -from dstack._internal.cli.models.endpoint_presets import ( - EndpointPreset, - EndpointPresetListOutput, +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.models.presets import ( + Preset, + PresetListOutput, ) -from dstack._internal.cli.models.endpoints import EndpointConfiguration from dstack._internal.cli.services.completion import ProjectNameCompleter -from dstack._internal.cli.services.endpoints.agent import ( +from dstack._internal.cli.services.presets.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, - plan_endpoint_preset, +from dstack._internal.cli.services.presets.apply import apply_preset +from dstack._internal.cli.services.presets.create import ( + create_preset, + plan_preset, reconcile_detached_sessions, - show_endpoint_session_logs, - stop_endpoint_session, + show_preset_session_logs, + stop_preset_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.presets.output import print_presets +from dstack._internal.cli.services.presets.store import ( + PresetStore, + load_preset_configuration, + resolve_preset_prompt, ) from dstack._internal.cli.services.profile import ( apply_profile_args, @@ -39,13 +39,13 @@ register_profile_args, ) from dstack._internal.cli.utils.common import confirm_ask, console -from dstack._internal.core.errors import CLIError +from dstack._internal.core.errors import CLIError, ConfigurationError from dstack._internal.core.models.profiles import ProfileParams from dstack._internal.core.services import is_valid_dstack_resource_name from dstack.api import Client -class EndpointCommand(BaseCommand): +class PresetCommand(BaseCommand): NAME = "preset" DESCRIPTION = "Manage model serving presets" @@ -213,20 +213,20 @@ def _reconcile(self) -> None: so a saved preset never depends on a foreground CLI surviving. Fully best-effort — it must never make a read command fail.""" with suppress(Exception): - reconcile_detached_sessions(EndpointPresetStore()) + reconcile_detached_sessions(PresetStore()) def _list(self, args: argparse.Namespace) -> None: base = getattr(args, "base", None) repo = getattr(args, "repo", None) if getattr(args, "json", False): self._reconcile() - presets = _filter_presets(EndpointPresetStore().list(), base=base, repo=repo) - print(EndpointPresetListOutput(presets=presets).json()) + presets = _filter_presets(PresetStore().list(), base=base, repo=repo) + print(PresetListOutput(presets=presets).json()) return verbose = getattr(args, "verbose", False) while True: self._reconcile() - presets = EndpointPresetStore().list() + presets = PresetStore().list() sessions = list_agent_sessions() if base or repo: repo_to_base = {preset.model: preset.base for preset in presets} @@ -240,16 +240,16 @@ def _list(self, args: argparse.Namespace) -> None: ] if getattr(args, "watch", False): console.clear() - print_endpoint_presets(presets, sessions=sessions, verbose=verbose) + print_presets(presets, sessions=sessions, verbose=verbose) if not getattr(args, "watch", False): return time.sleep(5) def _create(self, args: argparse.Namespace) -> None: - configuration_path, configuration = load_endpoint_configuration(args.configuration_file) + configuration_path, configuration = load_preset_configuration(args.configuration_file) configuration = _get_effective_configuration(configuration, args, require_name=False) - user_prompt = resolve_endpoint_prompt(configuration, configuration_path) - store = EndpointPresetStore() + user_prompt = resolve_preset_prompt(configuration, configuration_path) + store = PresetStore() resume_session = None if getattr(args, "resume", None): resume_session = load_resumable_agent_session(args.resume) @@ -261,12 +261,16 @@ def _create(self, args: argparse.Namespace) -> None: 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 configuration.max_trials is None: + raise ConfigurationError( + "max_trials is required. Set it in the configuration or pass --max-trials" + ) + allowed_fleets = plan_preset(api=api, configuration=configuration) if not _confirm_preset_creation(store, configuration.name, assume_yes=args.yes): console.print("\nExiting...") return try: - result = create_endpoint_preset( + result = create_preset( api=api, configuration=configuration, store=store, @@ -278,18 +282,15 @@ def _create(self, args: argparse.Namespace) -> None: ) except KeyboardInterrupt: return # the interrupt handler already reported detach / stop - console.print( - f"Preset [code]{result.preset.id}[/] for " - f"[code]{result.preset.base}[/] saved to [code]{result.path}[/]" - ) + console.print(f"Preset [code]{result.preset.id}[/] saved") if args.keep_service: console.print(f"Final service [code]{result.final_run_name}[/] kept running") def _logs(self, args: argparse.Namespace) -> None: try: - result = show_endpoint_session_logs( + result = show_preset_session_logs( project=args.project, - store=EndpointPresetStore(), + store=PresetStore(), preset_id=_resolve_session_ref(args.preset), follow=args.follow, keep_service=args.keep_service, @@ -297,10 +298,7 @@ def _logs(self, args: argparse.Namespace) -> None: except KeyboardInterrupt: return # a log viewer: Ctrl+C just stops watching, quietly if result is not None: - console.print( - f"Preset [code]{result.preset.id}[/] for " - f"[code]{result.preset.base}[/] saved to [code]{result.path}[/]" - ) + console.print(f"Preset [code]{result.preset.id}[/] saved") if args.keep_service: console.print(f"Final service [code]{result.final_run_name}[/] kept running") @@ -309,11 +307,11 @@ def _stop(self, args: argparse.Namespace) -> None: if not args.yes and not confirm_ask(f"Stop creating preset [code]{preset_id}[/]?"): console.print("\nExiting...") return - stop_endpoint_session(Client.from_config(project_name=args.project), preset_id) + stop_preset_session(Client.from_config(project_name=args.project), preset_id) def _get(self, args: argparse.Namespace) -> None: self._reconcile() - store = EndpointPresetStore() + store = PresetStore() preset = store.get(args.preset) or store.find_by_name(args.preset) if preset is None: raise CLIError(f"Preset {args.preset!r} does not exist") @@ -321,20 +319,20 @@ def _get(self, args: argparse.Namespace) -> None: def _apply(self, args: argparse.Namespace) -> None: self._reconcile() - configuration_path, configuration = load_endpoint_configuration(args.configuration_file) + configuration_path, configuration = load_preset_configuration(args.configuration_file) configuration = _get_effective_configuration(configuration, args) - apply_endpoint_preset( + apply_preset( api=Client.from_config(project_name=args.project), configuration=configuration, configuration_path=configuration_path, preset_ids=args.preset_ids, profile_name=args.profile, command_args=args, - store=EndpointPresetStore(), + store=PresetStore(), ) def _delete(self, args: argparse.Namespace) -> None: - store = EndpointPresetStore() + store = PresetStore() if args.preset is not None: preset = store.get(args.preset) or store.find_by_name(args.preset) if preset is None: @@ -411,11 +409,11 @@ def _add_list_args(parser: argparse.ArgumentParser) -> None: def _filter_presets( - presets: list[EndpointPreset], + presets: list[Preset], *, base: str | None, repo: str | None, -) -> list[EndpointPreset]: +) -> list[Preset]: return [ preset for preset in presets @@ -438,7 +436,7 @@ def _session_matches_model( return model == base or repo_to_base.get(model) == base -def _apply_name(configuration: EndpointConfiguration, name: str | None, *, required: bool) -> None: +def _apply_name(configuration: PresetConfiguration, name: str | None, *, required: bool) -> None: if name is not None: configuration.name = name if configuration.name is None: @@ -461,9 +459,7 @@ def _resolve_session_ref(ref: str) -> str: return ref -def _confirm_preset_creation( - store: EndpointPresetStore, name: str | None, *, assume_yes: bool -) -> bool: +def _confirm_preset_creation(store: PresetStore, name: str | None, *, assume_yes: bool) -> bool: """One apply-style confirmation; reassigns the name from any holder on yes.""" preset_holder = None session_holders = [] @@ -497,11 +493,11 @@ def _confirm_preset_creation( def _get_effective_configuration( - configuration: EndpointConfiguration, + configuration: PresetConfiguration, args: argparse.Namespace, *, require_name: bool = True, -) -> EndpointConfiguration: +) -> PresetConfiguration: _apply_name(configuration, args.name, required=require_name) if getattr(args, "max_trials", None) is not None: configuration.max_trials = args.max_trials @@ -510,4 +506,4 @@ def _get_effective_configuration( if getattr(configuration, field) is None: setattr(configuration, field, getattr(profile, field)) apply_profile_args(args, configuration) - return EndpointConfiguration.parse_obj(configuration.dict()) + return PresetConfiguration.parse_obj(configuration.dict()) diff --git a/src/dstack/_internal/cli/main.py b/src/dstack/_internal/cli/main.py index 335f7693f..e21d5d71a 100644 --- a/src/dstack/_internal/cli/main.py +++ b/src/dstack/_internal/cli/main.py @@ -8,7 +8,6 @@ from dstack._internal.cli.commands.attach import AttachCommand from dstack._internal.cli.commands.completion import CompletionCommand from dstack._internal.cli.commands.delete import DeleteCommand -from dstack._internal.cli.commands.endpoint import EndpointCommand from dstack._internal.cli.commands.event import EventCommand from dstack._internal.cli.commands.export import ExportCommand from dstack._internal.cli.commands.fleet import FleetCommand @@ -19,6 +18,7 @@ from dstack._internal.cli.commands.logs import LogsCommand from dstack._internal.cli.commands.metrics import MetricsCommand from dstack._internal.cli.commands.offer import OfferCommand +from dstack._internal.cli.commands.preset import PresetCommand from dstack._internal.cli.commands.project import ProjectCommand from dstack._internal.cli.commands.ps import PsCommand from dstack._internal.cli.commands.run import RunCommand @@ -68,7 +68,7 @@ def main(): ApplyCommand.register(subparsers) AttachCommand.register(subparsers) DeleteCommand.register(subparsers) - EndpointCommand.register(subparsers) + PresetCommand.register(subparsers) EventCommand.register(subparsers) ExportCommand.register(subparsers) FleetCommand.register(subparsers) diff --git a/src/dstack/_internal/cli/models/endpoints.py b/src/dstack/_internal/cli/models/configurations.py similarity index 88% rename from src/dstack/_internal/cli/models/endpoints.py rename to src/dstack/_internal/cli/models/configurations.py index 03762c41b..b0e08942b 100644 --- a/src/dstack/_internal/cli/models/endpoints.py +++ b/src/dstack/_internal/cli/models/configurations.py @@ -11,11 +11,10 @@ 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): +class PresetModelRepo(CoreModel): repo: Annotated[str, Field(description="The exact model repo or path to deploy")] name: Annotated[ Optional[str], Field(description="The client-facing model name. Defaults to `repo`") @@ -44,7 +43,7 @@ def validate_name(cls, value: Optional[str]) -> Optional[str]: return _validate_model(value, field="name") -class EndpointModelBase(CoreModel): +class PresetModelBase(CoreModel): base: Annotated[ str, Field(description="The base model for which the agent may select a compatible variant"), @@ -67,12 +66,12 @@ def validate_base(cls, value: str) -> str: return _validate_model(value, field="base") -EndpointModelSpec = Union[EndpointModelRepo, EndpointModelBase] +PresetModelSpec = Union[PresetModelRepo, PresetModelBase] MAX_PROMPT_LENGTH = 10_000 -class EndpointPromptFile(CoreModel): +class PresetPromptFile(CoreModel): path: Annotated[ str, Field(description="The path to a prompt file, relative to the configuration file"), @@ -85,7 +84,7 @@ def validate_path(cls, value: str) -> str: return value -class EndpointConfigurationConfig(ProfileParamsConfig): +class PresetConfigurationConfig(ProfileParamsConfig): @staticmethod def schema_extra(schema: dict[str, Any]): ProfileParamsConfig.schema_extra(schema) @@ -95,9 +94,9 @@ def schema_extra(schema: dict[str, Any]): ) -class EndpointConfiguration( +class PresetConfiguration( ProfileParams, - generate_dual_core_model(EndpointConfigurationConfig), + generate_dual_core_model(PresetConfigurationConfig), ): type: Annotated[Literal["preset"], Field(description="The configuration type")] = "preset" name: Annotated[ @@ -105,7 +104,7 @@ class EndpointConfiguration( Field(description="The service name. Required unless passed with `--name`"), ] = None model: Annotated[ - EndpointModelSpec, + PresetModelSpec, Field( description=( "The model to serve. Use a string or `repo` for an exact repo/path, " @@ -128,7 +127,7 @@ class EndpointConfiguration( Field(description="The exact model repo/path to serve. Shorthand for `model.repo`"), ] = None prompt: Annotated[ - Optional[Union[str, EndpointPromptFile]], + Optional[Union[str, PresetPromptFile]], Field( description=( "Additional instructions for the preset creation agent, inline or as a file `path`" @@ -143,7 +142,7 @@ class EndpointConfiguration( Field( description=( "The maximum number of benchmarked trials during preset creation" - f" before the best one is promoted. Defaults to `{DEFAULT_MAX_TRIALS}`" + " before the best one is promoted" ) ), ] = None @@ -170,10 +169,6 @@ 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 @@ -210,12 +205,12 @@ def validate_prompt(cls, value: Any) -> Any: return value -class EndpointPresetConstraints(CoreModel): - """The effective constraints for endpoint preset creation, saved as `constraints.json` +class PresetConstraints(CoreModel): + """The effective constraints for preset creation, saved as `constraints.json` in the agent workspace. Field semantics are documented in the agent system prompt.""" run_name_prefix: str - model: EndpointModelSpec + model: PresetModelSpec context_length: Optional[PositiveInt] = None max_trials: PositiveInt concurrency: PositiveInt @@ -225,5 +220,5 @@ class EndpointPresetConstraints(CoreModel): 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") + raise ValueError(f"Preset model {field} must be a non-empty string") return value diff --git a/src/dstack/_internal/cli/models/endpoint_agent.py b/src/dstack/_internal/cli/models/preset_agent.py similarity index 95% rename from src/dstack/_internal/cli/models/endpoint_agent.py rename to src/dstack/_internal/cli/models/preset_agent.py index 05dad6602..e75694124 100644 --- a/src/dstack/_internal/cli/models/endpoint_agent.py +++ b/src/dstack/_internal/cli/models/preset_agent.py @@ -3,7 +3,7 @@ from pydantic import PositiveInt, root_validator -from dstack._internal.cli.models.endpoint_presets import EndpointBenchmark +from dstack._internal.cli.models.presets import PresetBenchmark from dstack._internal.core.models.common import CoreModel _LATENCY_JSON_SCHEMA = { @@ -97,7 +97,7 @@ class AgentFinalReport(CoreModel): base: Optional[str] = None model: Optional[str] = None context_length: Optional[PositiveInt] = None - benchmark: Optional[EndpointBenchmark] = None + benchmark: Optional[PresetBenchmark] = None failure_summary: Optional[str] = None @root_validator @@ -120,7 +120,7 @@ def validate_report(cls, values: dict) -> dict: return values -class EndpointAgentInfo(CoreModel): +class PresetAgentInfo(CoreModel): """Base information about the agent runtime that ran a preset creation session, saved in the debug session directory.""" @@ -133,7 +133,7 @@ class ClaudeModelParams(CoreModel): effort: str -class ClaudeAgentInfo(EndpointAgentInfo): +class ClaudeAgentInfo(PresetAgentInfo): """Claude agent runtime information, saved as `agent.json`.""" model: ClaudeModelParams diff --git a/src/dstack/_internal/cli/models/endpoint_trials.py b/src/dstack/_internal/cli/models/preset_trials.py similarity index 77% rename from src/dstack/_internal/cli/models/endpoint_trials.py rename to src/dstack/_internal/cli/models/preset_trials.py index 1dd137553..b2f92668d 100644 --- a/src/dstack/_internal/cli/models/endpoint_trials.py +++ b/src/dstack/_internal/cli/models/preset_trials.py @@ -1,16 +1,16 @@ from typing import Optional -from dstack._internal.cli.models.endpoint_presets import EndpointBenchmark +from dstack._internal.cli.models.presets import PresetBenchmark from dstack._internal.core.models.common import CoreModel from dstack._internal.core.models.configurations import TaskConfiguration from dstack._internal.core.models.resources import ResourcesSpec -class EndpointPresetTrial(CoreModel): +class PresetTrial(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 + benchmark: Optional[PresetBenchmark] = None """Null only for a failed trial (the configuration never served).""" diff --git a/src/dstack/_internal/cli/models/endpoint_presets.py b/src/dstack/_internal/cli/models/presets.py similarity index 87% rename from src/dstack/_internal/cli/models/endpoint_presets.py rename to src/dstack/_internal/cli/models/presets.py index 75aa778e8..16866ecc3 100644 --- a/src/dstack/_internal/cli/models/endpoint_presets.py +++ b/src/dstack/_internal/cli/models/presets.py @@ -17,7 +17,7 @@ from dstack._internal.core.models.resources import CPUSpec, ResourcesSpec -class EndpointBenchmarkWorkload(CoreModel): +class PresetBenchmarkWorkload(CoreModel): api: Literal["chat_completions", "completions"] num_requests: PositiveInt input_tokens: PositiveInt @@ -25,38 +25,38 @@ class EndpointBenchmarkWorkload(CoreModel): concurrency: PositiveInt -class EndpointBenchmarkLatency(CoreModel): +class PresetBenchmarkLatency(CoreModel): mean: Annotated[float, Field(ge=0)] p50: Annotated[float, Field(ge=0)] p99: Annotated[float, Field(ge=0)] -class EndpointBenchmarkMetrics(CoreModel): +class PresetBenchmarkMetrics(CoreModel): successful_requests: Annotated[int, Field(ge=0)] failed_requests: Annotated[int, Field(ge=0)] duration_seconds: PositiveFloat total_input_tokens: Annotated[int, Field(ge=0)] total_output_tokens: Annotated[int, Field(ge=0)] - ttft_ms: EndpointBenchmarkLatency - tpot_ms: EndpointBenchmarkLatency + ttft_ms: PresetBenchmarkLatency + tpot_ms: PresetBenchmarkLatency -class EndpointBenchmarkTarget(CoreModel): +class PresetBenchmarkTarget(CoreModel): type: Literal["gateway", "server-proxy"] -class EndpointBenchmarkClient(CoreModel): +class PresetBenchmarkClient(CoreModel): type: Literal["local"] -class EndpointBenchmark(CoreModel): +class PresetBenchmark(CoreModel): tool: str tool_version: str command: str - workload: EndpointBenchmarkWorkload - metrics: EndpointBenchmarkMetrics - target: Optional[EndpointBenchmarkTarget] = None - client: Optional[EndpointBenchmarkClient] = None + workload: PresetBenchmarkWorkload + metrics: PresetBenchmarkMetrics + target: Optional[PresetBenchmarkTarget] = None + client: Optional[PresetBenchmarkClient] = None @validator("tool", "tool_version", "command") def validate_non_empty(cls, value: str) -> str: @@ -89,18 +89,18 @@ def validate_metrics(cls, values: dict) -> dict: return values -class EndpointPresetValidationReplica(CoreModel): +class PresetValidationReplica(CoreModel): resources: list[ResourcesSpec] """Exact resources for each running replica in this service replica group.""" -class EndpointPresetValidation(CoreModel): - replicas: list[EndpointPresetValidationReplica] +class PresetValidation(CoreModel): + replicas: list[PresetValidationReplica] """Ordered to match `ServiceConfiguration.replica_groups`.""" - benchmark: EndpointBenchmark + benchmark: PresetBenchmark -class EndpointPreset(CoreModel): +class Preset(CoreModel): base: str """Base model used for local preset lookup.""" id: str @@ -112,7 +112,7 @@ class EndpointPreset(CoreModel): """Token context length this preset was verified to support.""" created_at: datetime service: ServiceConfiguration - validations: list[EndpointPresetValidation] + validations: list[PresetValidation] @validator("base", "id", "model") def validate_non_empty(cls, value: str) -> str: @@ -151,8 +151,8 @@ def validate_preset(cls, values: dict) -> dict: return values -class EndpointPresetListOutput(CoreModel): - presets: list[EndpointPreset] +class PresetListOutput(CoreModel): + presets: list[Preset] def _validate_exact_resources(resources: ResourcesSpec) -> None: diff --git a/src/dstack/_internal/cli/services/endpoints/__init__.py b/src/dstack/_internal/cli/services/presets/__init__.py similarity index 100% rename from src/dstack/_internal/cli/services/endpoints/__init__.py rename to src/dstack/_internal/cli/services/presets/__init__.py diff --git a/src/dstack/_internal/cli/services/endpoints/agent.py b/src/dstack/_internal/cli/services/presets/agent.py similarity index 93% rename from src/dstack/_internal/cli/services/endpoints/agent.py rename to src/dstack/_internal/cli/services/presets/agent.py index 4369f9de2..1b8f888ca 100644 --- a/src/dstack/_internal/cli/services/endpoints/agent.py +++ b/src/dstack/_internal/cli/services/presets/agent.py @@ -18,11 +18,11 @@ import yaml from rich.text import Text -from dstack._internal.cli.models.endpoint_agent import ( +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.models.preset_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 from dstack._internal.core.errors import CLIError @@ -38,7 +38,7 @@ _FINAL_REPORT_FILENAME = "final_report.json" _SESSION_FILENAME = "session.json" _USER_PROMPT_FILENAME = "user_prompt.md" -_PROGRESS_ENV = "DSTACK_ENDPOINT_PROGRESS_LOG" +_PROGRESS_ENV = "DSTACK_PRESET_PROGRESS_LOG" _REDACTION = "[redacted]" _CLAUDE_TOOLS = "Bash,Read,Write,Edit,WebFetch,WebSearch,StructuredOutput" _CLAUDE_EFFORT_LEVELS = ("low", "medium", "high", "xhigh", "max") @@ -93,7 +93,7 @@ class SessionBusyError(CLIError): @dataclass(frozen=True) -class EndpointAgentWorkspace: +class PresetAgentWorkspace: path: Path dstack_home: Path @@ -135,7 +135,7 @@ def agent_stderr_path(self) -> Path: @dataclass -class EndpointAgentSession: +class PresetAgentSession: path: Path timestamp: str debug: bool @@ -238,14 +238,14 @@ class ClaudeAuth: @dataclass -class EndpointAgentProcessOutput: +class PresetAgentProcessOutput: report_data: Optional[dict[str, Any]] = None error: Optional[str] = None session_id: Optional[str] = None made_progress: bool = False -def create_agent_workspace(session: EndpointAgentSession) -> EndpointAgentWorkspace: +def create_agent_workspace(session: PresetAgentSession) -> PresetAgentWorkspace: real = session.path / "workspace" try: real.mkdir(mode=0o700) @@ -256,7 +256,7 @@ def create_agent_workspace(session: EndpointAgentSession) -> EndpointAgentWorksp else: alias = _create_workspace_alias(real) _validate_control_socket_path(alias) - workspace = EndpointAgentWorkspace(path=alias / "w", dstack_home=alias / "h") + workspace = PresetAgentWorkspace(path=alias / "w", dstack_home=alias / "h") _prepare_workspace(workspace) except OSError as e: raise CLIError(f"Could not create the agent workspace under {real}: {e}") from e @@ -264,7 +264,7 @@ def create_agent_workspace(session: EndpointAgentSession) -> EndpointAgentWorksp return workspace -def attach_agent_workspace(session: EndpointAgentSession) -> EndpointAgentWorkspace: +def attach_agent_workspace(session: PresetAgentSession) -> PresetAgentWorkspace: manifest = session.read_manifest() real_value, alias_value = manifest.get("workspace"), manifest.get("alias") if not real_value or not alias_value: @@ -277,10 +277,10 @@ def attach_agent_workspace(session: EndpointAgentSession) -> EndpointAgentWorksp ) if alias != real: _ensure_workspace_alias(alias, real) - return EndpointAgentWorkspace(path=alias / "w", dstack_home=alias / "h") + return PresetAgentWorkspace(path=alias / "w", dstack_home=alias / "h") -def remove_agent_workspace(session: EndpointAgentSession) -> None: +def remove_agent_workspace(session: PresetAgentSession) -> None: manifest = session.read_manifest() alias = manifest.get("alias") workspace = manifest.get("workspace") @@ -321,11 +321,11 @@ def get_presets_dir() -> Path: return get_dstack_dir() / "presets" -def create_endpoint_agent_session( - configuration: EndpointConfiguration, +def create_preset_agent_session( + configuration: PresetConfiguration, *, debug: bool = False, -) -> EndpointAgentSession: +) -> PresetAgentSession: if configuration.name is None: raise CLIError("The service name is required to save agent output") timestamp = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S-%fZ") @@ -347,11 +347,10 @@ def create_endpoint_agent_session( "status": "running", "pid": os.getpid(), "pid_started_at": _process_started_at(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, + "max_trials": configuration.max_trials, "created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"), "debug": debug, } @@ -372,7 +371,7 @@ def create_endpoint_agent_session( 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, preset_id=preset_id) + return PresetAgentSession(path=path, timestamp=timestamp, debug=debug, preset_id=preset_id) def _get_claude_version(auth: "ClaudeAuth") -> Optional[str]: @@ -406,9 +405,9 @@ def _get_claude_auth_status(auth: "ClaudeAuth") -> dict[str, Any]: return {"authMethod": "unknown"} -def load_resumable_agent_session(preset_id: str) -> EndpointAgentSession: +def load_resumable_agent_session(preset_id: str) -> PresetAgentSession: path = get_presets_dir() / preset_id - session = EndpointAgentSession(path=path, timestamp="", debug=False, preset_id=preset_id) + session = PresetAgentSession(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: {preset_id}") @@ -454,9 +453,9 @@ def session_process_alive(manifest: dict[str, Any]) -> bool: return _pid_alive(pid, manifest.get("pid_started_at")) -def load_attachable_agent_session(preset_id: str) -> EndpointAgentSession: +def load_attachable_agent_session(preset_id: str) -> PresetAgentSession: path = get_presets_dir() / preset_id - session = EndpointAgentSession(path=path, timestamp="", debug=False, preset_id=preset_id) + session = PresetAgentSession(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: {preset_id}") @@ -486,16 +485,16 @@ def load_attachable_agent_session(preset_id: str) -> EndpointAgentSession: return session -def load_agent_session(preset_id: str) -> EndpointAgentSession: +def load_agent_session(preset_id: str) -> PresetAgentSession: """Loads a session of any status for read-only inspection (its log).""" path = get_presets_dir() / preset_id - session = EndpointAgentSession(path=path, timestamp="", debug=False, preset_id=preset_id) + session = PresetAgentSession(path=path, timestamp="", debug=False, preset_id=preset_id) if not path.is_dir() or not session.read_manifest(): raise CLIError(f"Unknown preset: {preset_id}") return session -def print_session_log(session: EndpointAgentSession) -> None: +def print_session_log(session: PresetAgentSession) -> None: """Prints the session's redacted progress log verbatim (no markup).""" try: content = session.log_path.read_text(encoding="utf-8") @@ -508,7 +507,7 @@ def print_session_log(session: EndpointAgentSession) -> None: def mark_session_owner( - session: EndpointAgentSession, + session: PresetAgentSession, *, project: Optional[str] = None, keep_service: Optional[bool] = None, @@ -540,7 +539,7 @@ def session_report_exists(manifest: dict[str, Any]) -> bool: return (Path(workspace) / "w" / _FINAL_REPORT_FILENAME).is_file() -def try_claim_session(session: EndpointAgentSession) -> Optional[int]: +def try_claim_session(session: PresetAgentSession) -> Optional[int]: """Takes an exclusive, kernel-held lock for the duration of a session's finalization, so concurrent readers can't both finalize it. Returns an open file descriptor to release via `release_session_claim`, or None if another @@ -590,25 +589,22 @@ def _try_lock_fd(fd: int) -> bool: 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") + """The name this session holds.""" + value = manifest.get("name") return value if isinstance(value, str) and value else None -def iter_agent_sessions() -> Iterator[EndpointAgentSession]: +def iter_agent_sessions() -> Iterator[PresetAgentSession]: """Yields a handle for every session directory under the presets dir.""" root = get_presets_dir() if not root.is_dir(): return for path in sorted(root.iterdir()): if path.is_dir() and not path.name.startswith((".", "models--")): - yield EndpointAgentSession(path=path, timestamp="", debug=False, preset_id=path.name) + yield PresetAgentSession(path=path, timestamp="", debug=False, preset_id=path.name) -def find_session_name_claims(name: str) -> list[EndpointAgentSession]: +def find_session_name_claims(name: str) -> list[PresetAgentSession]: """Sessions of any status holding `name`, including failed ones.""" return [ session @@ -703,12 +699,12 @@ def get_claude_auth() -> ClaudeAuth: ) -def build_endpoint_agent_env( +def build_preset_agent_env( *, api: Client, - endpoint_env: dict[str, str], + preset_env: dict[str, str], auth: ClaudeAuth, - workspace: EndpointAgentWorkspace, + workspace: PresetAgentWorkspace, token: str, ) -> dict[str, str]: config_manager = ConfigManager(workspace.dstack_home / ".dstack") @@ -720,7 +716,7 @@ def build_endpoint_agent_env( ) config_manager.save() env = {name: value for name in _INHERITED_ENV_NAMES if (value := os.getenv(name))} - env.update(endpoint_env) + env.update(preset_env) if IS_WINDOWS: env.update( {name: value for name in _WINDOWS_INHERITED_ENV_NAMES if (value := os.getenv(name))} @@ -744,16 +740,16 @@ def build_endpoint_agent_env( return env -async def run_endpoint_agent( +async def run_preset_agent( *, prompt: str, env: dict[str, str], - workspace: EndpointAgentWorkspace, + workspace: PresetAgentWorkspace, auth: ClaudeAuth, redacted_values: Sequence[str], - agent_session: EndpointAgentSession, + agent_session: PresetAgentSession, initial_resume_session_id: Optional[str] = None, -) -> EndpointAgentProcessOutput: +) -> PresetAgentProcessOutput: async with _session_tailers( workspace=workspace, agent_session=agent_session, redacted_values=redacted_values ): @@ -794,7 +790,7 @@ async def run_endpoint_agent( action = "resuming" else: action = "retrying" - print_endpoint_progress( + print_preset_progress( f"Agent process exited without a report; {action} in {delay}s.", agent_session=agent_session, ) @@ -806,10 +802,10 @@ async def _run_claude_process( command: list[str], prompt: str, env: dict[str, str], - workspace: EndpointAgentWorkspace, + workspace: PresetAgentWorkspace, redacted_values: Sequence[str], - agent_session: EndpointAgentSession, -) -> tuple[EndpointAgentProcessOutput, int]: + agent_session: PresetAgentSession, +) -> tuple[PresetAgentProcessOutput, int]: proc: Optional[asyncio.subprocess.Process] = None try: # The agent's streams go to workspace files rather than pipes, so the @@ -862,7 +858,7 @@ def agent_alive() -> bool: return output, returncode -def print_endpoint_progress(message: str, *, agent_session: EndpointAgentSession) -> None: +def print_preset_progress(message: str, *, agent_session: PresetAgentSession) -> None: timestamp = datetime.now(timezone.utc).astimezone().strftime("%Y-%m-%d %H:%M:%S") message = message.rstrip("\r\n") agent_session.append_log(f"[{timestamp}] {message}") @@ -914,7 +910,7 @@ def _validate_control_socket_path(build_root: Path) -> None: raise CLIError(f"Temporary path is too long for an SSH control socket: {build_root}") -def _prepare_workspace(workspace: EndpointAgentWorkspace) -> None: +def _prepare_workspace(workspace: PresetAgentWorkspace) -> None: workspace.path.mkdir(mode=0o700, parents=True, exist_ok=False) workspace.dstack_home.mkdir(mode=0o700) workspace.temp_path.mkdir(mode=0o700) @@ -1086,7 +1082,7 @@ def _write_private_text(path: Path, content: str) -> None: def _write_debug_trace( - session: EndpointAgentSession, + session: PresetAgentSession, *, stream_name: str, text: str, @@ -1126,8 +1122,8 @@ def _redact_trace_value(value: Any, redacted_values: Sequence[str]) -> Any: @asynccontextmanager async def _session_tailers( *, - workspace: EndpointAgentWorkspace, - agent_session: EndpointAgentSession, + workspace: PresetAgentWorkspace, + agent_session: PresetAgentSession, redacted_values: Sequence[str], ) -> AsyncIterator[None]: """Mirrors the session's progress and record files while the body runs.""" @@ -1174,11 +1170,11 @@ async def _session_tailers( async def _collect_agent_output( *, - workspace: EndpointAgentWorkspace, - agent_session: EndpointAgentSession, + workspace: PresetAgentWorkspace, + agent_session: PresetAgentSession, redacted_values: Sequence[str], is_alive: Callable[[], bool], -) -> EndpointAgentProcessOutput: +) -> PresetAgentProcessOutput: """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") @@ -1216,14 +1212,14 @@ async def _collect_agent_output( return output -async def attach_endpoint_agent( +async def attach_preset_agent( *, - workspace: EndpointAgentWorkspace, + workspace: PresetAgentWorkspace, redacted_values: Sequence[str], - agent_session: EndpointAgentSession, -) -> EndpointAgentProcessOutput: + agent_session: PresetAgentSession, +) -> PresetAgentProcessOutput: """Follows a detached session's agent to completion, like - `run_endpoint_agent` without owning the process.""" + `run_preset_agent` without owning the process.""" async with _session_tailers( workspace=workspace, agent_session=agent_session, redacted_values=redacted_values ): @@ -1246,9 +1242,9 @@ async def _read_process_stream( stream_name: str, parse_result: bool, redacted_values: Sequence[str], - agent_session: EndpointAgentSession, -) -> EndpointAgentProcessOutput: - output = EndpointAgentProcessOutput() + agent_session: PresetAgentSession, +) -> PresetAgentProcessOutput: + output = PresetAgentProcessOutput() while True: line = await stream.readline() if not line: @@ -1464,7 +1460,7 @@ def __init__( *, path: Path, redacted_values: Sequence[str], - agent_session: EndpointAgentSession, + agent_session: PresetAgentSession, offset_store: Optional[_OffsetStore] = None, offset_key: str = "progress", ) -> None: @@ -1492,7 +1488,7 @@ def flush(self) -> None: for line in lines: message = _parse_progress(line) if message is not None: - print_endpoint_progress( + print_preset_progress( redact(message, self._redacted_values), agent_session=self._agent_session, ) diff --git a/src/dstack/_internal/cli/services/endpoints/apply.py b/src/dstack/_internal/cli/services/presets/apply.py similarity index 83% rename from src/dstack/_internal/cli/services/endpoints/apply.py rename to src/dstack/_internal/cli/services/presets/apply.py index 8aed115b7..6db11b79a 100644 --- a/src/dstack/_internal/cli/services/endpoints/apply.py +++ b/src/dstack/_internal/cli/services/presets/apply.py @@ -4,13 +4,13 @@ from rich.markup import escape -from dstack._internal.cli.models.endpoint_presets import EndpointPreset -from dstack._internal.cli.models.endpoints import EndpointConfiguration +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.models.presets import Preset from dstack._internal.cli.services.configurators.run import ServiceConfigurator -from dstack._internal.cli.services.endpoints.output import ( - format_endpoint_benchmark, +from dstack._internal.cli.services.presets.output import ( + format_preset_benchmark, ) -from dstack._internal.cli.services.endpoints.store import EndpointPresetStore +from dstack._internal.cli.services.presets.store import PresetStore from dstack._internal.core.errors import CLIError from dstack._internal.core.models.configurations import ServiceConfiguration from dstack._internal.core.models.profiles import ProfileParams @@ -21,20 +21,20 @@ @dataclass(frozen=True) class _PresetPlan: - preset: EndpointPreset + preset: Preset run_plan: RunPlan repo: Repo -def apply_endpoint_preset( +def apply_preset( *, api: Client, - configuration: EndpointConfiguration, + configuration: PresetConfiguration, configuration_path: str, preset_ids: Optional[Sequence[str]], profile_name: Optional[str], command_args: argparse.Namespace, - store: EndpointPresetStore, + store: PresetStore, ) -> None: candidates = _get_candidate_presets(store.list(), preset_ids=preset_ids) presets = _get_matching_presets(candidates, configuration=configuration) @@ -67,10 +67,10 @@ def apply_endpoint_preset( def _get_candidate_presets( - presets: list[EndpointPreset], + presets: list[Preset], *, preset_ids: Optional[Sequence[str]], -) -> list[EndpointPreset]: +) -> list[Preset]: if not preset_ids: return presets presets_by_ref = {preset.id: preset for preset in presets} @@ -92,10 +92,10 @@ def _get_candidate_presets( def _get_matching_presets( - presets: list[EndpointPreset], + presets: list[Preset], *, - configuration: EndpointConfiguration, -) -> list[EndpointPreset]: + configuration: PresetConfiguration, +) -> list[Preset]: model_name = configuration.model.api_model_name matches = [] for preset in presets: @@ -116,9 +116,9 @@ def _get_matching_presets( def _select_plan( *, - configuration: EndpointConfiguration, + configuration: PresetConfiguration, configuration_path: str, - presets: list[EndpointPreset], + presets: list[Preset], configurator: ServiceConfigurator, service_args: argparse.Namespace, ) -> _PresetPlan: @@ -140,8 +140,8 @@ def _select_plan( def _build_service( - configuration: EndpointConfiguration, - preset: EndpointPreset, + configuration: PresetConfiguration, + preset: Preset, ) -> ServiceConfiguration: service = preset.service.copy(deep=True) service.name = configuration.name @@ -161,13 +161,13 @@ def _has_available_offers(plan: RunPlan) -> bool: ) -def _format_requested_model(configuration: EndpointConfiguration) -> str: +def _format_requested_model(configuration: PresetConfiguration) -> str: model = escape(configuration.model.api_model_name) if configuration.model.allows_variant_selection: return f"{model} ([secondary]base[/])" return model -def _format_selected_preset(preset: EndpointPreset) -> str: - details = format_endpoint_benchmark(preset, verbose=True) +def _format_selected_preset(preset: Preset) -> str: + details = format_preset_benchmark(preset, verbose=True) return f"{escape(preset.id)} ([secondary]{details}[/])" diff --git a/src/dstack/_internal/cli/services/endpoints/create.py b/src/dstack/_internal/cli/services/presets/create.py similarity index 86% rename from src/dstack/_internal/cli/services/endpoints/create.py rename to src/dstack/_internal/cli/services/presets/create.py index 263b57c76..167e78578 100644 --- a/src/dstack/_internal/cli/services/endpoints/create.py +++ b/src/dstack/_internal/cli/services/presets/create.py @@ -14,23 +14,23 @@ from rich.table import Table from rich.text import Text -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, - EndpointPresetConstraints, +from dstack._internal.cli.models.configurations import ( + PresetConfiguration, + PresetConstraints, ) -from dstack._internal.cli.services.endpoints.agent import ( - EndpointAgentSession, - EndpointAgentWorkspace, +from dstack._internal.cli.models.preset_agent import AgentFinalReport +from dstack._internal.cli.models.presets import Preset +from dstack._internal.cli.services.presets.agent import ( + PresetAgentSession, + PresetAgentWorkspace, SessionBusyError, attach_agent_workspace, - attach_endpoint_agent, - build_endpoint_agent_env, + attach_preset_agent, + build_preset_agent_env, claimed_session_name, contains_redacted_value, create_agent_workspace, - create_endpoint_agent_session, + create_preset_agent_session, get_claude_auth, get_redacted_values, get_sensitive_inherited_env_values, @@ -38,23 +38,23 @@ load_agent_session, load_attachable_agent_session, mark_session_owner, - print_endpoint_progress, + print_preset_progress, print_session_log, redact, release_session_claim, remove_agent_workspace, - run_endpoint_agent, + run_preset_agent, session_process_alive, session_report_exists, terminate_agent_process, try_claim_session, ) -from dstack._internal.cli.services.endpoints.presets import endpoint_preset_to_data -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.services.presets.presets import preset_to_data +from dstack._internal.cli.services.presets.prompt import get_preset_agent_system_prompt +from dstack._internal.cli.services.presets.store import PresetStore +from dstack._internal.cli.services.presets.verify import ( + build_verified_preset, + load_preset_agent_report, ) from dstack._internal.cli.utils.common import NO_OFFERS_WARNING, confirm_ask, console, warn from dstack._internal.cli.utils.offers import print_offers_table @@ -69,8 +69,8 @@ @dataclass(frozen=True) -class EndpointPresetCreateResult: - preset: EndpointPreset +class PresetCreateResult: + preset: Preset path: Path final_run_id: uuid.UUID final_run_name: str @@ -85,16 +85,16 @@ def __init__(self, error: Optional[str]) -> None: self.error = error -def follow_endpoint_preset( +def follow_preset( *, api: Client, - store: EndpointPresetStore, + store: PresetStore, preset_id: str, keep_service: bool = False, wait_for_run_stop: bool = True, echo: bool = True, claim: bool = False, -) -> EndpointPresetCreateResult: +) -> PresetCreateResult: """Re-owns a detached session: follows its agent to completion, then verifies and saves the preset (the finalize role, which must run CLI-side for secret-scrubbing and server-verified preset building). @@ -111,9 +111,9 @@ def follow_endpoint_preset( configuration = _load_session_configuration(agent_session) try: result = asyncio.run( - _create_endpoint_preset( + _create_preset( api=api, - configuration=_resolve_endpoint_env_best_effort(configuration), + configuration=_resolve_preset_env_best_effort(configuration), source_configuration=configuration, store=store, keep_service=keep_service, @@ -146,7 +146,7 @@ def follow_endpoint_preset( release_session_claim(lock) -def _load_session_configuration(agent_session: EndpointAgentSession) -> EndpointConfiguration: +def _load_session_configuration(agent_session: PresetAgentSession) -> PresetConfiguration: configuration_path = agent_session.path / "preset.dstack.yml" if not configuration_path.is_file(): raise CLIError( @@ -156,21 +156,21 @@ def _load_session_configuration(agent_session: EndpointAgentSession) -> Endpoint # The session copy is canonical output, not user input: parse it without # the user-facing deprecation warnings. try: - return EndpointConfiguration.parse_obj( + return PresetConfiguration.parse_obj( yaml.safe_load(configuration_path.read_text(encoding="utf-8")) ) except (OSError, ValueError) as e: raise CLIError(f"Could not read the preset configuration: {e}") from e -def show_endpoint_session_logs( +def show_preset_session_logs( *, project: Optional[str], - store: EndpointPresetStore, + store: PresetStore, preset_id: str, follow: bool, keep_service: bool, -) -> Optional[EndpointPresetCreateResult]: +) -> Optional[PresetCreateResult]: """`logs`: dump a session's log (any status). With `follow`, a still-live session is re-owned, followed to completion, and its preset saved; a finished session just prints its log. Returns the saved preset, if any.""" @@ -189,7 +189,7 @@ def show_endpoint_session_logs( # read-only dump never needs a server or authentication. print_session_log(session) try: - return follow_endpoint_preset( + return follow_preset( api=Client.from_config(project_name=project), store=store, preset_id=preset_id, @@ -203,7 +203,7 @@ def show_endpoint_session_logs( return None -def _follow_session_log_readonly(session: EndpointAgentSession) -> None: +def _follow_session_log_readonly(session: PresetAgentSession) -> None: """Read-only follow: another CLI owns the finalize, so just stream the log it writes until the preset reaches a terminal state.""" try: @@ -225,7 +225,7 @@ def _follow_session_log_readonly(session: EndpointAgentSession) -> None: time.sleep(1) -def reconcile_detached_sessions(store: EndpointPresetStore) -> None: +def reconcile_detached_sessions(store: PresetStore) -> None: """Finalizes sessions whose agent completed while no CLI was attached (graceful detach, or an ungraceful CLI death). This is what makes the saved preset independent of a foreground process: any read command runs it, and @@ -253,18 +253,18 @@ def _is_reconcilable(manifest: dict[str, Any]) -> bool: ) -def _reconcile_session(session: EndpointAgentSession, store: EndpointPresetStore) -> None: +def _reconcile_session(session: PresetAgentSession, store: PresetStore) -> None: manifest = session.read_manifest() try: api = Client.from_config(project_name=str(manifest.get("project") or "")) except Exception: # noqa: BLE001 — offline/misconfigured must not break the read command return - # follow_endpoint_preset takes the finalize claim (so a concurrent `logs -f` + # follow_preset takes the finalize claim (so a concurrent `logs -f` # or reconcile can't double-finalize), records the terminal status itself, # and leaves the session intact on a transient error. Every outcome is silent # here — the result shows in the list that follows. with suppress(Exception): - follow_endpoint_preset( + follow_preset( api=api, store=store, preset_id=session.preset_id, @@ -275,7 +275,7 @@ def _reconcile_session(session: EndpointAgentSession, store: EndpointPresetStore ) -def stop_endpoint_session(api: Client, preset_id: str) -> None: +def stop_preset_session(api: Client, preset_id: str) -> None: session = load_attachable_agent_session(preset_id) manifest = session.read_manifest() # If the agent already finished, there is nothing to stop. Leave the session @@ -289,7 +289,7 @@ def stop_endpoint_session(api: Client, preset_id: str) -> None: _suspend_agent_session(session) -def _stop_active_session_runs(api: Client, session: EndpointAgentSession) -> None: +def _stop_active_session_runs(api: Client, session: PresetAgentSession) -> None: """Stops the session's non-terminal runs (with a spinner), like `dstack stop`. Keeping a trial instance warm for resume is the detach path, not this.""" try: @@ -318,9 +318,9 @@ def _stop_active_session_runs(api: Client, session: EndpointAgentSession) -> Non api.client.runs.stop(api.project, active, abort=False) -def _resolve_endpoint_env_best_effort( - configuration: EndpointConfiguration, -) -> EndpointConfiguration: +def _resolve_preset_env_best_effort( + configuration: PresetConfiguration, +) -> PresetConfiguration: """For attach: env values only feed redaction, and the agent already runs, so missing variables are tolerable.""" configuration = configuration.copy(deep=True) @@ -336,23 +336,23 @@ def _resolve_endpoint_env_best_effort( return configuration -def create_endpoint_preset( +def create_preset( *, api: Client, - configuration: EndpointConfiguration, - store: EndpointPresetStore, + configuration: PresetConfiguration, + store: PresetStore, keep_service: bool = False, build_name: Optional[str] = None, debug: bool = False, - resume_session: Optional[EndpointAgentSession] = None, + resume_session: Optional[PresetAgentSession] = None, user_prompt: Optional[str] = None, allowed_fleets: Optional[tuple[str, ...]] = None, -) -> EndpointPresetCreateResult: - agent_session = resume_session or create_endpoint_agent_session(configuration, debug=debug) +) -> PresetCreateResult: + agent_session = resume_session or create_preset_agent_session(configuration, debug=debug) try: - resolved_configuration = _resolve_endpoint_env(configuration) + resolved_configuration = _resolve_preset_env(configuration) result = asyncio.run( - _create_endpoint_preset( + _create_preset( api=api, configuration=resolved_configuration, source_configuration=configuration, @@ -377,21 +377,21 @@ def create_endpoint_preset( return result -async def _create_endpoint_preset( +async def _create_preset( *, api: Client, - configuration: EndpointConfiguration, - store: EndpointPresetStore, - source_configuration: Optional[EndpointConfiguration] = None, + configuration: PresetConfiguration, + store: PresetStore, + source_configuration: Optional[PresetConfiguration] = None, keep_service: bool = False, build_name: Optional[str] = None, - agent_session: EndpointAgentSession, + agent_session: PresetAgentSession, resume: bool = False, attach: bool = False, wait_for_run_stop: bool = True, user_prompt: Optional[str] = None, allowed_fleets: Optional[tuple[str, ...]] = None, -) -> EndpointPresetCreateResult: +) -> PresetCreateResult: source_configuration = source_configuration or configuration initial_resume_session_id: Optional[str] = None if attach: @@ -437,7 +437,7 @@ async def _create_endpoint_preset( claude_model=auth.model if auth is not None else None, ) - endpoint_env = configuration.env.as_dict() + preset_env = configuration.env.as_dict() token = getattr(api.client, "_token", None) if not isinstance(token, str) or not token: raise CLIError("The configured dstack client has no authentication token") @@ -445,26 +445,26 @@ async def _create_endpoint_preset( [ token, (auth.api_key if auth is not None else None) or "", - *endpoint_env.values(), + *preset_env.values(), *get_sensitive_inherited_env_values(), ] ) env: dict[str, str] = {} report: Optional[AgentFinalReport] = None - preset: Optional[EndpointPreset] = None + preset: Optional[Preset] = None preset_path: Optional[Path] = None creation_succeeded = False interrupted = False cleanup_error: Optional[str] = None if auth is not None: - env = build_endpoint_agent_env( + env = build_preset_agent_env( api=api, - endpoint_env=endpoint_env, + preset_env=preset_env, auth=auth, workspace=workspace, token=token, ) - prompt = get_endpoint_agent_system_prompt(user_prompt=user_prompt) + prompt = get_preset_agent_system_prompt(user_prompt=user_prompt) if not resume and not attach: if user_prompt: agent_session.write_user_prompt(user_prompt) @@ -481,7 +481,7 @@ async def _create_endpoint_preset( agent_session.write_agent_info(auth) try: if attach: - process_output = await attach_endpoint_agent( + process_output = await attach_preset_agent( workspace=workspace, redacted_values=redacted_values, agent_session=agent_session, @@ -490,7 +490,7 @@ async def _create_endpoint_preset( raise AgentExitedWithoutReport(process_output.error) else: assert auth is not None - process_output = await run_endpoint_agent( + process_output = await run_preset_agent( prompt=prompt, env=env, workspace=workspace, @@ -499,20 +499,20 @@ async def _create_endpoint_preset( agent_session=agent_session, initial_resume_session_id=initial_resume_session_id, ) - report = load_endpoint_agent_report( + report = load_preset_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( + preset = build_verified_preset( run=run, - endpoint_configuration=source_configuration, + preset_configuration=source_configuration, report=report, preset_id=agent_session.preset_id or None, name=claimed_session_name(agent_session.read_manifest()), ) - if contains_redacted_value(endpoint_preset_to_data(preset), redacted_values): + if contains_redacted_value(preset_to_data(preset), redacted_values): raise CLIError("Generated preset contains a secret value") preset_path = store.save(preset) creation_succeeded = True @@ -562,7 +562,7 @@ async def _create_endpoint_preset( assert report is not None assert report.run_id is not None assert report.run_name is not None - return EndpointPresetCreateResult( + return PresetCreateResult( preset=preset, path=preset_path, final_run_id=report.run_id, @@ -570,7 +570,7 @@ async def _create_endpoint_preset( ) -def _resolve_endpoint_env(configuration: EndpointConfiguration) -> EndpointConfiguration: +def _resolve_preset_env(configuration: PresetConfiguration) -> PresetConfiguration: configuration = configuration.copy(deep=True) for key, value in configuration.env.items(): if not isinstance(value, EnvSentinel): @@ -583,7 +583,7 @@ def _resolve_endpoint_env(configuration: EndpointConfiguration) -> EndpointConfi def _finish_agent_session( - session: EndpointAgentSession, + session: PresetAgentSession, status: str, ) -> None: try: @@ -593,7 +593,7 @@ def _finish_agent_session( warn(f"Could not finalize agent output. Files remain at {session.path}: {e}") -def _detach_agent_session(session: EndpointAgentSession) -> None: +def _detach_agent_session(session: PresetAgentSession) -> None: """Releases ownership but leaves the agent running — it stays visible and reconcilable in `dstack preset`. Silent: `logs -f` calls this on Ctrl+C, and a viewer that just stops watching shouldn't announce anything.""" @@ -601,7 +601,7 @@ def _detach_agent_session(session: EndpointAgentSession) -> None: def _stop_or_detach_agent_session( - session: EndpointAgentSession, api: Optional[Client] = None + session: PresetAgentSession, api: Optional[Client] = None ) -> None: """`create` interrupt: stop the session, or detach and leave the agent working — it stays visible as a running session in `dstack preset`.""" @@ -626,7 +626,7 @@ def _stop_or_detach_agent_session( _suspend_agent_session(session) -def _suspend_agent_session(session: EndpointAgentSession) -> None: +def _suspend_agent_session(session: PresetAgentSession) -> None: try: session.finish("interrupted") except OSError as e: @@ -654,7 +654,7 @@ def _model_slug(model_name: str) -> str: return slug -def _load_build_name(workspace: EndpointAgentWorkspace) -> str: +def _load_build_name(workspace: PresetAgentWorkspace) -> str: try: data = json.loads(workspace.constraints_path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError) as e: @@ -665,7 +665,7 @@ def _load_build_name(workspace: EndpointAgentWorkspace) -> str: return prefix -def plan_endpoint_preset(*, api: Client, configuration: EndpointConfiguration) -> tuple[str, ...]: +def plan_preset(*, api: Client, configuration: PresetConfiguration) -> tuple[str, ...]: """Resolves the allowed fleets and shows what the agent will have to work with — Project, User, the effective fleets, and their offers. Agent-free.""" allowed_fleets = _get_allowed_fleets(api, configuration) @@ -706,7 +706,7 @@ def _print_fleet_offers(api: Client, allowed_fleets: tuple[str, ...]) -> None: warn(f"Could not list offers for the allowed fleets: {e}") -def _get_allowed_fleets(api: Client, configuration: EndpointConfiguration) -> tuple[str, ...]: +def _get_allowed_fleets(api: Client, configuration: PresetConfiguration) -> tuple[str, ...]: if configuration.fleets is not None: return tuple( fleet.format() if hasattr(fleet, "format") else str(fleet) @@ -722,16 +722,16 @@ def _get_allowed_fleets(api: Client, configuration: EndpointConfiguration) -> tu def _build_constraints( *, - configuration: EndpointConfiguration, + configuration: PresetConfiguration, build_name: str, allowed_fleets: Sequence[str], ) -> str: - constraints = EndpointPresetConstraints.parse_obj( + constraints = PresetConstraints.parse_obj( { "run_name_prefix": build_name, "model": json.loads(configuration.model.json(exclude_none=True)), "context_length": configuration.context_length, - "max_trials": configuration.effective_max_trials, + "max_trials": configuration.max_trials, "concurrency": configuration.effective_concurrency, "fleets": list(allowed_fleets), "env": list(configuration.env), @@ -743,8 +743,8 @@ def _build_constraints( def _save_final_report_copy( *, - workspace: EndpointAgentWorkspace, - agent_session: EndpointAgentSession, + workspace: PresetAgentWorkspace, + agent_session: PresetAgentSession, redacted_values: Sequence[str], ) -> None: if not workspace.final_report_path.exists(): @@ -760,9 +760,9 @@ async def _cleanup_runs( *, api: Client, build_name: str, - workspace: EndpointAgentWorkspace, + workspace: PresetAgentWorkspace, final_run_name: Optional[str], - agent_session: EndpointAgentSession, + agent_session: PresetAgentSession, keep_final_service: bool = False, wait_for_stop: bool = True, ) -> None: @@ -798,7 +798,7 @@ async def _cleanup_runs( if pending: await asyncio.sleep(2) if agent_session.debug: - print_endpoint_progress("All preset creation runs stopped.", agent_session=agent_session) + print_preset_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/presets/output.py similarity index 94% rename from src/dstack/_internal/cli/services/endpoints/output.py rename to src/dstack/_internal/cli/services/presets/output.py index 460248917..34034b6df 100644 --- a/src/dstack/_internal/cli/services/endpoints/output.py +++ b/src/dstack/_internal/cli/services/presets/output.py @@ -4,8 +4,8 @@ from rich.table import Table -from dstack._internal.cli.models.endpoint_presets import ( - EndpointPreset, +from dstack._internal.cli.models.presets import ( + Preset, ) from dstack._internal.cli.utils.common import add_row_from_dict, console from dstack._internal.utils.common import pretty_date, pretty_resources @@ -50,8 +50,8 @@ def _format_trial_progress(session: Optional[dict[str, Any]]) -> str: return f" [secondary]({progress})[/]" -def print_endpoint_presets( - presets: list[EndpointPreset], +def print_presets( + presets: list[Preset], sessions: Optional[list[dict[str, Any]]] = None, verbose: bool = False, ) -> None: @@ -63,7 +63,7 @@ def print_endpoint_presets( table.add_column("STATUS", no_wrap=True) table.add_column("SUBMITTED", no_wrap=True, style="secondary") table.add_column("NAME", no_wrap=True, style="secondary") - presets_by_base: dict[str, list[EndpointPreset]] = defaultdict(list) + presets_by_base: dict[str, list[Preset]] = defaultdict(list) repo_to_base: dict[str, str] = {} for preset in presets: presets_by_base[preset.base].append(preset) @@ -137,7 +137,7 @@ def _add_session(table: Table, session: dict[str, Any], base_label: Optional[str def _add_preset( table: Table, - preset: EndpointPreset, + preset: Preset, *, verbose: bool, creation: Optional[dict[str, Any]] = None, @@ -149,7 +149,7 @@ def _add_preset( "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), + "BENCHMARK": format_preset_benchmark(preset, verbose=verbose), "SUBMITTED": pretty_date(preset.created_at), } if verbose and preset.model != preset.base: @@ -167,7 +167,7 @@ def _add_preset( ) -def format_endpoint_benchmark(preset: EndpointPreset, *, verbose: bool = False) -> str: +def format_preset_benchmark(preset: Preset, *, verbose: bool = False) -> str: benchmark = preset.validations[0].benchmark workload = benchmark.workload metrics = benchmark.metrics diff --git a/src/dstack/_internal/cli/services/endpoints/presets.py b/src/dstack/_internal/cli/services/presets/presets.py similarity index 91% rename from src/dstack/_internal/cli/services/endpoints/presets.py rename to src/dstack/_internal/cli/services/presets/presets.py index 553730346..3b4db39ee 100644 --- a/src/dstack/_internal/cli/services/endpoints/presets.py +++ b/src/dstack/_internal/cli/services/presets/presets.py @@ -4,11 +4,11 @@ import gpuhunt -from dstack._internal.cli.models.endpoint_presets import ( - EndpointBenchmark, - EndpointPreset, - EndpointPresetValidation, - EndpointPresetValidationReplica, +from dstack._internal.cli.models.presets import ( + Preset, + PresetBenchmark, + PresetValidation, + PresetValidationReplica, ) from dstack._internal.core.models.configurations import ServiceConfiguration from dstack._internal.core.models.envs import EnvSentinel @@ -18,31 +18,31 @@ from dstack._internal.utils.common import format_mib_as_gb, get_current_datetime -def build_endpoint_preset( +def build_preset( *, service: ServiceConfiguration, - validation_replicas: list[EndpointPresetValidationReplica], + validation_replicas: list[PresetValidationReplica], base_model: str, model: str, context_length: int, - benchmark: EndpointBenchmark, + benchmark: PresetBenchmark, preset_id: Optional[str] = None, name: Optional[str] = None, -) -> EndpointPreset: +) -> Preset: service = service.copy(deep=True) service.name = None service.gateway = None for field in ProfileParams.__fields__: setattr(service, field, None) - validation = EndpointPresetValidation( + validation = PresetValidation( replicas=validation_replicas, benchmark=benchmark, ) set_service_gpu_vendors_from_validations(service, [validation]) - return EndpointPreset( + return Preset( name=name, base=base_model, - id=preset_id or make_endpoint_preset_id(service, context_length=context_length), + id=preset_id or make_preset_id(service, context_length=context_length), model=model, context_length=context_length, created_at=get_current_datetime(), @@ -51,7 +51,7 @@ def build_endpoint_preset( ) -def make_endpoint_preset_id( +def make_preset_id( service: ServiceConfiguration, context_length: int, ) -> str: @@ -66,7 +66,7 @@ def make_endpoint_preset_id( return hashlib.sha256(payload.encode()).hexdigest()[:8] -def endpoint_preset_to_data(preset: EndpointPreset) -> dict[str, Any]: +def preset_to_data(preset: Preset) -> dict[str, Any]: return { "base": preset.base, "id": preset.id, @@ -134,7 +134,7 @@ def resources_spec_from_instance_resources(resources: Resources) -> ResourcesSpe def set_service_gpu_vendors_from_validations( service: ServiceConfiguration, - validations: list[EndpointPresetValidation], + validations: list[PresetValidation], ) -> None: for group_num, group in enumerate(service.replica_groups): resources = group.resources @@ -157,7 +157,7 @@ def _env_item_to_preset_data(key: str, value: str | EnvSentinel) -> str: def _get_validation_group_gpu_vendor( - validations: list[EndpointPresetValidation], + validations: list[PresetValidation], group_num: int, ) -> gpuhunt.AcceleratorVendor | None: vendors = { diff --git a/src/dstack/_internal/cli/services/endpoints/prompt.py b/src/dstack/_internal/cli/services/presets/prompt.py similarity index 94% rename from src/dstack/_internal/cli/services/endpoints/prompt.py rename to src/dstack/_internal/cli/services/presets/prompt.py index 0123047ee..0d0052c32 100644 --- a/src/dstack/_internal/cli/services/endpoints/prompt.py +++ b/src/dstack/_internal/cli/services/presets/prompt.py @@ -14,7 +14,7 @@ # 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: +def get_preset_agent_system_prompt(user_prompt: Optional[str] = None) -> str: text = _SYSTEM_PROMPT_PATH.read_text(encoding="utf-8").strip() variables = {"prompt": user_prompt.strip() if user_prompt else None} applied = 0 diff --git a/src/dstack/_internal/cli/services/endpoints/resources/system_prompt.md b/src/dstack/_internal/cli/services/presets/resources/system_prompt.md similarity index 100% rename from src/dstack/_internal/cli/services/endpoints/resources/system_prompt.md rename to src/dstack/_internal/cli/services/presets/resources/system_prompt.md diff --git a/src/dstack/_internal/cli/services/endpoints/store.py b/src/dstack/_internal/cli/services/presets/store.py similarity index 82% rename from src/dstack/_internal/cli/services/endpoints/store.py rename to src/dstack/_internal/cli/services/presets/store.py index 561d88a47..7c2580fe6 100644 --- a/src/dstack/_internal/cli/services/endpoints/store.py +++ b/src/dstack/_internal/cli/services/presets/store.py @@ -9,19 +9,19 @@ import yaml from pydantic import ValidationError -from dstack._internal.cli.models.endpoint_presets import EndpointPreset -from dstack._internal.cli.models.endpoints import ( +from dstack._internal.cli.models.configurations import ( MAX_PROMPT_LENGTH, - EndpointConfiguration, - EndpointPromptFile, + PresetConfiguration, + PresetPromptFile, ) -from dstack._internal.cli.services.endpoints.presets import endpoint_preset_to_data +from dstack._internal.cli.models.presets import Preset +from dstack._internal.cli.services.presets.presets import preset_to_data from dstack._internal.cli.utils.common import warn from dstack._internal.core.errors import CLIError, ConfigurationError from dstack._internal.utils.common import get_dstack_dir -class EndpointPresetStore: +class PresetStore: """Presets live at `//` — one directory per preset holding the artifact (`preset.yaml`) next to the creation session internals. Deleted presets are archived under `/.archive/`.""" @@ -29,14 +29,14 @@ class EndpointPresetStore: def __init__(self, root: Path | None = None) -> None: self.root = root or get_dstack_dir() / "presets" - def list(self) -> list[EndpointPreset]: + def list(self) -> list[Preset]: if not self.root.exists(): return [] 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: + def get(self, preset_id: str) -> Preset | None: _validate_preset_id(preset_id) if not self.root.exists(): return None @@ -49,13 +49,13 @@ def get(self, preset_id: str) -> EndpointPreset | None: raise CLIError(f"Preset file {path} does not match its path") return preset - def save(self, preset: EndpointPreset) -> Path: + def save(self, preset: Preset) -> Path: _validate_preset_id(preset.id) self._migrate_legacy() directory = self.root / preset.id directory.mkdir(parents=True, exist_ok=True) path = directory / "preset.yaml" - content = yaml.safe_dump(endpoint_preset_to_data(preset), sort_keys=False) + content = yaml.safe_dump(preset_to_data(preset), sort_keys=False) fd, temporary_path = tempfile.mkstemp( dir=directory, prefix=f".{preset.id}.", @@ -74,13 +74,13 @@ def save(self, preset: EndpointPreset) -> Path: pass return path - def find_by_name(self, name: str) -> EndpointPreset | None: + def find_by_name(self, name: str) -> Preset | None: for preset in self.list(): if preset.name == name: return preset return None - def release_name(self, name: str) -> EndpointPreset | None: + def release_name(self, name: str) -> Preset | None: """Releases `name` from the preset holding it, keeping the preset.""" preset = self.find_by_name(name) if preset is None: @@ -119,10 +119,10 @@ def _migrate_legacy(self) -> None: with suppress(OSError): directory.rmdir() - def _load(self, path: Path) -> EndpointPreset: + def _load(self, path: Path) -> Preset: try: with path.open(encoding="utf-8") as f: - return EndpointPreset.parse_obj(yaml.safe_load(f)) + return Preset.parse_obj(yaml.safe_load(f)) except (OSError, ValidationError, yaml.YAMLError) as e: raise CLIError(f"Invalid preset file {path}: {e}") from e @@ -132,26 +132,26 @@ def _validate_preset_id(preset_id: str) -> None: raise CLIError(f"Invalid preset ID: {preset_id!r}") -def load_endpoint_configuration(path: str) -> tuple[str, EndpointConfiguration]: +def load_preset_configuration(path: str) -> tuple[str, PresetConfiguration]: if path == "-": - return "-", _parse_endpoint_configuration(sys.stdin) + return "-", _parse_preset_configuration(sys.stdin) configuration_path = Path(path) if not configuration_path.is_file(): raise ConfigurationError(f"Configuration file {path} does not exist") try: with configuration_path.open(encoding="utf-8") as f: - configuration = _parse_endpoint_configuration(f) + configuration = _parse_preset_configuration(f) except OSError as e: raise ConfigurationError(f"Failed to load configuration from {path}") from e return str(configuration_path.resolve()), configuration -def _parse_endpoint_configuration(stream: TextIO) -> EndpointConfiguration: +def _parse_preset_configuration(stream: TextIO) -> PresetConfiguration: try: data = yaml.safe_load(stream) if not isinstance(data, dict): raise ConfigurationError("Preset configuration must be a YAML object") - configuration = EndpointConfiguration.parse_obj(data) + configuration = PresetConfiguration.parse_obj(data) except ValidationError as e: raise ConfigurationError(e) from e except yaml.YAMLError as e: @@ -166,15 +166,15 @@ def _parse_endpoint_configuration(stream: TextIO) -> EndpointConfiguration: return configuration -def resolve_endpoint_prompt( - configuration: EndpointConfiguration, configuration_path: str +def resolve_preset_prompt( + configuration: PresetConfiguration, configuration_path: str ) -> str | None: """The resolved user prompt text; file paths are relative to the configuration file.""" if configuration.prompt is None: return None if isinstance(configuration.prompt, str): return configuration.prompt.strip() - assert isinstance(configuration.prompt, EndpointPromptFile) + assert isinstance(configuration.prompt, PresetPromptFile) base = Path.cwd() if configuration_path == "-" else Path(configuration_path).parent path = base / configuration.prompt.path try: diff --git a/src/dstack/_internal/cli/services/endpoints/verify.py b/src/dstack/_internal/cli/services/presets/verify.py similarity index 75% rename from src/dstack/_internal/cli/services/endpoints/verify.py rename to src/dstack/_internal/cli/services/presets/verify.py index f1917b1c4..c69a216d2 100644 --- a/src/dstack/_internal/cli/services/endpoints/verify.py +++ b/src/dstack/_internal/cli/services/presets/verify.py @@ -5,21 +5,21 @@ from pydantic import ValidationError -from dstack._internal.cli.models.endpoint_agent import AgentFinalReport -from dstack._internal.cli.models.endpoint_presets import ( - EndpointBenchmarkClient, - EndpointBenchmarkTarget, - EndpointPreset, - EndpointPresetValidationReplica, +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.models.preset_agent import AgentFinalReport +from dstack._internal.cli.models.presets import ( + Preset, + PresetBenchmarkClient, + PresetBenchmarkTarget, + PresetValidationReplica, ) -from dstack._internal.cli.models.endpoints import EndpointConfiguration -from dstack._internal.cli.services.endpoints.agent import ( - EndpointAgentProcessOutput, - EndpointAgentWorkspace, +from dstack._internal.cli.services.presets.agent import ( + PresetAgentProcessOutput, + PresetAgentWorkspace, redact, ) -from dstack._internal.cli.services.endpoints.presets import ( - build_endpoint_preset, +from dstack._internal.cli.services.presets.presets import ( + build_preset, resources_spec_from_instance_resources, ) from dstack._internal.core.errors import CLIError @@ -38,10 +38,10 @@ def _redact_structure(value: Any, redacted_values: Sequence[str]) -> Any: return value -def load_endpoint_agent_report( +def load_preset_agent_report( *, - output: EndpointAgentProcessOutput, - workspace: EndpointAgentWorkspace, + output: PresetAgentProcessOutput, + workspace: PresetAgentWorkspace, redacted_values: Sequence[str], ) -> AgentFinalReport: report_data = output.report_data or _load_json_object(workspace.final_report_path) @@ -70,14 +70,14 @@ def load_endpoint_agent_report( return report -def build_verified_endpoint_preset( +def build_verified_preset( *, run: Run, - endpoint_configuration: EndpointConfiguration, + preset_configuration: PresetConfiguration, report: AgentFinalReport, preset_id: Optional[str] = None, name: Optional[str] = None, -) -> EndpointPreset: +) -> Preset: if run.id != report.run_id or run.run_spec.run_name != report.run_name: raise CLIError("Claude final report identifies a different service run") if run.status != RunStatus.RUNNING or run.service is None: @@ -85,20 +85,20 @@ def build_verified_endpoint_preset( service = run.run_spec.configuration if not isinstance(service, ServiceConfiguration) or service.model is None: raise CLIError("Claude final run is not a model service") - if service.model.name != endpoint_configuration.model.api_model_name: + if service.model.name != preset_configuration.model.api_model_name: raise CLIError("Claude final service model name does not match the requested model") assert report.base is not None assert report.model is not None assert report.context_length is not None assert report.benchmark is not None - if endpoint_configuration.model.allows_variant_selection: - if report.base != endpoint_configuration.model.api_model_name: + if preset_configuration.model.allows_variant_selection: + if report.base != preset_configuration.model.api_model_name: raise CLIError("Claude final report base does not match the requested model") - elif report.model != endpoint_configuration.model.exact_repo: + elif report.model != preset_configuration.model.exact_repo: raise CLIError("Claude changed an exact model request") if ( - endpoint_configuration.context_length is not None - and report.context_length < endpoint_configuration.context_length + preset_configuration.context_length is not None + and report.context_length < preset_configuration.context_length ): raise CLIError("Claude final service does not meet the requested context length") @@ -107,16 +107,16 @@ def build_verified_endpoint_preset( ) benchmark = report.benchmark.copy( update={ - "target": EndpointBenchmarkTarget(type=target_type), - "client": EndpointBenchmarkClient(type="local"), + "target": PresetBenchmarkTarget(type=target_type), + "client": PresetBenchmarkClient(type="local"), } ) portable_service = service.copy(deep=True) - # The CLI resolved endpoint env references before submission; presets retain the references. - for key, value in endpoint_configuration.env.items(): + # The CLI resolved preset env references before submission; presets retain the references. + for key, value in preset_configuration.env.items(): if isinstance(value, EnvSentinel) and key in portable_service.env: portable_service.env[key] = value - return build_endpoint_preset( + return build_preset( name=name, service=portable_service, validation_replicas=_get_validation_replicas(run, service), @@ -131,8 +131,8 @@ def build_verified_endpoint_preset( def _get_validation_replicas( run: Run, service: ServiceConfiguration, -) -> list[EndpointPresetValidationReplica]: - replicas: list[EndpointPresetValidationReplica] = [] +) -> list[PresetValidationReplica]: + replicas: list[PresetValidationReplica] = [] for group in service.replica_groups: resources = [] for job in sorted(run.jobs, key=lambda job: job.job_spec.replica_num): @@ -154,7 +154,7 @@ def _get_validation_replicas( ) if not resources: raise CLIError(f"Final service replica group {group.name!r} has no running replicas") - replicas.append(EndpointPresetValidationReplica(resources=resources)) + replicas.append(PresetValidationReplica(resources=resources)) return replicas diff --git a/src/tests/_internal/cli/commands/test_endpoint.py b/src/tests/_internal/cli/commands/test_preset.py similarity index 77% rename from src/tests/_internal/cli/commands/test_endpoint.py rename to src/tests/_internal/cli/commands/test_preset.py index 0fac79b71..000ed6ff6 100644 --- a/src/tests/_internal/cli/commands/test_endpoint.py +++ b/src/tests/_internal/cli/commands/test_preset.py @@ -7,28 +7,28 @@ from rich.console import Console from rich.theme import Theme -from dstack._internal.cli.services.endpoints import output as endpoint_presets_utils -from dstack._internal.cli.services.endpoints.store import EndpointPresetStore +from dstack._internal.cli.services.presets import output as presets_utils +from dstack._internal.cli.services.presets.store import PresetStore from tests._internal.cli.common import run_dstack_cli -from tests._internal.cli.endpoint_presets import get_endpoint_preset +from tests._internal.cli.preset_factories import get_preset pytestmark = pytest.mark.windows -class TestEndpointPresetLocalCommands: +class TestPresetLocalCommands: @pytest.fixture(autouse=True) def mock_ssh_client_info(self): with patch("dstack._internal.cli.main.get_ssh_client_info"): yield def test_formats_second_scale_ttft_without_scientific_notation(self): - preset = get_endpoint_preset() + preset = get_preset() ttft = preset.validations[0].benchmark.metrics.ttft_ms ttft.mean = 8148.3 ttft.p50 = 8151.4 ttft.p99 = 8334.2 - output = endpoint_presets_utils.format_endpoint_benchmark(preset, verbose=True) + output = presets_utils.format_preset_benchmark(preset, verbose=True) assert output.startswith("ctx=32K ") assert "TTFT 8.15s" in output @@ -37,7 +37,7 @@ def test_formats_second_scale_ttft_without_scientific_notation(self): def test_preserves_benchmark_concurrency_at_narrow_width(self, monkeypatch): output = StringIO() monkeypatch.setattr( - endpoint_presets_utils, + presets_utils, "console", Console( file=output, @@ -47,14 +47,14 @@ def test_preserves_benchmark_concurrency_at_narrow_width(self, monkeypatch): ), ) - endpoint_presets_utils.print_endpoint_presets([get_endpoint_preset()]) + presets_utils.print_presets([get_preset()]) assert "con=1" in "".join(output.getvalue().split()) def test_prints_submitted_column(self, monkeypatch): output = StringIO() monkeypatch.setattr( - endpoint_presets_utils, + presets_utils, "console", Console( file=output, @@ -63,27 +63,27 @@ def test_prints_submitted_column(self, monkeypatch): theme=Theme({"secondary": "grey58"}), ), ) - monkeypatch.setattr(endpoint_presets_utils, "pretty_date", lambda _: "2 months ago") + monkeypatch.setattr(presets_utils, "pretty_date", lambda _: "2 months ago") - endpoint_presets_utils.print_endpoint_presets([get_endpoint_preset()]) + presets_utils.print_presets([get_preset()]) 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 = tmp_path / "preset.dstack.yml" configuration_path.write_text( - "type: preset\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\n" + "type: preset\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\nmax_trials: 1\n" ) with ( patch("dstack.api.Client.from_config"), patch( - "dstack._internal.cli.commands.endpoint.plan_endpoint_preset", + "dstack._internal.cli.commands.preset.plan_preset", return_value=("fleet-a",), ), patch( - "dstack._internal.cli.commands.endpoint.create_endpoint_preset", + "dstack._internal.cli.commands.preset.create_preset", side_effect=KeyboardInterrupt, ), ): @@ -97,15 +97,38 @@ def test_handles_keyboard_interrupt(self, tmp_path, capsys): # interrupt handler, so the command layer stays silent — no generic line. assert "Operation interrupted by user" not in capsys.readouterr().out + def test_create_requires_max_trials(self, tmp_path, capsys): + configuration_path = tmp_path / "preset.dstack.yml" + configuration_path.write_text("type: preset\nname: qwen\nbase: Qwen/Qwen3.5-27B\n") + + with ( + patch("dstack.api.Client.from_config"), + patch( + "dstack._internal.cli.commands.preset.plan_preset", + return_value=("fleet-a",), + ), + patch("dstack._internal.cli.commands.preset.create_preset") as create, + ): + exit_code = run_dstack_cli( + ["preset", "create", "-y", "-f", str(configuration_path)], + home_dir=tmp_path, + repo_dir=tmp_path, + ) + + assert exit_code == 1 + captured = capsys.readouterr() + assert "max_trials is required" in captured.out + captured.err + create.assert_not_called() + def test_lists_and_deletes_preset_without_api_client(self, tmp_path, capsys): - preset = get_endpoint_preset() - EndpointPresetStore(tmp_path / ".dstack" / "presets").save(preset) + preset = get_preset() + PresetStore(tmp_path / ".dstack" / "presets").save(preset) list_output = StringIO() with ( patch("dstack.api.Client.from_config") as from_config, patch.object( - endpoint_presets_utils, + presets_utils, "console", Console( file=list_output, @@ -114,9 +137,7 @@ def test_lists_and_deletes_preset_without_api_client(self, tmp_path, capsys): theme=Theme({"secondary": "grey58"}), ), ), - patch.object( - endpoint_presets_utils, "pretty_date", return_value="2 months ago" - ) as pretty_date, + patch.object(presets_utils, "pretty_date", return_value="2 months ago") as pretty_date, ): assert run_dstack_cli(["preset", "list"], home_dir=tmp_path) == 0 from_config.assert_not_called() @@ -162,12 +183,12 @@ def test_lists_and_deletes_preset_without_api_client(self, tmp_path, capsys): ) from_config.assert_not_called() - assert EndpointPresetStore(tmp_path / ".dstack" / "presets").list() == [] + assert PresetStore(tmp_path / ".dstack" / "presets").list() == [] assert not (tmp_path / ".dstack" / "presets" / "models--Qwen--Qwen3.5-27B").exists() def test_gets_complete_preset_as_json_without_api_client(self, tmp_path, capsys): - preset = get_endpoint_preset() - EndpointPresetStore(tmp_path / ".dstack" / "presets").save(preset) + preset = get_preset() + PresetStore(tmp_path / ".dstack" / "presets").save(preset) with patch("dstack.api.Client.from_config") as from_config: assert ( @@ -193,8 +214,8 @@ def test_gets_complete_preset_as_json_without_api_client(self, tmp_path, capsys) ], ) def test_lists_complete_presets_as_json(self, tmp_path, capsys, args): - preset = get_endpoint_preset() - EndpointPresetStore(tmp_path / ".dstack" / "presets").save(preset) + preset = get_preset() + PresetStore(tmp_path / ".dstack" / "presets").save(preset) assert run_dstack_cli(args, home_dir=tmp_path) == 0 @@ -209,8 +230,8 @@ def test_lists_complete_presets_as_json(self, tmp_path, capsys, args): @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") + preset = get_preset() + store = PresetStore(tmp_path / ".dstack" / "presets") store.save(preset) store.save(preset.copy(update={"id": "01234567"})) @@ -227,8 +248,8 @@ def test_deletes_presets_by_model_without_api_client(self, tmp_path, flag_attrib assert store.list() == [] def test_delete_by_model_keeps_other_presets(self, tmp_path): - preset = get_endpoint_preset() - store = EndpointPresetStore(tmp_path / ".dstack" / "presets") + preset = get_preset() + store = PresetStore(tmp_path / ".dstack" / "presets") store.save(preset) other = preset.copy( update={"id": "01234567", "base": "meta/Llama-4", "model": "meta/Llama-4"} @@ -248,8 +269,8 @@ def test_delete_by_model_keeps_other_presets(self, tmp_path): @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") + preset = get_preset() + store = PresetStore(tmp_path / ".dstack" / "presets") store.save(preset) store.save( preset.copy(update={"id": "01234567", "base": "meta/Llama-4", "model": "meta/Llama-4"}) @@ -273,7 +294,7 @@ def test_merges_profile_configuration_and_cli_args(self, tmp_path): spot_policy: spot """ ) - configuration_path = tmp_path / "endpoint.dstack.yml" + configuration_path = tmp_path / "preset.dstack.yml" configuration_path.write_text( """type: preset name: file-name @@ -281,11 +302,12 @@ def test_merges_profile_configuration_and_cli_args(self, tmp_path): base: Qwen/Qwen3.5-27B regions: [file-region] max_price: 0.5 +max_trials: 1 env: - HF_TOKEN """ ) - preset = get_endpoint_preset() + preset = get_preset() result = SimpleNamespace( preset=preset, path=tmp_path / "preset.yaml", @@ -295,11 +317,11 @@ 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", + "dstack._internal.cli.commands.preset.plan_preset", return_value=("fleet-a",), ), patch( - "dstack._internal.cli.commands.endpoint.create_endpoint_preset", + "dstack._internal.cli.commands.preset.create_preset", return_value=result, ) as create, ): @@ -356,7 +378,7 @@ def test_apply_passes_selected_profile_and_preset_ids( with ( patch("dstack.api.Client.from_config"), - patch("dstack._internal.cli.commands.endpoint.apply_endpoint_preset") as apply, + patch("dstack._internal.cli.commands.preset.apply_preset") as apply, ): args = [ "preset", @@ -381,11 +403,13 @@ def test_apply_passes_selected_profile_and_preset_ids( 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") + preset = get_preset().copy(update={"name": "qwen"}) + store = PresetStore(tmp_path / ".dstack" / "presets") store.save(preset) configuration_path = tmp_path / "preset.dstack.yml" - configuration_path.write_text("type: preset\nname: qwen\nbase: Qwen/Qwen3.5-27B\n") + configuration_path.write_text( + "type: preset\nname: qwen\nbase: Qwen/Qwen3.5-27B\nmax_trials: 1\n" + ) result = SimpleNamespace( preset=preset, path=tmp_path / "preset.yaml", final_run_name="qwen-1" ) @@ -393,11 +417,11 @@ def test_create_detaches_the_name_from_the_old_preset(self, tmp_path, monkeypatc with ( patch("dstack.api.Client.from_config"), patch( - "dstack._internal.cli.commands.endpoint.plan_endpoint_preset", + "dstack._internal.cli.commands.preset.plan_preset", return_value=("fleet-a",), ), patch( - "dstack._internal.cli.commands.endpoint.create_endpoint_preset", + "dstack._internal.cli.commands.preset.create_preset", return_value=result, ) as create, ): @@ -412,20 +436,22 @@ def test_create_detaches_the_name_from_the_old_preset(self, tmp_path, monkeypatc 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") + preset = get_preset().copy(update={"name": "qwen"}) + store = PresetStore(tmp_path / ".dstack" / "presets") store.save(preset) configuration_path = tmp_path / "preset.dstack.yml" - configuration_path.write_text("type: preset\nname: qwen\nbase: Qwen/Qwen3.5-27B\n") + configuration_path.write_text( + "type: preset\nname: qwen\nbase: Qwen/Qwen3.5-27B\nmax_trials: 1\n" + ) with ( patch("dstack.api.Client.from_config"), patch( - "dstack._internal.cli.commands.endpoint.plan_endpoint_preset", + "dstack._internal.cli.commands.preset.plan_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, + patch("dstack._internal.cli.commands.preset.confirm_ask", return_value=False), + patch("dstack._internal.cli.commands.preset.create_preset") as create, ): exit_code = run_dstack_cli( ["preset", "create", "-f", str(configuration_path)], @@ -438,29 +464,29 @@ def test_create_without_confirmation_exits_before_creating(self, tmp_path, capsy 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) + preset = get_preset().copy(update={"name": "qwen"}) + PresetStore(tmp_path / ".dstack" / "presets").save(preset) assert run_dstack_cli(["preset", "get", "qwen", "--json"], home_dir=tmp_path) == 0 assert json.loads(capsys.readouterr().out)["id"] == preset.id assert run_dstack_cli(["preset", "delete", "qwen", "-y"], home_dir=tmp_path) == 0 - assert EndpointPresetStore(tmp_path / ".dstack" / "presets").list() == [] + assert PresetStore(tmp_path / ".dstack" / "presets").list() == [] def test_create_always_asks_even_without_a_name_conflict(self, tmp_path): configuration_path = tmp_path / "preset.dstack.yml" - configuration_path.write_text("type: preset\nbase: Qwen/Qwen3.5-27B\n") + configuration_path.write_text("type: preset\nbase: Qwen/Qwen3.5-27B\nmax_trials: 1\n") with ( patch("dstack.api.Client.from_config"), patch( - "dstack._internal.cli.commands.endpoint.plan_endpoint_preset", + "dstack._internal.cli.commands.preset.plan_preset", return_value=("fleet-a",), ), patch( - "dstack._internal.cli.commands.endpoint.confirm_ask", return_value=False + "dstack._internal.cli.commands.preset.confirm_ask", return_value=False ) as confirm, - patch("dstack._internal.cli.commands.endpoint.create_endpoint_preset") as create, + patch("dstack._internal.cli.commands.preset.create_preset") as create, ): exit_code = run_dstack_cli( ["preset", "create", "-f", str(configuration_path)], diff --git a/src/tests/_internal/cli/models/test_endpoints.py b/src/tests/_internal/cli/models/test_configurations.py similarity index 53% rename from src/tests/_internal/cli/models/test_endpoints.py rename to src/tests/_internal/cli/models/test_configurations.py index 5aa4b6cc1..81b17146d 100644 --- a/src/tests/_internal/cli/models/test_endpoints.py +++ b/src/tests/_internal/cli/models/test_configurations.py @@ -1,42 +1,42 @@ import pytest from pydantic import ValidationError -from dstack._internal.cli.models.endpoints import ( - EndpointConfiguration, - EndpointModelBase, - EndpointModelRepo, +from dstack._internal.cli.models.configurations import ( + PresetConfiguration, + PresetModelBase, + PresetModelRepo, ) pytestmark = pytest.mark.windows -class TestEndpointConfiguration: +class TestPresetConfiguration: def test_schema_documents_supported_input(self): assert all( - field.field_info.description for field in EndpointConfiguration.__fields__.values() + field.field_info.description for field in PresetConfiguration.__fields__.values() ) - assert all(field.field_info.description for field in EndpointModelBase.__fields__.values()) - assert all(field.field_info.description for field in EndpointModelRepo.__fields__.values()) - assert {"type": "string"} in EndpointConfiguration.schema()["properties"]["model"]["anyOf"] + assert all(field.field_info.description for field in PresetModelBase.__fields__.values()) + assert all(field.field_info.description for field in PresetModelRepo.__fields__.values()) + assert {"type": "string"} in PresetConfiguration.schema()["properties"]["model"]["anyOf"] def test_parses_string_as_exact_repo(self): - configuration = EndpointConfiguration(model="Qwen/Qwen3.5-27B") + configuration = PresetConfiguration(model="Qwen/Qwen3.5-27B") - assert isinstance(configuration.model, EndpointModelRepo) + assert isinstance(configuration.model, PresetModelRepo) assert configuration.model.exact_repo == "Qwen/Qwen3.5-27B" assert configuration.model.api_model_name == "Qwen/Qwen3.5-27B" assert not configuration.model.allows_variant_selection def test_parses_base_model(self): - configuration = EndpointConfiguration(model={"base": "Qwen/Qwen3.5-27B"}) + configuration = PresetConfiguration(model={"base": "Qwen/Qwen3.5-27B"}) - assert isinstance(configuration.model, EndpointModelBase) + assert isinstance(configuration.model, PresetModelBase) assert configuration.model.exact_repo is None assert configuration.model.api_model_name == "Qwen/Qwen3.5-27B" assert configuration.model.allows_variant_selection def test_parses_exact_repo_with_client_facing_name(self): - configuration = EndpointConfiguration( + configuration = PresetConfiguration( model={ "repo": "community/Qwen3.5-27B-GPTQ-Int4", "name": "Qwen/Qwen3.5-27B", @@ -48,37 +48,37 @@ 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"}) + PresetConfiguration(model={"base": "Qwen/base", "repo": "Qwen/repo"}) def test_parses_top_level_base_shorthand(self): - configuration = EndpointConfiguration(base="Qwen/Qwen3.5-27B") + configuration = PresetConfiguration(base="Qwen/Qwen3.5-27B") - assert isinstance(configuration.model, EndpointModelBase) + assert isinstance(configuration.model, PresetModelBase) assert configuration.model.api_model_name == "Qwen/Qwen3.5-27B" assert configuration.base is None def test_parses_top_level_repo_shorthand(self): - configuration = EndpointConfiguration(repo="community/Qwen3.5-27B-GPTQ-Int4") + configuration = PresetConfiguration(repo="community/Qwen3.5-27B-GPTQ-Int4") - assert isinstance(configuration.model, EndpointModelRepo) + assert isinstance(configuration.model, PresetModelRepo) assert configuration.model.exact_repo == "community/Qwen3.5-27B-GPTQ-Int4" assert configuration.repo is None def test_shorthand_round_trips_through_dict(self): - configuration = EndpointConfiguration(base="Qwen/Qwen3.5-27B") + configuration = PresetConfiguration(base="Qwen/Qwen3.5-27B") - round_tripped = EndpointConfiguration.parse_obj(configuration.dict()) + round_tripped = PresetConfiguration.parse_obj(configuration.dict()) assert round_tripped.model == configuration.model def test_rejects_combined_base_and_repo_shorthand(self): with pytest.raises(ValidationError): - EndpointConfiguration(base="Qwen/base", repo="Qwen/repo") + PresetConfiguration(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"}) + PresetConfiguration(base="Qwen/base", model={"repo": "Qwen/repo"}) def test_requires_model(self): with pytest.raises(ValidationError): - EndpointConfiguration() + PresetConfiguration() diff --git a/src/tests/_internal/cli/endpoint_presets.py b/src/tests/_internal/cli/preset_factories.py similarity index 81% rename from src/tests/_internal/cli/endpoint_presets.py rename to src/tests/_internal/cli/preset_factories.py index fe654857e..3104447d7 100644 --- a/src/tests/_internal/cli/endpoint_presets.py +++ b/src/tests/_internal/cli/preset_factories.py @@ -2,14 +2,14 @@ from types import SimpleNamespace from uuid import uuid4 -from dstack._internal.cli.models.endpoint_agent import AgentFinalReport -from dstack._internal.cli.models.endpoint_presets import ( - EndpointBenchmark, - EndpointBenchmarkClient, - EndpointBenchmarkTarget, - EndpointPreset, - EndpointPresetValidation, - EndpointPresetValidationReplica, +from dstack._internal.cli.models.preset_agent import AgentFinalReport +from dstack._internal.cli.models.presets import ( + Preset, + PresetBenchmark, + PresetBenchmarkClient, + PresetBenchmarkTarget, + PresetValidation, + PresetValidationReplica, ) from dstack._internal.core.models.configurations import ServiceConfiguration from dstack._internal.core.models.instances import Disk, Gpu, Resources @@ -17,8 +17,8 @@ from dstack._internal.core.models.runs import JobStatus, Run, RunStatus, ServiceSpec -def get_endpoint_benchmark(*, verified: bool = True) -> EndpointBenchmark: - benchmark = EndpointBenchmark( +def get_preset_benchmark(*, verified: bool = True) -> PresetBenchmark: + benchmark = PresetBenchmark( tool="vllm bench serve", tool_version="0.11.0", command="vllm bench serve --base-url $SERVICE_URL", @@ -43,17 +43,17 @@ def get_endpoint_benchmark(*, verified: bool = True) -> EndpointBenchmark: return benchmark return benchmark.copy( update={ - "target": EndpointBenchmarkTarget(type="server-proxy"), - "client": EndpointBenchmarkClient(type="local"), + "target": PresetBenchmarkTarget(type="server-proxy"), + "client": PresetBenchmarkClient(type="local"), } ) -def get_endpoint_preset( +def get_preset( *, preset_id: str = "8f3a12c4", context_length: int = 32768, -) -> EndpointPreset: +) -> Preset: resources = ResourcesSpec.parse_obj( { "cpu": "16", @@ -62,7 +62,7 @@ def get_endpoint_preset( "gpu": {"name": "A6000", "memory": "48GB", "count": 1}, } ) - return EndpointPreset( + return Preset( base="Qwen/Qwen3.5-27B", id=preset_id, model="community/Qwen3.5-27B-GPTQ-Int4", @@ -79,9 +79,9 @@ def get_endpoint_preset( } ), validations=[ - EndpointPresetValidation( - replicas=[EndpointPresetValidationReplica(resources=[resources])], - benchmark=get_endpoint_benchmark(), + PresetValidation( + replicas=[PresetValidationReplica(resources=[resources])], + benchmark=get_preset_benchmark(), ) ], ) @@ -136,7 +136,7 @@ def get_running_service_run() -> Run: ) -def get_successful_endpoint_report(run: Run) -> AgentFinalReport: +def get_successful_preset_report(run: Run) -> AgentFinalReport: return AgentFinalReport( success=True, run_id=run.id, @@ -145,5 +145,5 @@ def get_successful_endpoint_report(run: Run) -> AgentFinalReport: base="Qwen/Qwen3.5-27B", model="community/Qwen3.5-27B-GPTQ-Int4", context_length=32768, - benchmark=get_endpoint_benchmark(verified=False), + benchmark=get_preset_benchmark(verified=False), ) diff --git a/src/tests/_internal/cli/services/endpoints/__init__.py b/src/tests/_internal/cli/services/presets/__init__.py similarity index 100% rename from src/tests/_internal/cli/services/endpoints/__init__.py rename to src/tests/_internal/cli/services/presets/__init__.py diff --git a/src/tests/_internal/cli/services/endpoints/test_agent.py b/src/tests/_internal/cli/services/presets/test_agent.py similarity index 89% rename from src/tests/_internal/cli/services/endpoints/test_agent.py rename to src/tests/_internal/cli/services/presets/test_agent.py index a4e7c3b4e..9144025ba 100644 --- a/src/tests/_internal/cli/services/endpoints/test_agent.py +++ b/src/tests/_internal/cli/services/presets/test_agent.py @@ -13,11 +13,11 @@ import pytest import yaml -from dstack._internal.cli.models.endpoints import EndpointConfiguration -from dstack._internal.cli.services.endpoints.agent import ( +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.services.presets.agent import ( ClaudeAuth, - EndpointAgentSession, - EndpointAgentWorkspace, + PresetAgentSession, + PresetAgentWorkspace, _build_claude_command, _prepare_subprocess_command, _ProgressTailer, @@ -25,16 +25,16 @@ _summarize_session_trials, _terminate_process, attach_agent_workspace, - build_endpoint_agent_env, + build_preset_agent_env, contains_redacted_value, create_agent_workspace, - create_endpoint_agent_session, + create_preset_agent_session, get_claude_auth, load_resumable_agent_session, - print_endpoint_progress, + print_preset_progress, redact, remove_agent_workspace, - run_endpoint_agent, + run_preset_agent, ) from dstack._internal.compat import IS_WINDOWS from dstack._internal.core.errors import CLIError @@ -103,11 +103,11 @@ def test_inherits_only_required_environment(self, tmp_path, monkeypatch): client=SimpleNamespace(base_url="http://127.0.0.1:3000"), ) - env = build_endpoint_agent_env( + env = build_preset_agent_env( api=api, - endpoint_env={"HF_TOKEN": "hf-secret"}, + preset_env={"HF_TOKEN": "hf-secret"}, auth=_claude_auth(), - workspace=EndpointAgentWorkspace( + workspace=PresetAgentWorkspace( path=tmp_path, dstack_home=tmp_path / "home", ), @@ -163,7 +163,7 @@ 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( + session = PresetAgentSession( path=session_dir, timestamp="t", debug=False, preset_id="abcd1234" ) return create_agent_workspace(session) @@ -173,14 +173,14 @@ class TestAgentSession: def test_saves_log_and_debug_files(self, tmp_path, monkeypatch, capsys): monkeypatch.setenv("HOME", str(tmp_path)) monkeypatch.setenv("USERPROFILE", str(tmp_path)) - configuration = EndpointConfiguration( + configuration = PresetConfiguration( name="qwen", model={"base": "Qwen/Qwen3.5-27B"}, max_price=0.5, env=["HF_TOKEN", "TOKENIZERS_PARALLELISM=false"], ) - session = create_endpoint_agent_session(configuration) + session = create_preset_agent_session(configuration) assert session.path.parent == tmp_path / ".dstack" / "presets" assert session.path.name == session.preset_id @@ -193,17 +193,17 @@ def test_saves_log_and_debug_files(self, tmp_path, monkeypatch, capsys): 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["name"] == "qwen" assert manifest["model"] == "Qwen/Qwen3.5-27B" assert manifest["pid"] == os.getpid() - print_endpoint_progress("creating preset", agent_session=session) + print_preset_progress("creating preset", agent_session=session) assert "creating preset" in session.log_path.read_text() assert "creating preset" in capsys.readouterr().out if not IS_WINDOWS: assert session.path.stat().st_mode & 0o777 == 0o700 assert session.log_path.stat().st_mode & 0o777 == 0o600 - debug_session = create_endpoint_agent_session(configuration, debug=True) + debug_session = create_preset_agent_session(configuration, debug=True) data = yaml.safe_load((debug_session.path / "preset.dstack.yml").read_text()) assert {path.name for path in debug_session.path.iterdir()} == { @@ -219,7 +219,7 @@ def test_saves_log_and_debug_files(self, tmp_path, monkeypatch, capsys): 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_session = create_preset_agent_session(configuration) failed_path = failed_session.finish("failed") assert failed_path == failed_session.path assert json.loads((failed_session.path / "session.json").read_text())["status"] == "failed" @@ -228,7 +228,7 @@ 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( + session = PresetAgentSession( path=session_dir, timestamp="20260714-120000-000000Z", debug=False, @@ -246,15 +246,15 @@ def test_reports_invalid_existing_parent(self, tmp_path, monkeypatch): (tmp_path / ".dstack" / "presets").write_text("not a directory") with pytest.raises(CLIError, match="Could not create agent output"): - create_endpoint_agent_session( - EndpointConfiguration(name="qwen", model={"base": "Qwen/Qwen3.5-27B"}) + create_preset_agent_session( + PresetConfiguration(name="qwen", model={"base": "Qwen/Qwen3.5-27B"}) ) def test_log_write_failure_warns_once(self, tmp_path, capsys): path = tmp_path / "agent-running" path.mkdir() (path / "agent.log").touch() - session = EndpointAgentSession( + session = PresetAgentSession( path=path, timestamp="20260714-120000-000000Z", debug=False, @@ -294,22 +294,22 @@ async def test_sends_prompt_and_redacts_raw_output(self, tmp_path, monkeypatch, ) (tmp_path / "progress.jsonl").touch() monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.agent._build_claude_command", + "dstack._internal.cli.services.presets.agent._build_claude_command", lambda **_: [sys.executable, str(script)], ) - workspace = EndpointAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home") + workspace = PresetAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home") session_path = tmp_path / "debug-running" session_path.mkdir() (session_path / "agent.log").touch() (session_path / "trace.jsonl").touch() - agent_session = EndpointAgentSession( + agent_session = PresetAgentSession( path=session_path, timestamp="20260714-120000Z", debug=True, ) - output = await run_endpoint_agent( - prompt="full endpoint prompt", + output = await run_preset_agent( + prompt="full preset prompt", env=os.environ.copy(), workspace=workspace, auth=_claude_auth(), @@ -317,7 +317,7 @@ async def test_sends_prompt_and_redacts_raw_output(self, tmp_path, monkeypatch, agent_session=agent_session, ) - assert output.report_data == {"prompt": "full endpoint prompt"} + assert output.report_data == {"prompt": "full preset prompt"} assert output.error == "bad [redacted]" trace = [json.loads(line) for line in agent_session.trace_path.read_text().splitlines()] assert len(trace) == 1 @@ -344,17 +344,17 @@ async def test_accepts_stream_event_larger_than_64_kib(self, tmp_path, monkeypat session_path.mkdir() (session_path / "agent.log").touch() monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.agent._build_claude_command", + "dstack._internal.cli.services.presets.agent._build_claude_command", lambda **_: [sys.executable, str(script)], ) - output = await run_endpoint_agent( + output = await run_preset_agent( prompt="prompt", env=os.environ.copy(), - workspace=EndpointAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home"), + workspace=PresetAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home"), auth=_claude_auth(), redacted_values=(), - agent_session=EndpointAgentSession( + agent_session=PresetAgentSession( path=session_path, timestamp="20260714-120000Z", debug=False, @@ -370,7 +370,7 @@ def test_progress_stream_prints_only_redacted_messages(self, tmp_path, capsys): session_path = tmp_path / "agent-running" session_path.mkdir() (session_path / "agent.log").touch() - agent_session = EndpointAgentSession( + agent_session = PresetAgentSession( path=session_path, timestamp="20260714-120000Z", debug=False, @@ -451,16 +451,16 @@ def test_missing_source_is_no_op(self, tmp_path): class TestWriteAgentInfo: def test_writes_model_params_and_auth(self, tmp_path, monkeypatch): monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.agent._get_claude_version", + "dstack._internal.cli.services.presets.agent._get_claude_version", lambda auth: "2.1.0 (Claude Code)", ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.agent._get_claude_auth_status", + "dstack._internal.cli.services.presets.agent._get_claude_auth_status", lambda auth: {"authMethod": "claude.ai", "loggedIn": True}, ) session_dir = tmp_path / "session" session_dir.mkdir() - session = EndpointAgentSession(path=session_dir, timestamp="t", debug=True) + session = PresetAgentSession(path=session_dir, timestamp="t", debug=True) session.write_agent_info( ClaudeAuth(api_key=None, executable="claude", effort=None, model="claude-opus-4-8") @@ -505,11 +505,11 @@ def _write_flaky_claude(self, tmp_path): def _agent_setup(self, tmp_path): (tmp_path / "progress.jsonl").touch() - workspace = EndpointAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home") + workspace = PresetAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home") session_path = tmp_path / "session" session_path.mkdir() (session_path / "agent.log").touch() - agent_session = EndpointAgentSession( + agent_session = PresetAgentSession( path=session_path, timestamp="20260714-120000Z", debug=False, @@ -520,10 +520,10 @@ def _agent_setup(self, tmp_path): 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,) + "dstack._internal.cli.services.presets.agent._RESUME_DELAYS_SECONDS", (0,) ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.agent._build_claude_command", + "dstack._internal.cli.services.presets.agent._build_claude_command", lambda **kwargs: [sys.executable, str(script)] + ( ["--resume", kwargs["resume_session_id"]] @@ -533,7 +533,7 @@ async def test_resumes_after_connection_error(self, tmp_path, monkeypatch, capsy ) workspace, agent_session = self._agent_setup(tmp_path) - output = await run_endpoint_agent( + output = await run_preset_agent( prompt="system prompt", env={}, workspace=workspace, @@ -570,10 +570,10 @@ async def test_resumes_on_any_unreported_death(self, tmp_path, monkeypatch): """ ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.agent._RESUME_DELAYS_SECONDS", (0,) + "dstack._internal.cli.services.presets.agent._RESUME_DELAYS_SECONDS", (0,) ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.agent._build_claude_command", + "dstack._internal.cli.services.presets.agent._build_claude_command", lambda **kwargs: [sys.executable, str(script)] + ( ["--resume", kwargs["resume_session_id"]] @@ -583,7 +583,7 @@ async def test_resumes_on_any_unreported_death(self, tmp_path, monkeypatch): ) workspace, agent_session = self._agent_setup(tmp_path) - output = await run_endpoint_agent( + output = await run_preset_agent( prompt="system prompt", env={}, workspace=workspace, @@ -617,22 +617,22 @@ async def test_gives_up_after_repeated_no_progress_failures(self, tmp_path, monk """ ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.agent._RESUME_DELAYS_SECONDS", (0, 0) + "dstack._internal.cli.services.presets.agent._RESUME_DELAYS_SECONDS", (0, 0) ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.agent._build_claude_command", + "dstack._internal.cli.services.presets.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") + workspace = PresetAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home") session_path = tmp_path / "session" session_path.mkdir() (session_path / "agent.log").touch() - agent_session = EndpointAgentSession( + agent_session = PresetAgentSession( path=session_path, timestamp="20260714-120000Z", debug=False ) - output = await run_endpoint_agent( + output = await run_preset_agent( prompt="system prompt", env={}, workspace=workspace, @@ -651,7 +651,7 @@ class TestWorkspaceLifecycle: def _session(self, tmp_path): session_dir = tmp_path / "sessions" / "ab12cd34" session_dir.mkdir(parents=True) - return EndpointAgentSession( + return PresetAgentSession( path=session_dir, timestamp="t", debug=False, preset_id="ab12cd34" ) @@ -699,7 +699,7 @@ def test_attach_fails_when_workspace_is_gone(self, tmp_path): class TestOffsetPersistence: def test_mirror_does_not_duplicate_after_restart(self, tmp_path): - from dstack._internal.cli.services.endpoints.agent import _OffsetStore + from dstack._internal.cli.services.presets.agent import _OffsetStore source = tmp_path / "runs.jsonl" target = tmp_path / "mirror.jsonl" @@ -768,7 +768,7 @@ def test_treats_dead_running_session_as_resumable(self, tmp_path, monkeypatch): }, ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.agent.psutil.pid_exists", + "dstack._internal.cli.services.presets.agent.psutil.pid_exists", lambda pid: False, ) @@ -799,7 +799,7 @@ def test_refusals(self, tmp_path, monkeypatch): }, ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.agent.psutil.pid_exists", + "dstack._internal.cli.services.presets.agent.psutil.pid_exists", lambda pid: True, ) with pytest.raises(CLIError, match="still being created"): @@ -841,7 +841,7 @@ def test_counts_records_even_when_trials_share_a_task(self, tmp_path): class TestFileLineReader: @pytest.mark.asyncio async def test_reads_lines_and_continues_from_persisted_offset(self, tmp_path): - from dstack._internal.cli.services.endpoints.agent import _FileLineReader, _OffsetStore + from dstack._internal.cli.services.presets.agent import _FileLineReader, _OffsetStore stream = tmp_path / "stdout.jsonl" state = tmp_path / ".offsets.json" @@ -875,13 +875,13 @@ async def test_reads_lines_and_continues_from_persisted_offset(self, tmp_path): 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 + from dstack._internal.cli.services.presets import create as create_module + from dstack._internal.cli.services.presets.create import _stop_or_detach_agent_session session_dir = tmp_path / "ab12cd34" session_dir.mkdir() (session_dir / "agent.log").touch() - session = EndpointAgentSession( + session = PresetAgentSession( path=session_dir, timestamp="t", debug=False, preset_id="ab12cd34" ) agent = subprocess.Popen( @@ -907,7 +907,7 @@ def test_detach_keeps_the_agent_and_stop_terminates_it(self, tmp_path, monkeypat class TestOffsetStoreSharing: def test_disjoint_stores_do_not_clobber_each_other(self, tmp_path): - from dstack._internal.cli.services.endpoints.agent import _OffsetStore + from dstack._internal.cli.services.presets.agent import _OffsetStore state = tmp_path / ".offsets.json" readers = _OffsetStore(state) diff --git a/src/tests/_internal/cli/services/endpoints/test_apply.py b/src/tests/_internal/cli/services/presets/test_apply.py similarity index 77% rename from src/tests/_internal/cli/services/endpoints/test_apply.py rename to src/tests/_internal/cli/services/presets/test_apply.py index fab87e900..c413c70b8 100644 --- a/src/tests/_internal/cli/services/endpoints/test_apply.py +++ b/src/tests/_internal/cli/services/presets/test_apply.py @@ -3,17 +3,17 @@ import pytest -from dstack._internal.cli.models.endpoints import EndpointConfiguration -from dstack._internal.cli.services.endpoints.apply import ( +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.services.presets.apply import ( _build_service, _get_candidate_presets, _get_matching_presets, _select_plan, - apply_endpoint_preset, + apply_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 +from tests._internal.cli.preset_factories import get_preset pytestmark = pytest.mark.windows @@ -21,10 +21,10 @@ class TestGetMatchingPresets: def test_matches_base_model_context_and_preset(self): presets = [ - get_endpoint_preset(preset_id="small", context_length=4096), - get_endpoint_preset(preset_id="large", context_length=32768), + get_preset(preset_id="small", context_length=4096), + get_preset(preset_id="large", context_length=32768), ] - configuration = EndpointConfiguration( + configuration = PresetConfiguration( name="qwen", model={"base": "Qwen/Qwen3.5-27B"}, context_length=8192, @@ -40,8 +40,8 @@ def test_matches_base_model_context_and_preset(self): def test_candidates_preserve_id_order_and_reject_unknown_ids(self): presets = [ - get_endpoint_preset(preset_id="aa11bb22"), - get_endpoint_preset(preset_id="cc33dd44"), + get_preset(preset_id="aa11bb22"), + get_preset(preset_id="cc33dd44"), ] ordered = _get_candidate_presets(presets, preset_ids=["cc33dd44", "aa11bb22"]) @@ -51,8 +51,8 @@ def test_candidates_preserve_id_order_and_reject_unknown_ids(self): _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") - configuration = EndpointConfiguration( + matching = get_preset(preset_id="matching") + configuration = PresetConfiguration( name="qwen", model={ "repo": "community/Qwen3.5-27B-GPTQ-Int4", @@ -68,8 +68,8 @@ def test_exact_request_matches_repo_and_client_facing_name(self): class TestBuildService: - def test_applies_endpoint_name_env_gateway_and_constraints(self): - configuration = EndpointConfiguration( + def test_applies_preset_name_env_gateway_and_constraints(self): + configuration = PresetConfiguration( name="qwen-production", model={"base": "Qwen/Qwen3.5-27B"}, gateway="inference", @@ -78,7 +78,7 @@ def test_applies_endpoint_name_env_gateway_and_constraints(self): max_price=1, ) - service = _build_service(configuration, get_endpoint_preset()) + service = _build_service(configuration, get_preset()) assert service.name == "qwen-production" assert service.gateway == "inference" @@ -90,9 +90,9 @@ def test_applies_endpoint_name_env_gateway_and_constraints(self): class TestSelectPlan: def test_selects_first_preset_with_available_offer(self): presets = [ - get_endpoint_preset(preset_id="unavailable"), - get_endpoint_preset(preset_id="available"), - get_endpoint_preset(preset_id="never-probed"), + get_preset(preset_id="unavailable"), + get_preset(preset_id="available"), + get_preset(preset_id="never-probed"), ] configurator = Mock() plans = [ @@ -103,8 +103,8 @@ def test_selects_first_preset_with_available_offer(self): service_args = SimpleNamespace(profile=None) selected = _select_plan( - configuration=EndpointConfiguration(name="qwen", model={"base": "Qwen/Qwen3.5-27B"}), - configuration_path="endpoint.dstack.yml", + configuration=PresetConfiguration(name="qwen", model={"base": "Qwen/Qwen3.5-27B"}), + configuration_path="preset.dstack.yml", presets=presets, configurator=configurator, service_args=service_args, @@ -116,7 +116,7 @@ def test_selects_first_preset_with_available_offer(self): assert configurator.get_plan.call_count == 2 def test_applies_the_selected_plan(self, monkeypatch): - preset = get_endpoint_preset() + preset = get_preset() run_plan = _plan(InstanceAvailability.AVAILABLE) repo = Mock() service_args = SimpleNamespace(profile="gpu") @@ -124,14 +124,14 @@ def test_applies_the_selected_plan(self, monkeypatch): configurator.get_parser.return_value.parse_args.return_value = service_args configurator.get_plan.return_value = (run_plan, repo) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.apply.ServiceConfigurator", + "dstack._internal.cli.services.presets.apply.ServiceConfigurator", lambda api_client: configurator, ) command_args = SimpleNamespace() - apply_endpoint_preset( + apply_preset( api=Mock(), - configuration=EndpointConfiguration( + configuration=PresetConfiguration( name="qwen", model={"base": "Qwen/Qwen3.5-27B"}, ), diff --git a/src/tests/_internal/cli/services/endpoints/test_create.py b/src/tests/_internal/cli/services/presets/test_create.py similarity index 82% rename from src/tests/_internal/cli/services/endpoints/test_create.py rename to src/tests/_internal/cli/services/presets/test_create.py index d49b1b76b..5957ffbdf 100644 --- a/src/tests/_internal/cli/services/endpoints/test_create.py +++ b/src/tests/_internal/cli/services/presets/test_create.py @@ -6,44 +6,44 @@ import pytest -from dstack._internal.cli.models.endpoints import EndpointConfiguration -from dstack._internal.cli.services.endpoints.agent import ( +from dstack._internal.cli.models.configurations import PresetConfiguration +from dstack._internal.cli.services.presets.agent import ( ClaudeAuth, - EndpointAgentProcessOutput, - EndpointAgentSession, - EndpointAgentWorkspace, + PresetAgentProcessOutput, + PresetAgentSession, + PresetAgentWorkspace, create_agent_workspace, load_agent_session, mark_session_owner, - print_endpoint_progress, + print_preset_progress, print_session_log, release_session_claim, remove_agent_workspace, session_process_alive, try_claim_session, ) -from dstack._internal.cli.services.endpoints.create import ( - EndpointPresetCreateResult, +from dstack._internal.cli.services.presets.create import ( + PresetCreateResult, SessionBusyError, _build_constraints, _cleanup_runs, - _create_endpoint_preset, + _create_preset, _get_build_name, _print_fleet_offers, _save_final_report_copy, _stop_active_session_runs, - create_endpoint_preset, - follow_endpoint_preset, + create_preset, + follow_preset, reconcile_detached_sessions, ) -from dstack._internal.cli.services.endpoints.store import EndpointPresetStore +from dstack._internal.cli.services.presets.store import PresetStore from dstack._internal.core.errors import CLIError from dstack._internal.core.models.envs import EnvSentinel from dstack._internal.core.models.runs import Run, RunStatus -from tests._internal.cli.endpoint_presets import ( - get_endpoint_preset, +from tests._internal.cli.preset_factories import ( + get_preset, get_running_service_run, - get_successful_endpoint_report, + get_successful_preset_report, ) pytestmark = pytest.mark.windows @@ -71,26 +71,28 @@ def creation_context(tmp_path, monkeypatch): runs=run_apis, ), ) - configuration = EndpointConfiguration( + configuration = PresetConfiguration( name="qwen-build", model={"base": "Qwen/Qwen3.5-27B"}, context_length=8192, + max_trials=1, fleets=["gpu-fleet"], env={"LICENSE": "license-secret", "TOKENIZERS_PARALLELISM": "false"}, ) - source_configuration = EndpointConfiguration( + source_configuration = PresetConfiguration( name="qwen-build", model={"base": "Qwen/Qwen3.5-27B"}, context_length=8192, + max_trials=1, fleets=["gpu-fleet"], env=["LICENSE", "TOKENIZERS_PARALLELISM=false"], ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.get_claude_auth", + "dstack._internal.cli.services.presets.create.get_claude_auth", _claude_auth, ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create._get_build_name", + "dstack._internal.cli.services.presets.create._get_build_name", lambda *_: "qwen-build", ) return SimpleNamespace( @@ -99,19 +101,19 @@ def creation_context(tmp_path, monkeypatch): source_configuration=source_configuration, run=run, run_apis=run_apis, - store=EndpointPresetStore(tmp_path / "presets"), + store=PresetStore(tmp_path / "presets"), ) -class TestCreateEndpointPreset: +class TestCreatePreset: def test_saves_agent_log_without_debug(self, tmp_path, monkeypatch, capsys): monkeypatch.setenv("HOME", str(tmp_path)) monkeypatch.setenv("USERPROFILE", str(tmp_path)) - preset = get_endpoint_preset() + preset = get_preset() async def create(**kwargs): - print_endpoint_progress("testing preset", agent_session=kwargs["agent_session"]) - return EndpointPresetCreateResult( + print_preset_progress("testing preset", agent_session=kwargs["agent_session"]) + return PresetCreateResult( preset=preset, path=tmp_path / "preset.yaml", final_run_id=uuid.uuid4(), @@ -119,17 +121,17 @@ async def create(**kwargs): ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create._create_endpoint_preset", + "dstack._internal.cli.services.presets.create._create_preset", create, ) - create_endpoint_preset( + create_preset( api=SimpleNamespace(), - configuration=EndpointConfiguration( + configuration=PresetConfiguration( name="qwen", model={"base": "Qwen/Qwen3.5-27B"}, ), - store=EndpointPresetStore(tmp_path / "presets"), + store=PresetStore(tmp_path / "presets"), ) paths = [ @@ -152,7 +154,7 @@ def test_debug_finalization_error_does_not_mask_success(self, tmp_path, monkeypa monkeypatch.setenv("HOME", str(tmp_path)) monkeypatch.setenv("USERPROFILE", str(tmp_path)) monkeypatch.setenv("HF_TOKEN", "hf-secret") - preset = get_endpoint_preset() + preset = get_preset() async def create(**kwargs): assert kwargs["configuration"].env.as_dict() == { @@ -162,7 +164,7 @@ async def create(**kwargs): assert isinstance(kwargs["source_configuration"].env["HF_TOKEN"], EnvSentinel) assert kwargs["source_configuration"].env["TOKENIZERS_PARALLELISM"] == "false" kwargs["agent_session"].write_prompt("test prompt") - return EndpointPresetCreateResult( + return PresetCreateResult( preset=preset, path=tmp_path / "preset.yaml", final_run_id=uuid.uuid4(), @@ -173,19 +175,19 @@ def fail_finish(self, preset_id=None): raise OSError("rename failed") monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create._create_endpoint_preset", + "dstack._internal.cli.services.presets.create._create_preset", create, ) - monkeypatch.setattr(EndpointAgentSession, "finish", fail_finish) + monkeypatch.setattr(PresetAgentSession, "finish", fail_finish) - result = create_endpoint_preset( + result = create_preset( api=SimpleNamespace(), - configuration=EndpointConfiguration( + configuration=PresetConfiguration( name="qwen", model={"base": "Qwen/Qwen3.5-27B"}, env=["HF_TOKEN", "TOKENIZERS_PARALLELISM=false"], ), - store=EndpointPresetStore(tmp_path / "presets"), + store=PresetStore(tmp_path / "presets"), debug=True, ) @@ -218,19 +220,19 @@ def fail_finish(self, preset_id=None): raise OSError("rename failed") monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create._create_endpoint_preset", + "dstack._internal.cli.services.presets.create._create_preset", create, ) - monkeypatch.setattr(EndpointAgentSession, "finish", fail_finish) + monkeypatch.setattr(PresetAgentSession, "finish", fail_finish) with pytest.raises(RuntimeError, match="creation failed"): - create_endpoint_preset( + create_preset( api=SimpleNamespace(), - configuration=EndpointConfiguration( + configuration=PresetConfiguration( name="qwen", model={"base": "Qwen/Qwen3.5-27B"}, ), - store=EndpointPresetStore(tmp_path / "presets"), + store=PresetStore(tmp_path / "presets"), debug=True, ) @@ -241,18 +243,18 @@ async def test_checks_active_fleets_before_claude_auth(self, tmp_path, monkeypat client=SimpleNamespace(fleets=SimpleNamespace(list=lambda *args, **kwargs: [])), ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.get_claude_auth", + "dstack._internal.cli.services.presets.create.get_claude_auth", lambda: pytest.fail("Claude auth must not be checked without an active fleet"), ) with pytest.raises(CLIError, match="no fleets"): - await _create_endpoint_preset( + await _create_preset( api=api, - configuration=EndpointConfiguration( + configuration=PresetConfiguration( name="qwen-build", model={"base": "Qwen/Qwen3.5-27B"}, ), - store=EndpointPresetStore(tmp_path / "presets"), + store=PresetStore(tmp_path / "presets"), agent_session=_agent_session(tmp_path), ) @@ -267,16 +269,16 @@ async def cleanup_runs(**kwargs): cleanup_calls.append(kwargs) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.run_endpoint_agent", + "dstack._internal.cli.services.presets.create.run_preset_agent", run_agent, ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create._cleanup_runs", + "dstack._internal.cli.services.presets.create._cleanup_runs", cleanup_runs, ) with pytest.raises(asyncio.CancelledError): - await _create_endpoint_preset( + await _create_preset( api=creation_context.api, configuration=creation_context.configuration, source_configuration=creation_context.source_configuration, @@ -298,7 +300,7 @@ async def test_saves_preset_and_cleans_up_runs( session_path.mkdir() (session_path / "agent.log").touch() (session_path / "trace.jsonl").touch() - agent_session = EndpointAgentSession( + agent_session = PresetAgentSession( path=session_path, timestamp="20260714-120000Z", debug=True, @@ -307,15 +309,15 @@ async def test_saves_preset_and_cleans_up_runs( async def run_agent(**kwargs): assert kwargs["agent_session"] is agent_session assert (session_path / "prompt.md").is_file() - return EndpointAgentProcessOutput( - report_data=json.loads(get_successful_endpoint_report(creation_context.run).json()) + return PresetAgentProcessOutput( + report_data=json.loads(get_successful_preset_report(creation_context.run).json()) ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.run_endpoint_agent", + "dstack._internal.cli.services.presets.create.run_preset_agent", run_agent, ) - result = await _create_endpoint_preset( + result = await _create_preset( api=creation_context.api, configuration=creation_context.configuration, source_configuration=creation_context.source_configuration, @@ -338,7 +340,7 @@ def test_derives_slug_for_nameless_and_keeps_prefix_bounded(self): assert _get_build_name(None, "Qwen/Qwen3.5-27B", "a1b2c3d4") == "qwen3-5-27b-a1b2c3d4" build_name = _get_build_name( - "qwen-endpoint-with-a-name-that-is-forty-one", "Qwen/Qwen3.5-27B", "a1b2c3d4" + "qwen-preset-with-a-name-that-is-forty-one", "Qwen/Qwen3.5-27B", "a1b2c3d4" ) assert build_name.endswith("-a1b2c3d4") @@ -365,7 +367,7 @@ async def no_sleep(_): await _cleanup_runs( api=api, build_name="qwen-build", - workspace=EndpointAgentWorkspace( + workspace=PresetAgentWorkspace( path=tmp_path, dstack_home=tmp_path / "home", ), @@ -377,13 +379,13 @@ async def no_sleep(_): assert runs.stopped_names == ["qwen-build-1"] -def _agent_session(tmp_path, *, debug: bool = False) -> EndpointAgentSession: +def _agent_session(tmp_path, *, debug: bool = False) -> PresetAgentSession: path = tmp_path / "agent-running" path.mkdir() (path / "agent.log").touch() if debug: (path / "trace.jsonl").touch() - return EndpointAgentSession( + return PresetAgentSession( path=path, timestamp="20260714-120000Z", debug=debug, @@ -423,9 +425,10 @@ def stop(self, project, names, abort): class TestBuildConstraints: def test_renders_all_fields_with_explicit_nulls_and_defaults(self): - configuration = EndpointConfiguration( + configuration = PresetConfiguration( name="qwen", model={"base": "Qwen/Qwen3-32B"}, + max_trials=3, env=["HF_TOKEN"], ) @@ -447,7 +450,7 @@ def test_renders_all_fields_with_explicit_nulls_and_defaults(self): } def test_renders_configured_values(self): - configuration = EndpointConfiguration( + configuration = PresetConfiguration( name="qwen", model={"repo": "Qwen/Qwen3-32B-AWQ", "name": "qwen3"}, context_length=32768, @@ -472,7 +475,7 @@ def test_renders_configured_values(self): class TestSaveFinalReportCopy: def test_copies_report_redacted(self, tmp_path): - workspace = EndpointAgentWorkspace(path=tmp_path / "w", dstack_home=tmp_path / "h") + workspace = PresetAgentWorkspace(path=tmp_path / "w", dstack_home=tmp_path / "h") workspace.path.mkdir() workspace.final_report_path.write_text( '{"success": true, "note": "token dstack-secret"}', encoding="utf-8" @@ -490,7 +493,7 @@ def test_copies_report_redacted(self, tmp_path): 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 = PresetAgentWorkspace(path=tmp_path / "w", dstack_home=tmp_path / "h") workspace.path.mkdir() session = _agent_session(tmp_path, debug=True) @@ -512,17 +515,15 @@ async def create(**kwargs): raise KeyboardInterrupt monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create._create_endpoint_preset", + "dstack._internal.cli.services.presets.create._create_preset", create, ) with pytest.raises(KeyboardInterrupt): - create_endpoint_preset( + create_preset( api=SimpleNamespace(), - configuration=EndpointConfiguration( - name="qwen", model={"base": "Qwen/Qwen3.5-27B"} - ), - store=EndpointPresetStore(tmp_path / "presets"), + configuration=PresetConfiguration(name="qwen", model={"base": "Qwen/Qwen3.5-27B"}), + store=PresetStore(tmp_path / "presets"), ) sessions = [ @@ -542,7 +543,7 @@ async def test_resume_uses_saved_claude_session(self, creation_context, monkeypa session_dir = tmp_path / "sessions" / "fe98dc76" session_dir.mkdir(parents=True) (session_dir / "agent.log").touch() - agent_session = EndpointAgentSession( + agent_session = PresetAgentSession( path=session_dir, timestamp="t", debug=False, preset_id="fe98dc76" ) workspace = create_agent_workspace(agent_session) @@ -554,16 +555,16 @@ async def test_resume_uses_saved_claude_session(self, creation_context, monkeypa async def run_agent(**kwargs): captured.update(kwargs) - return EndpointAgentProcessOutput( - report_data=json.loads(get_successful_endpoint_report(creation_context.run).json()) + return PresetAgentProcessOutput( + report_data=json.loads(get_successful_preset_report(creation_context.run).json()) ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.run_endpoint_agent", + "dstack._internal.cli.services.presets.create.run_preset_agent", run_agent, ) - result = await _create_endpoint_preset( + result = await _create_preset( api=creation_context.api, configuration=creation_context.configuration, source_configuration=creation_context.source_configuration, @@ -583,23 +584,23 @@ async def test_pins_user_prompt_on_create(self, creation_context, monkeypatch, t session_dir = tmp_path / "ab34ef12" session_dir.mkdir() (session_dir / "agent.log").touch() - agent_session = EndpointAgentSession( + agent_session = PresetAgentSession( 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()) + return PresetAgentProcessOutput( + report_data=json.loads(get_successful_preset_report(creation_context.run).json()) ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.run_endpoint_agent", + "dstack._internal.cli.services.presets.create.run_preset_agent", run_agent, ) - await _create_endpoint_preset( + await _create_preset( api=creation_context.api, configuration=creation_context.configuration, source_configuration=creation_context.source_configuration, @@ -620,7 +621,7 @@ async def test_resume_keeps_the_pinned_user_prompt( session_dir = tmp_path / "ab34ef12" session_dir.mkdir() (session_dir / "agent.log").touch() - agent_session = EndpointAgentSession( + agent_session = PresetAgentSession( path=session_dir, timestamp="t", debug=False, preset_id="ab34ef12" ) workspace = create_agent_workspace(agent_session) @@ -633,16 +634,16 @@ async def test_resume_keeps_the_pinned_user_prompt( async def run_agent(**kwargs): captured.update(kwargs) - return EndpointAgentProcessOutput( - report_data=json.loads(get_successful_endpoint_report(creation_context.run).json()) + return PresetAgentProcessOutput( + report_data=json.loads(get_successful_preset_report(creation_context.run).json()) ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.run_endpoint_agent", + "dstack._internal.cli.services.presets.create.run_preset_agent", run_agent, ) - await _create_endpoint_preset( + await _create_preset( api=creation_context.api, configuration=creation_context.configuration, source_configuration=creation_context.source_configuration, @@ -679,11 +680,11 @@ def test_no_offers_shows_the_shared_warning_without_failing(self, capsys): class TestSessionLog: - def _session(self, tmp_path, preset_id: str, status: str, log: str) -> EndpointAgentSession: + def _session(self, tmp_path, preset_id: str, status: str, log: str) -> PresetAgentSession: session_dir = tmp_path / preset_id session_dir.mkdir() (session_dir / "agent.log").write_text(log) - session = EndpointAgentSession( + session = PresetAgentSession( path=session_dir, timestamp="t", debug=False, preset_id=preset_id ) session.update_manifest(status=status) @@ -692,7 +693,7 @@ def _session(self, tmp_path, preset_id: str, status: str, log: str) -> EndpointA def test_load_agent_session_reads_any_status(self, tmp_path, monkeypatch): self._session(tmp_path, "dead0000", "failed", "[t] boom\n") monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.agent.get_presets_dir", + "dstack._internal.cli.services.presets.agent.get_presets_dir", lambda: tmp_path, ) # A failed session is off-limits to follow/resume, but its log is readable. @@ -716,13 +717,13 @@ def test_print_session_log_notes_empty_log(self, tmp_path, capsys): assert "No log output yet" in capsys.readouterr().out -class TestFollowEndpointPreset: - def _detached_session(self, tmp_path, configuration_yaml: str) -> EndpointAgentSession: +class TestFollowPreset: + def _detached_session(self, tmp_path, configuration_yaml: str) -> PresetAgentSession: session_dir = tmp_path / "ab12cd34" session_dir.mkdir() (session_dir / "agent.log").touch() (session_dir / "preset.dstack.yml").write_text(configuration_yaml) - session = EndpointAgentSession( + session = PresetAgentSession( path=session_dir, timestamp="t", debug=False, preset_id="ab12cd34" ) workspace = create_agent_workspace(session) @@ -735,25 +736,25 @@ def test_finalizes_a_detached_session(self, creation_context, monkeypatch, tmp_p 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", + "dstack._internal.cli.services.presets.agent.get_presets_dir", lambda: tmp_path, ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.load_attachable_agent_session", + "dstack._internal.cli.services.presets.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()) + return PresetAgentProcessOutput( + report_data=json.loads(get_successful_preset_report(creation_context.run).json()) ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.attach_endpoint_agent", + "dstack._internal.cli.services.presets.create.attach_preset_agent", fake_attach, ) - result = follow_endpoint_preset( + result = follow_preset( api=creation_context.api, store=creation_context.store, preset_id="ab12cd34", @@ -770,20 +771,20 @@ def test_agent_death_without_report_suspends_instead_of_failing( tmp_path, "type: preset\nname: qwen\nmodel:\n base: Qwen/Qwen3.5-27B\n" ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.load_attachable_agent_session", + "dstack._internal.cli.services.presets.create.load_attachable_agent_session", lambda preset_id: session, ) async def fake_attach(**kwargs): - return EndpointAgentProcessOutput(error="agent died") + return PresetAgentProcessOutput(error="agent died") monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.attach_endpoint_agent", + "dstack._internal.cli.services.presets.create.attach_preset_agent", fake_attach, ) with pytest.raises(CLIError, match="agent died"): - follow_endpoint_preset( + follow_preset( api=creation_context.api, store=creation_context.store, preset_id="ab12cd34", @@ -799,12 +800,12 @@ def test_backs_off_when_claim_is_held(self, creation_context, monkeypatch, tmp_p held = try_claim_session(session) assert held is not None monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.load_attachable_agent_session", + "dstack._internal.cli.services.presets.create.load_attachable_agent_session", lambda preset_id: session, ) # claim=True must refuse rather than double-finalize; the session is untouched. with pytest.raises(SessionBusyError): - follow_endpoint_preset( + follow_preset( api=creation_context.api, store=creation_context.store, preset_id="ab12cd34", @@ -815,13 +816,13 @@ def test_backs_off_when_claim_is_held(self, creation_context, monkeypatch, tmp_p class TestStopActiveSessionRuns: - def _session(self, tmp_path) -> EndpointAgentSession: + def _session(self, tmp_path) -> PresetAgentSession: session_dir = tmp_path / "ab12cd34" session_dir.mkdir() (session_dir / "runs.jsonl").write_text( '{"name":"qwen-build-1","id":"a"}\n{"name":"qwen-build-2","id":"b"}\n' ) - return EndpointAgentSession( + return PresetAgentSession( path=session_dir, timestamp="t", debug=False, preset_id="ab12cd34" ) @@ -895,15 +896,15 @@ def _session_dir( def _patch(self, monkeypatch, tmp_path, follow): # reconcile iterates via agent.iter_agent_sessions -> agent.get_presets_dir. monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.agent.get_presets_dir", + "dstack._internal.cli.services.presets.agent.get_presets_dir", lambda: tmp_path, ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.Client", + "dstack._internal.cli.services.presets.create.Client", SimpleNamespace(from_config=lambda project_name=None: SimpleNamespace()), ) monkeypatch.setattr( - "dstack._internal.cli.services.endpoints.create.follow_endpoint_preset", + "dstack._internal.cli.services.presets.create.follow_preset", follow, ) @@ -920,7 +921,7 @@ def test_finalizes_eligible_detached_session(self, tmp_path, monkeypatch): self._session_dir(tmp_path, keep_service=True) calls: list = [] self._patch(monkeypatch, tmp_path, self._recording_follow(calls)) - reconcile_detached_sessions(EndpointPresetStore(tmp_path / "store")) + reconcile_detached_sessions(PresetStore(tmp_path / "store")) assert len(calls) == 1 assert calls[0]["preset_id"] == "dead0001" # Honors persisted keep-service; non-interactive, non-blocking, silent, @@ -947,7 +948,7 @@ def test_skips_ineligible(self, tmp_path, monkeypatch, kwargs): self._session_dir(tmp_path, **kwargs) calls: list = [] self._patch(monkeypatch, tmp_path, self._recording_follow(calls)) - reconcile_detached_sessions(EndpointPresetStore(tmp_path / "store")) + reconcile_detached_sessions(PresetStore(tmp_path / "store")) assert calls == [] def test_never_raises_when_finalize_fails(self, tmp_path, monkeypatch): @@ -958,13 +959,13 @@ def boom(**kwargs): self._patch(monkeypatch, tmp_path, boom) # A read command must never fail because reconcile did. - reconcile_detached_sessions(EndpointPresetStore(tmp_path / "store")) + reconcile_detached_sessions(PresetStore(tmp_path / "store")) class TestSessionClaim: def _session(self, tmp_path): (tmp_path / "sess").mkdir() - return EndpointAgentSession( + return PresetAgentSession( path=tmp_path / "sess", timestamp="", debug=False, preset_id="sess" ) @@ -1001,9 +1002,7 @@ def test_dead_pids_are_not_alive(self): class TestMarkSessionOwner: def test_persists_finalize_context(self, tmp_path): (tmp_path / "s").mkdir() - session = EndpointAgentSession( - path=tmp_path / "s", timestamp="", debug=False, preset_id="s" - ) + session = PresetAgentSession(path=tmp_path / "s", timestamp="", debug=False, preset_id="s") session.update_manifest(status="running") mark_session_owner(session, project="main", keep_service=True) manifest = session.read_manifest() diff --git a/src/tests/_internal/cli/services/endpoints/test_output.py b/src/tests/_internal/cli/services/presets/test_output.py similarity index 91% rename from src/tests/_internal/cli/services/endpoints/test_output.py rename to src/tests/_internal/cli/services/presets/test_output.py index eccad45f4..484623426 100644 --- a/src/tests/_internal/cli/services/endpoints/test_output.py +++ b/src/tests/_internal/cli/services/presets/test_output.py @@ -6,9 +6,9 @@ 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 +from dstack._internal.cli.services.presets import output as output_module +from dstack._internal.cli.services.presets.output import _add_session, _format_number +from tests._internal.cli.preset_factories import get_preset pytestmark = pytest.mark.windows @@ -81,7 +81,7 @@ def test_sorts_presets_and_sessions_newest_first(self, monkeypatch): file=buffer, width=200, color_system=None, theme=Theme({"secondary": "grey58"}) ), ) - old = get_endpoint_preset() + old = get_preset() new = old.copy(update={"id": "11aa22bb", "created_at": old.created_at + timedelta(days=2)}) sessions = [ { @@ -98,7 +98,7 @@ def test_sorts_presets_and_sessions_newest_first(self, monkeypatch): }, ] - output_module.print_endpoint_presets([old, new], sessions=sessions) + output_module.print_presets([old, new], sessions=sessions) text = buffer.getvalue() assert text.index("11aa22bb") < text.index(old.id) @@ -115,7 +115,7 @@ def test_completed_creation_decorates_preset_row_without_extra_session_row(self, file=buffer, width=200, color_system=None, theme=Theme({"secondary": "grey58"}) ), ) - preset = get_endpoint_preset() + preset = get_preset() sessions = [ { "id": preset.id, @@ -126,7 +126,7 @@ def test_completed_creation_decorates_preset_row_without_extra_session_row(self, } ] - output_module.print_endpoint_presets([preset], sessions=sessions) + output_module.print_presets([preset], sessions=sessions) text = buffer.getvalue() assert "verified (3/4)" in text diff --git a/src/tests/_internal/cli/services/endpoints/test_prompt.py b/src/tests/_internal/cli/services/presets/test_prompt.py similarity index 70% rename from src/tests/_internal/cli/services/endpoints/test_prompt.py rename to src/tests/_internal/cli/services/presets/test_prompt.py index 01c4a30c3..111576340 100644 --- a/src/tests/_internal/cli/services/endpoints/test_prompt.py +++ b/src/tests/_internal/cli/services/presets/test_prompt.py @@ -1,7 +1,7 @@ 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.cli.services.presets import prompt as prompt_module +from dstack._internal.cli.services.presets.prompt import get_preset_agent_system_prompt from dstack._internal.core.errors import CLIError pytestmark = pytest.mark.windows @@ -9,17 +9,15 @@ class TestSystemPrompt: def test_stays_byte_identical_without_user_prompt(self): - text = get_endpoint_agent_system_prompt() + text = get_preset_agent_system_prompt() - assert ( - text == get_endpoint_agent_system_prompt(None) == get_endpoint_agent_system_prompt("") - ) + assert text == get_preset_agent_system_prompt(None) == get_preset_agent_system_prompt("") assert "## Additional instructions" not in text assert "