From 1c96e107f830c20e859d9f36d1d0a48cb3253215 Mon Sep 17 00:00:00 2001 From: Andrey Cheptsov Date: Mon, 20 Jul 2026 12:39:58 -0700 Subject: [PATCH 01/10] [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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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/10] 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