diff --git a/mkdocs/docs/concepts/presets.md b/mkdocs/docs/concepts/presets.md index a10f5a40c..74bc07b88 100644 --- a/mkdocs/docs/concepts/presets.md +++ b/mkdocs/docs/concepts/presets.md @@ -9,6 +9,8 @@ A preset configuration lets you use an agent to create a preset: a verified and The value of presets comes from combining two fundamental features: agent-driven model inference optimization and the `dstack` [service](services.md) primitive, which can deploy model inference to any cloud, Kubernetes, or on-prem cluster. +To get the best performance for the given model, hardware, and other constraints, the agent selects the serving framework, quantization, and serving parameters, and can patch the framework's source code, generate custom kernels, and patch drivers. + > The presets feature is experimental and may change. ??? info "Prerequisites" @@ -169,6 +171,23 @@ prompt: | Set `baseline: true` to make the first trial a baseline: the agent serves the model the way the chosen serving framework recommends, without tuning it for performance. Later trials are optimization attempts. +### Previous sessions + +Set `previous` to a list of preset IDs to give the agent the results of earlier creation sessions. It analyzes what they tried and how it worked, and aims to improve on them instead of rediscovering it. + +
+ +```yaml +previous: + - c83375b4 +``` + +
+ +Alternatively, pass `--previous` (repeatable) to `dstack preset create`. + +With `baseline: true`, the first trial reproduces the best comparable previous result to confirm it still holds before optimizing further. + !!! info "Reference" The `preset` configuration supports many more options. See the [`.dstack.yml` reference](../reference/dstack.yml/preset.md). @@ -251,13 +270,11 @@ $ dstack preset delete c83375b4 For command options and agent settings, see the [`dstack preset` CLI reference](../reference/cli/dstack/preset.md). -!!! info "Roadmap and feedback" - Here's what is coming soon: - - * Allow the agent to change the source code, compile binaries, etc. - * Support for PD disaggregation - * Allow passing multiple `--previous ` to `dstack preset create` to reuse the insights from previous sessions - * Allow passing ranges to `concurrency` +!!! info "Limitations" + * Currently, the agent doesn't upload compiled binaries anywhere; patches compile at runtime + * Doesn't support PD disaggregation (coming soon) + * Doesn't allow a custom dataset; always uses `random` + * Doesn't support ranges for `concurrency` Report bugs and request features on [GitHub](https://github.com/dstackai/dstack/issues), and ask questions on [Discord](https://discord.gg/u8SmfwPpMd). diff --git a/src/dstack/_internal/cli/commands/preset.py b/src/dstack/_internal/cli/commands/preset.py index 562b3f31b..5f9952cc8 100644 --- a/src/dstack/_internal/cli/commands/preset.py +++ b/src/dstack/_internal/cli/commands/preset.py @@ -22,6 +22,7 @@ plan_preset, reassign_preset_name, reconcile_detached_sessions, + resolve_previous_sessions, show_preset_session_logs, stop_preset_session, ) @@ -114,6 +115,13 @@ def _register(self) -> None: action="store_true", help="Save the agent prompt and raw trace", ) + create_parser.add_argument( + "--previous", + action="append", + metavar="ID", + help="Give the agent a previous session's results to analyze and improve on." + " Repeat for several", + ) create_parser.add_argument( "--resume", metavar="ID", @@ -286,6 +294,14 @@ def _create(self, args: argparse.Namespace) -> None: "[warning]--trials is ignored when resuming: " "the constraints are fixed at creation[/]" ) + if configuration.previous: + console.print( + "[warning]previous is ignored when resuming: " + "the previous sessions are fixed at creation[/]" + ) + previous = () + if resume_session is None and configuration.previous: + previous = resolve_previous_sessions(configuration.previous) api = Client.from_config(project_name=args.project) allowed_fleets = None if resume_session is None: @@ -310,6 +326,7 @@ def _create(self, args: argparse.Namespace) -> None: resume_session=resume_session, user_prompt=user_prompt, allowed_fleets=allowed_fleets, + previous=previous, ) except KeyboardInterrupt: return # the interrupt handler already reported detach / stop @@ -524,6 +541,8 @@ def _get_effective_configuration( _apply_name(configuration, args.name, required=require_name) if getattr(args, "trials", None) is not None: configuration.trials = args.trials + if getattr(args, "previous", None): + configuration.previous = list(args.previous) profile = load_profile_from_args(args=args, repo_dir=Path.cwd()) for field in ProfileParams.model_fields: if getattr(configuration, field) is None: diff --git a/src/dstack/_internal/cli/models/configurations.py b/src/dstack/_internal/cli/models/configurations.py index e1bb40a7c..79c860ef9 100644 --- a/src/dstack/_internal/cli/models/configurations.py +++ b/src/dstack/_internal/cli/models/configurations.py @@ -154,6 +154,15 @@ class PresetConfiguration( ) ), ] = None + previous: Annotated[ + Optional[list[str]], + Field( + description=( + "The IDs of previous presets whose creation results the agent" + " analyzes and improves on" + ) + ), + ] = None concurrency: Annotated[ Optional[PositiveInt], Field( diff --git a/src/dstack/_internal/cli/models/preset_agent.py b/src/dstack/_internal/cli/models/preset_agent.py index 25c740083..ba29d32c2 100644 --- a/src/dstack/_internal/cli/models/preset_agent.py +++ b/src/dstack/_internal/cli/models/preset_agent.py @@ -85,6 +85,7 @@ "run_id": {"type": "string"}, "run_name": {"type": "string"}, "service_yaml": {"type": "string"}, + "trial": {"type": "integer", "minimum": 1}, "base": {"type": "string"}, "model": {"type": "string"}, "context_length": {"type": "integer", "minimum": 1}, @@ -101,6 +102,7 @@ class AgentFinalReport(CoreModel): run_id: Optional[uuid.UUID] = None run_name: Optional[str] = None service_yaml: Optional[str] = None + trial: Optional[PositiveInt] = None base: Optional[str] = None model: Optional[str] = None context_length: Optional[PositiveInt] = None @@ -114,6 +116,7 @@ def validate_report(self) -> Self: "run_id", "run_name", "service_yaml", + "trial", "base", "model", "context_length", diff --git a/src/dstack/_internal/cli/models/presets.py b/src/dstack/_internal/cli/models/presets.py index 8f2c0fdf7..6d824cb40 100644 --- a/src/dstack/_internal/cli/models/presets.py +++ b/src/dstack/_internal/cli/models/presets.py @@ -67,15 +67,10 @@ class PresetBenchmark(CoreModel): @property def effective_output_tok_per_s(self) -> float: - """Performance as defined in the agent prompt's `## Performance`. Derived - rather than read, so a miscomputed field cannot become the displayed truth.""" return self.metrics.total_output_tokens / self.metrics.duration_seconds @property def effective_per_user_tok_per_s(self) -> float: - """Per-user output speed as the serving literature defines it: the steady - decode rate, `1/TPOT`, which excludes time to first token. Dividing the - aggregate by concurrency instead folds TTFT and the ramp into it.""" return 1000 / self.metrics.tpot_ms.p50 @field_validator("tool", "tool_version", "command") @@ -113,25 +108,22 @@ def validate_metrics(self) -> Self: class PresetValidationReplica(CoreModel): resources: list[ResourcesSpec] - """Exact resources for each running replica in this service replica group.""" class PresetValidation(CoreModel): replicas: list[PresetValidationReplica] - """Ordered to match `ServiceConfiguration.replica_groups`.""" benchmark: PresetBenchmark class Preset(CoreModel): base: str - """Base model used for local preset lookup.""" id: str name: Optional[str] = None - """Mutable human name; at most one preset or in-flight session holds it.""" model: str - """Exact repo/path loaded by the service command.""" context_length: PositiveInt - """Token context length this preset was verified to support.""" + trial: Optional[PositiveInt] = None + min_context_length: Optional[PositiveInt] = None + max_ttft: Optional[PositiveInt] = None created_at: datetime service: ServiceConfiguration validations: list[PresetValidation] diff --git a/src/dstack/_internal/cli/services/presets/agent.py b/src/dstack/_internal/cli/services/presets/agent.py index 6dd1d8bbb..b28470891 100644 --- a/src/dstack/_internal/cli/services/presets/agent.py +++ b/src/dstack/_internal/cli/services/presets/agent.py @@ -23,6 +23,7 @@ print_preset_progress, ) from dstack._internal.cli.services.presets.tail import ( + _DirectoryMirror, _FileLineReader, _OffsetStore, _ProgressTailer, @@ -171,6 +172,9 @@ def build_preset_agent_env( env[_PROGRESS_ENV] = str(workspace.progress_path) for name in ["TMPDIR", "TEMP", "TMP"]: env[name] = str(workspace.temp_path) + # Sandbox the agent's Claude config under the workspace home when we pass our + # own API key; under subscription auth keep the real HOME so it reuses the + # user's existing `claude` login. if auth.api_key is not None: env["ANTHROPIC_API_KEY"] = auth.api_key env["HOME"] = str(workspace.dstack_home) @@ -222,14 +226,12 @@ async def run_preset_agent( # 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. + # Only reset the retry budget when the last attempt made progress; a + # run that keeps stalling exhausts its retries instead of retrying a + # stuck agent forever. if output.made_progress: retry_delays = list(_RESUME_DELAYS_SECONDS) - # An externally recorded stop is a decision, not an outage: never - # resurrect an agent another CLI just terminated. + # Another process marked this session interrupted; don't restart it. if agent_session.read_manifest().get("status") == "interrupted": return output if not retry_delays: @@ -276,9 +278,9 @@ async def _run_claude_process( stdout=stdout_file, stderr=stderr_file, start_new_session=not IS_WINDOWS, - # Inherit only the redirected std handles, not the CLI's other fds. - # Without this the untrusted agent inherits our open descriptors, and - # on Windows the broad inheritance flakes CreateProcess (WinError 87). + # So the untrusted agent inherits only the redirected std handles, + # not our other descriptors; broad inheritance also flakes + # CreateProcess on Windows (WinError 87). close_fds=True, ) agent_session.update_manifest( @@ -358,6 +360,8 @@ def _build_claude_command( def _prepare_subprocess_command(command: list[str]) -> list[str]: + """On Windows a `.bat`/`.cmd` Claude launcher can't be exec'd directly; wrap + it in `cmd.exe /c`. Every other case is returned unchanged.""" if not IS_WINDOWS or Path(command[0]).suffix.lower() not in {".bat", ".cmd"}: return command comspec = os.getenv("COMSPEC") or shutil.which("cmd.exe") @@ -399,7 +403,6 @@ async def _session_tailers( redacted_values: Sequence[str], offset_store: _OffsetStore, ) -> AsyncIterator[None]: - """Mirrors the session's progress and record files while the body runs.""" progress_tailer = _ProgressTailer( path=workspace.progress_path, redacted_values=redacted_values, @@ -415,20 +418,16 @@ async def _session_tailers( offset_key="runs", echo=agent_session.echo, ), - _RecordMirror( - source=workspace.trials_path, - target=agent_session.trials_path, + _DirectoryMirror( + source=workspace.trials_dir, + target=agent_session.trials_dir, redacted_values=redacted_values, - offset_store=offset_store, - offset_key="trials", echo=agent_session.echo, ), - _RecordMirror( - source=workspace.verifications_path, - target=agent_session.verifications_path, + _DirectoryMirror( + source=workspace.service_dir, + target=agent_session.service_dir, redacted_values=redacted_values, - offset_store=offset_store, - offset_key="verifications", echo=agent_session.echo, ), ] @@ -456,8 +455,7 @@ async def _collect_agent_output( is_alive: Callable[[], bool], offset_store: _OffsetStore, ) -> PresetAgentProcessOutput: - """Parses the agent's stream files until it exits; safe alongside a live - process or over the remains of a finished one.""" + """Safe to run alongside a live agent or over the stream files a finished one left behind.""" stdout_output, _ = await asyncio.gather( _read_process_stream( stream=_FileLineReader( @@ -495,8 +493,7 @@ async def attach_preset_agent( redacted_values: Sequence[str], agent_session: PresetAgentSession, ) -> PresetAgentProcessOutput: - """Follows a detached session's agent to completion, like - `run_preset_agent` without owning the process.""" + """Like `run_preset_agent`, but tails a detached agent it does not own.""" offset_store = open_session_offsets(agent_session) async with _session_tailers( workspace=workspace, @@ -574,8 +571,7 @@ async def _read_process_stream( async def _terminate_process(proc: asyncio.subprocess.Process) -> None: - """SIGTERM, a grace period, then SIGKILL — the same ladder as - `terminate_agent_process`, driven through the owned process handle.""" + """Twin of `terminate_agent_process` for a process this CLI owns, driven through its handle.""" if IS_WINDOWS: await asyncio.to_thread(_terminate_windows_process_tree, proc.pid) await proc.wait() @@ -601,9 +597,7 @@ async def _terminate_process(proc: asyncio.subprocess.Process) -> None: def terminate_agent_process(manifest: dict[str, Any]) -> None: - """Terminates the session's agent process tree, if alive. The same - SIGTERM-grace-SIGKILL ladder as `_terminate_process`, driven by pid because - the caller (`preset stop`) never owned the process.""" + """Twin of `_terminate_process` driven by pid, because the caller (`preset stop`) never owned the process.""" agent_pid = manifest.get("agent_pid") if not isinstance(agent_pid, int) or not _pid_alive( agent_pid, manifest.get("agent_started_at") diff --git a/src/dstack/_internal/cli/services/presets/apply.py b/src/dstack/_internal/cli/services/presets/apply.py index f5a15d22f..738cb9183 100644 --- a/src/dstack/_internal/cli/services/presets/apply.py +++ b/src/dstack/_internal/cli/services/presets/apply.py @@ -11,6 +11,7 @@ format_preset_objective, ) from dstack._internal.cli.services.presets.store import PresetStore +from dstack._internal.cli.utils.common import warn from dstack._internal.core.errors import CLIError from dstack._internal.core.models.configurations import ServiceConfiguration from dstack._internal.core.models.profiles import ProfileParams @@ -54,15 +55,15 @@ def apply_preset( def _validate_preset_matches(preset: Preset, *, configuration: PresetConfiguration) -> None: - """The referenced preset must serve what the configuration asks for.""" model_name = configuration.model.api_model_name service_model = preset.service.model if service_model is None or service_model.name.lower() != model_name.lower(): raise CLIError(f"Preset {preset.id} does not serve {model_name}") if configuration.min_context_length is not None: if preset.context_length < configuration.min_context_length: - raise CLIError( - f"Preset {preset.id} does not support context length" + warn( + f"Preset {preset.id} is verified for context length" + f" {preset.context_length}, below the requested" f" {configuration.min_context_length}" ) if configuration.model.allows_variant_selection: @@ -95,7 +96,5 @@ def _format_requested_model(configuration: PresetConfiguration) -> str: def _format_selected_preset(preset: Preset) -> str: - # The formatter dims its own keys; wrapping it again would flatten that. - # One line, so the objective and the result are joined rather than columned. details = f"{format_preset_objective(preset)} {format_preset_benchmark(preset, verbose=True)}" return f"{escape(preset.id)} ({details})" diff --git a/src/dstack/_internal/cli/services/presets/create.py b/src/dstack/_internal/cli/services/presets/create.py index c5ac21316..2d4b66200 100644 --- a/src/dstack/_internal/cli/services/presets/create.py +++ b/src/dstack/_internal/cli/services/presets/create.py @@ -50,6 +50,7 @@ print_preset_progress, print_session_log, release_session_claim, + resolve_session_ref, session_process_alive, session_report_exists, try_claim_session, @@ -63,6 +64,7 @@ PresetAgentWorkspace, attach_agent_workspace, create_agent_workspace, + install_previous_records, remove_agent_workspace, scrub_workspace_token, ) @@ -92,8 +94,7 @@ class CreationStopped(Exception): class AgentExitedWithoutReport(Exception): - """A detached agent died without submitting a report; the session is - resumable rather than failed.""" + """A detached agent exited without a report; the session can be resumed.""" def __init__(self, error: Optional[str]) -> None: super().__init__(error or "The agent exited without a report") @@ -109,13 +110,11 @@ def follow_preset( wait_for_run_stop: bool = True, echo: bool = True, ) -> PresetCreateResult: - """Re-owns a detached session: follows its agent to completion, then - verifies and saves the preset (the finalize role, which must run CLI-side - for secret-scrubbing and server-verified preset building). + """Finalizes a detached session CLI-side (secret-scrubbing and + server-verified preset building must run here, not on the server). - Always takes the exclusive finalize lock so a concurrent `logs -f` and - reconcile can't both finalize the same session. `wait_for_run_stop=False` - and `echo=False` make it non-blocking and silent for reconcile.""" + Takes the exclusive finalize lock so a concurrent `logs -f` and reconcile + can't both finalize the same session.""" agent_session = load_attachable_agent_session(preset_id) agent_session.echo = echo lock = try_claim_session(agent_session) @@ -146,8 +145,8 @@ def follow_preset( _suspend_agent_session(agent_session) raise CLIError(str(e)) from e except CLIError: - # Definitive: a failure/invalid report, an unverifiable service, or a - # leaked secret — the preset genuinely cannot be built, so fail it. + # A CLIError is definitive (bad report, unverifiable service, leaked + # secret): the preset cannot be built, so fail it. _close_agent_session(agent_session, "failed") raise # A transient error (network / OS) propagates untouched: the completed @@ -165,8 +164,6 @@ def _load_session_configuration(agent_session: PresetAgentSession) -> PresetConf f"Preset {agent_session.preset_id} has no saved configuration and cannot be" f" followed; resume it with --resume {agent_session.preset_id} instead" ) - # The session copy is canonical output, not user input: parse it without - # the user-facing deprecation warnings. try: return PresetConfiguration.model_validate( yaml.safe_load(configuration_path.read_text(encoding="utf-8")) @@ -183,17 +180,11 @@ def show_preset_session_logs( follow: bool, keep_service: bool, ) -> Optional[PresetCreateResult]: - """`logs`: dump a session's log (any status). With `follow`, a still-live - session is re-owned, followed to completion, and its preset saved; a - finished session just prints its log. Returns the saved preset, if any.""" session = load_agent_session(preset_id) status = session.read_manifest().get("status") if not follow or status in ("success", "failed", "interrupted"): print_session_log(session) return None - # Following a live session: print the log so far, then stream new progress - # (a future --since could bound this). The client is built only here, so a - # read-only dump never needs a server or authentication. print_session_log(session) try: return follow_preset( @@ -210,8 +201,6 @@ def show_preset_session_logs( def _follow_session_log_readonly(session: PresetAgentSession) -> None: - """Read-only follow: another CLI owns the finalize, so just stream the log it - writes until the preset reaches a terminal state.""" try: offset = session.log_path.stat().st_size except OSError: @@ -241,14 +230,11 @@ def _follow_session_log_readonly(session: PresetAgentSession) -> None: def reconcile_detached_sessions(store: PresetStore) -> None: - """Finalizes sessions whose agent completed while no CLI was attached - (graceful detach, or an ungraceful CLI death). This is what makes the saved - preset independent of a foreground process: any read command runs it, and - the work materializes from the on-disk report. - - Best-effort and parallel-safe — finalize takes an exclusive claim, and every - error is swallowed so the calling read command never fails. - """ + """Finalizes sessions whose agent completed while no CLI was attached, so + the saved preset never depends on a foreground process staying alive. + + Best-effort and parallel-safe: finalize takes an exclusive claim, and every + error is swallowed so the calling read command never fails.""" for session in iter_agent_sessions(): if _is_reconcilable(session.read_manifest()): _reconcile_session(session, store) @@ -256,8 +242,6 @@ def reconcile_detached_sessions(store: PresetStore) -> None: def _is_reconcilable(manifest: dict[str, Any]) -> bool: # An orphaned session (no live owner) whose agent left a completion report. - # A session interrupted mid-work has no report and stays resumable; one - # stopped *after* the agent finished is finalized by `stop` itself, not here. # Sessions created before finalize context was persisted lack `project` and # are skipped — they finalize interactively via `logs -f`. return ( @@ -274,10 +258,8 @@ def _reconcile_session(session: PresetAgentSession, store: PresetStore) -> None: api = Client.from_config(project_name=str(manifest.get("project") or "")) except Exception: # noqa: BLE001 — offline/misconfigured must not break the read command return - # follow_preset takes the finalize claim (so a concurrent `logs -f` - # or reconcile can't double-finalize), records the terminal status itself, - # and leaves the session intact on a transient error. Every outcome is silent - # here — the result shows in the list that follows. + # follow_preset records the terminal status and is claim-safe; suppress every + # error and stay silent here since the result shows up in the list that follows. with suppress(Exception): follow_preset( api=api, @@ -323,8 +305,6 @@ def stop_preset_session(api: Client, preset_id: str) -> None: def _stop_active_session_runs(api: Client, session: PresetAgentSession) -> None: - """Stops the session's non-terminal runs (with a spinner), like `dstack - stop`. Keeping a trial instance warm for resume is the detach path, not this.""" names = _load_submitted_run_names(session.runs_path) active = [] for name in names: @@ -343,9 +323,9 @@ def _stop_active_session_runs(api: Client, session: PresetAgentSession) -> None: def _resolve_preset_env( configuration: PresetConfiguration, *, strict: bool = True ) -> PresetConfiguration: - """Resolves `EnvSentinel` entries from the process environment. Non-strict - drops unresolvable entries instead of raising — for attach, where env values - only feed redaction and the agent already runs.""" + """Non-strict mode drops unresolvable `EnvSentinel` entries instead of + raising — for attach, where env values only feed redaction and the agent + already runs.""" configuration = configuration.model_copy(deep=True) resolved: dict[str, str] = {} for key, value in configuration.env.items(): @@ -361,6 +341,42 @@ def _resolve_preset_env( return configuration +def resolve_previous_sessions(refs: Sequence[str]) -> tuple[PresetAgentSession, ...]: + sessions: list[PresetAgentSession] = [] + for ref in refs: + try: + session = load_agent_session(resolve_session_ref(ref)) + except CLIError: + raise CLIError(f"Previous session {ref!r} does not exist") + if all(existing.preset_id != session.preset_id for existing in sessions): + sessions.append(session) + included = {session.preset_id for session in sessions} + for session in sessions: + manifest = session.read_manifest() + if manifest.get("status") == "running" and session_process_alive(manifest): + raise CLIError( + f"Previous session {session.preset_id} is still running;" + " wait for it to finish or stop it" + ) + for parent in manifest.get("previous") or []: + if parent not in included: + warn( + f"{session.preset_id} was created with --previous {parent}," + " which is not included" + ) + return tuple(sessions) + + +def _load_pinned_previous_sessions(ids: Sequence[str]) -> tuple[PresetAgentSession, ...]: + sessions = [] + for preset_id in ids: + try: + sessions.append(load_agent_session(preset_id)) + except CLIError: + warn(f"Previous session {preset_id} no longer exists; keeping its copied records") + return tuple(sessions) + + def create_preset( *, api: Client, @@ -372,6 +388,7 @@ def create_preset( resume_session: Optional[PresetAgentSession] = None, user_prompt: Optional[str] = None, allowed_fleets: Optional[tuple[str, ...]] = None, + previous: Sequence[PresetAgentSession] = (), ) -> PresetCreateResult: agent_session = resume_session or create_preset_agent_session(configuration, debug=debug) try: @@ -388,6 +405,7 @@ def create_preset( resume=resume_session is not None, user_prompt=user_prompt, allowed_fleets=allowed_fleets, + previous=previous, ) ) except KeyboardInterrupt: @@ -415,6 +433,7 @@ class _CreationSetup: user_prompt: Optional[str] initial_resume_session_id: Optional[str] write_constraints: bool # True only for fresh creations + previous: tuple[str, ...] = () # session IDs, pinned at creation def _fresh_setup( @@ -424,6 +443,7 @@ def _fresh_setup( build_name: Optional[str], allowed_fleets: Optional[tuple[str, ...]], user_prompt: Optional[str], + previous: Sequence[PresetAgentSession] = (), ) -> _CreationSetup: if allowed_fleets is None: allowed_fleets = _get_allowed_fleets(api, configuration) @@ -431,6 +451,10 @@ def _fresh_setup( raise CLIError(_NO_FLEETS_ERROR) auth = get_claude_auth() workspace = create_agent_workspace(agent_session) + previous_ids = tuple(session.preset_id for session in previous) + if previous_ids: + agent_session.update_manifest(previous=list(previous_ids)) + install_previous_records(workspace, previous) build_name = build_name or _get_build_name( configuration.name, configuration.model.api_model_name, agent_session.preset_id ) @@ -442,6 +466,7 @@ def _fresh_setup( user_prompt=user_prompt, initial_resume_session_id=None, write_constraints=True, + previous=previous_ids, ) @@ -467,6 +492,9 @@ def _resume_setup( claude_session_id = manifest.get("claude_session_id") if isinstance(claude_session_id, str) and claude_session_id: initial_resume_session_id = claude_session_id + previous_ids = tuple(manifest.get("previous") or []) + if previous_ids: + install_previous_records(workspace, _load_pinned_previous_sessions(previous_ids)) return _CreationSetup( auth=auth, workspace=workspace, @@ -475,6 +503,7 @@ def _resume_setup( user_prompt=user_prompt, initial_resume_session_id=initial_resume_session_id, write_constraints=False, + previous=previous_ids, ) @@ -491,6 +520,7 @@ def _attach_setup( user_prompt=None, initial_resume_session_id=None, write_constraints=False, + previous=tuple(agent_session.read_manifest().get("previous") or []), ) @@ -508,6 +538,7 @@ async def _create_preset( wait_for_run_stop: bool = True, user_prompt: Optional[str] = None, allowed_fleets: Optional[tuple[str, ...]] = None, + previous: Sequence[PresetAgentSession] = (), ) -> PresetCreateResult: source_configuration = source_configuration or configuration if attach: @@ -516,7 +547,7 @@ async def _create_preset( setup = _resume_setup(agent_session, build_name, user_prompt) else: setup = _fresh_setup( - api, configuration, agent_session, build_name, allowed_fleets, user_prompt + api, configuration, agent_session, build_name, allowed_fleets, user_prompt, previous ) # Record ownership + the finalize context (project, keep-service) so a later # detached reconcile can complete this session from disk alone. @@ -557,6 +588,7 @@ async def _create_preset( prompt = get_preset_agent_system_prompt( user_prompt=setup.user_prompt, baseline=configuration.effective_baseline, + previous=", ".join(setup.previous) if setup.previous else None, ) if setup.write_constraints: if setup.user_prompt: @@ -567,9 +599,11 @@ async def _create_preset( allowed_fleets=setup.allowed_fleets, ) setup.workspace.constraints_path.write_text(constraints_text, encoding="utf-8") + # A second, persistent copy: the workspace above is deleted with the run, + # while the listing and `--previous` read constraints from the session dir. + agent_session.write_constraints(constraints_text) if agent_session.debug: agent_session.write_prompt(prompt) - agent_session.write_constraints(constraints_text) if setup.auth is not None: agent_session.write_agent_info(setup.auth) try: @@ -607,6 +641,8 @@ async def _create_preset( run=run, preset_configuration=source_configuration, report=report, + workspace_path=setup.workspace.path, + session_path=agent_session.path, preset_id=agent_session.preset_id or None, name=claimed_session_name(agent_session.read_manifest()), ) @@ -640,9 +676,8 @@ async def _create_preset( cleanup_error = str(e) if cleanup_error is not None: - # The preset is already saved by this point; a failed cleanup only means - # trial runs may still be running. Warn rather than fail the (successful) - # session — otherwise a transient blip would discard completed work. + # The preset is already saved; a failed cleanup only means trial runs may + # still be running. Warn rather than fail — else a blip discards the work. if agent_session.echo: warn(f"Failed to stop preset creation runs: {cleanup_error}") assert preset is not None @@ -677,15 +712,13 @@ def _finish_agent_session( def _close_agent_session(session: PresetAgentSession, status: str) -> None: - """Records the terminal status and removes the workspace alias.""" _finish_agent_session(session, status) remove_agent_workspace(session) def _detach_agent_session(session: PresetAgentSession) -> None: - """Releases ownership but leaves the agent running — it stays visible and - reconcilable in `dstack preset`. Silent: `logs -f` calls this on Ctrl+C, and - a viewer that just stops watching shouldn't announce anything.""" + """Releases ownership but leaves the agent running (still reconcilable in + `dstack preset`), and stays silent since `logs -f` calls this on Ctrl+C.""" session.update_manifest(pid=None) @@ -693,7 +726,7 @@ def _stop_or_detach_agent_session( session: PresetAgentSession, api: Optional[Client] = None ) -> None: """`create` interrupt: stop the session, or detach and leave the agent - working — it stays visible as a running session in `dstack preset`.""" + working as a running session in `dstack preset`.""" manifest = session.read_manifest() agent_alive = session_process_alive({**manifest, "pid": None}) stop = True @@ -737,7 +770,6 @@ def _get_build_name(name: Optional[str], model_name: str, suffix: str) -> str: 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(): @@ -783,7 +815,6 @@ def find_preset_name_holders(store: PresetStore, name: str) -> PresetNameHolders def reassign_preset_name(store: PresetStore, holders: PresetNameHolders) -> None: - """Releases the name from every holder so a new preset can claim it.""" if holders.preset is not None: store.release_name(holders.name) for session in holders.sessions: @@ -791,8 +822,7 @@ def reassign_preset_name(store: PresetStore, holders: PresetNameHolders) -> None def plan_preset(*, api: Client, configuration: PresetConfiguration) -> tuple[str, ...]: - """Resolves the allowed fleets and shows what the agent will have to work - with — Project, User, the effective fleets, and their offers. Agent-free.""" + """Agent-free preview of the fleets and offers the agent would be given.""" allowed_fleets = _get_allowed_fleets(api, configuration) if not allowed_fleets: raise CLIError(_NO_FLEETS_ERROR) @@ -917,8 +947,7 @@ async def _cleanup_runs( return deadline = asyncio.get_running_loop().time() + _RUN_STOP_TIMEOUT_SECONDS pending = set(active_names) - # The same spinner the stop command shows: without it the CLI looks hung - # for however long the runs take to terminate. + # Without a spinner the CLI looks hung while the runs terminate. spinner = console.status("Stopping runs...") if agent_session.echo else nullcontext() with spinner: while pending: diff --git a/src/dstack/_internal/cli/services/presets/output.py b/src/dstack/_internal/cli/services/presets/output.py index a15b63bdf..b361c8084 100644 --- a/src/dstack/_internal/cli/services/presets/output.py +++ b/src/dstack/_internal/cli/services/presets/output.py @@ -25,10 +25,6 @@ def _format_status(status: str) -> str: def _verifying(session: dict[str, Any]) -> bool: - """Whether the agent has moved on to the final service. Read from the session's - verification records rather than inferred from a spent trial budget, which - misses every run that stopped early. An attempt that failed still counts: the - agent is picking the next trial to verify, not trialing again.""" return isinstance(session.get("verification"), dict) @@ -36,9 +32,6 @@ def _verifying(session: dict[str, Any]) -> bool: def _format_trial_spark(session: Optional[dict[str, Any]]) -> str: - """One glyph per trial, scaled within the run: the shape of the search. - A red `·` marks a trial that produced no benchmark at all; a yellow bar marks - one that measured but broke a constraint, since its number is real.""" if not isinstance(session, dict): return "" trials = session.get("trials") @@ -51,8 +44,9 @@ def _format_trial_spark(session: Optional[dict[str, Any]]) -> str: values = [v for v in series if isinstance(v, (int, float))] if not values: return "·" * len(series) - low, high = min(values), max(values) - span = high - low + high = max(values) + passed = [v for v, f in zip(series, failed) if isinstance(v, (int, float)) and not f] + best = max(passed) if passed else max(values) out = [] for value, is_failed in zip(series, failed): if not isinstance(value, (int, float)): @@ -60,31 +54,28 @@ def _format_trial_spark(session: Optional[dict[str, Any]]) -> str: continue glyph = ( _SPARK_BLOCKS[-1] - if span <= 0 - else _SPARK_BLOCKS[round((value - low) / span * (len(_SPARK_BLOCKS) - 1))] + if high <= 0 + else _SPARK_BLOCKS[round(max(value, 0) / high * (len(_SPARK_BLOCKS) - 1))] ) - # The best trial is the answer the run found; everything else is context. - # Yellow, not red: the trial measured, its number is real, and only the - # constraint breach makes it unusable. Red is reserved for `·`, where - # nothing came back at all. if is_failed: - style = "gold1" + style = "gold1" if not passed and value >= best else "indian_red1" else: - style = "bold sea_green3" if value >= high else "secondary" + style = "bold sea_green3" if value >= best else "secondary" out.append(f"[{style}]{glyph}[/]") return "".join(out) -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.""" +def _format_trial_progress(session: Optional[dict[str, Any]], *, in_flight: bool = False) -> str: if not isinstance(session, dict): return "" trials = session.get("trials") trials_num = session.get("trials_num") if not isinstance(trials, dict) or not (trials.get("count") or isinstance(trials_num, int)): return "" - progress = str(trials.get("count") or 0) + count = trials.get("count") or 0 + if in_flight: + count = min(count + 1, trials_num) if isinstance(trials_num, int) else count + 1 + progress = str(count) if isinstance(trials_num, int): progress += f"/{trials_num}" return f" [secondary]({progress})[/]" @@ -113,16 +104,22 @@ def get_presets_table( limit: Optional[int] = None, ) -> Table: table = Table(box=None) + compact = not verbose table.add_column("ID", no_wrap=True) - table.add_column("BASE", no_wrap=True, style="secondary") - table.add_column("RESOURCES" if verbose else "GPU", style="secondary") - # CONSTRAINTS is the test that was asked for; BENCHMARK is the best trial under it. - table.add_column("CONSTRAINTS", no_wrap=True) + # Compact-view caps keep a long model name from starving the wrapping + # CONSTRAINTS and BENCHMARK columns below. + table.add_column("BASE", no_wrap=True, max_width=24 if compact else None, style="secondary") + table.add_column( + "RESOURCES" if verbose else "GPU", + no_wrap=compact, + max_width=18 if compact else None, + style="secondary", + ) + table.add_column("CONSTRAINTS", min_width=len("io=1K/1K"), overflow="fold") table.add_column("BENCHMARK", min_width=len("tps=1"), overflow="fold") - # The search shape, one glyph per trial. Unlabelled: it reads on sight. table.add_column("", no_wrap=True) - table.add_column("STATUS", no_wrap=True) - table.add_column("SUBMITTED", no_wrap=True, style="secondary") + table.add_column("STATUS") + table.add_column("SUBMITTED", style="secondary") if verbose: table.add_column("NAME", no_wrap=True, style="secondary") presets_by_base: dict[str, list[Preset]] = defaultdict(list) @@ -140,10 +137,8 @@ def get_presets_table( model = str(session.get("model") or "unknown") sessions_by_model[repo_to_base.get(model, model)].append(session) - # One flat list, newest first, as in `dstack ps`. The base is a column, so - # runs of different models still sort together by when they were submitted. - # Same contract as `dstack ps`: only active by default, or the single most - # recent row when nothing is active. `-a` and `-n` show everything. + # Same contract as `dstack ps`: one flat list, newest first (base is a column, + # so different models interleave); active-only by default, else the latest row. rows: list[tuple[str, Any, bool]] = [] for preset_list in presets_by_base.values(): rows += [(preset.created_at.isoformat(), preset, True) for preset in preset_list] @@ -180,7 +175,9 @@ def _add_session(table: Table, session: dict[str, Any], *, verbose: bool = False status_key = str(session.get("status", "")) if status_key == "running" and _verifying(session): status_key = "verifying" - status = _format_status(status_key) + _format_trial_progress(session) + status = _format_status(status_key) + _format_trial_progress( + session, in_flight=status_key == "running" + ) trials = session.get("trials") best = trials.get("best") if isinstance(trials, dict) else None # Nothing passed: fall back to the fastest attempt that did not. @@ -215,15 +212,10 @@ def _add_session(table: Table, session: dict[str, Any], *, verbose: bool = False max_ttft = constraints.get("max_ttft") if verbose and isinstance(max_ttft, (int, float)): objective.append(f"ttft<={_format_duration_ms(max_ttft)}") - # Stays empty until a trial has produced a benchmark: a run that has measured - # nothing yet has no best, and `n/a` is noise in a column of numbers. if isinstance(best, dict): tps = _format_number(best["tok_s"]) if objective: - # Per-user leads: aggregate rises with concurrency, so it makes rows at - # different concurrencies look better or worse than they serve. - # Same definition as `effective_per_user_tok_per_s`: the steady decode - # rate, not the aggregate divided by concurrency. + # Lead with per-user tok/s: comparable across rows regardless of concurrency. tpot_ms = best.get("tpot_ms") if isinstance(tpot_ms, (int, float)) and tpot_ms > 0: parts.append(f"tok/s/user={_format_number(1000 / tpot_ms)}") @@ -251,8 +243,6 @@ def _add_session(table: Table, session: dict[str, Any], *, verbose: bool = False "GPU": gpu, "RESOURCES": gpu, "": _format_trial_spark(session), - # The constraints are context for the number, so the whole cell recedes; - # the benchmark beside it is what the reader came for and stays bright. "CONSTRAINTS": f"[secondary]{' '.join(objective)}[/]" if objective else "", "BENCHMARK": benchmark, "STATUS": status, @@ -279,8 +269,11 @@ def _add_preset( "": _format_trial_spark(creation), "CONSTRAINTS": format_preset_objective( preset, - min_context_length=(creation or {}).get("constraints", {}).get("min_context_length"), - max_ttft=(creation or {}).get("constraints", {}).get("max_ttft"), + # Fall back to the creation record for presets saved before the preset + # itself carried the requested values. + min_context_length=preset.min_context_length + or (creation or {}).get("constraints", {}).get("min_context_length"), + max_ttft=preset.max_ttft or (creation or {}).get("constraints", {}).get("max_ttft"), verbose=verbose, ), "BENCHMARK": format_preset_benchmark(preset, verbose=verbose), @@ -308,10 +301,6 @@ def format_preset_objective( max_ttft: Optional[float] = None, verbose: bool = False, ) -> str: - """What was asked for. The context the configuration actually reached is a - result and sits next to the numbers; the context that was *required* is shown - here under `-v`. Two runs can share every constraint and still serve different - context lengths, so both are worth seeing.""" workload = preset.validations[0].benchmark.workload parts = [ f"io={_format_token_count(workload.input_tokens)}" @@ -323,18 +312,23 @@ def format_preset_objective( # Absent for presets saved before the creation record was consulted. if verbose and min_context_length is not None: parts.append(f"ctx>={_format_token_count(min_context_length)}") - # The latency ceiling only explains a number that is near it, so it waits for `-v`. if verbose and max_ttft is not None: parts.append(f"ttft<={_format_duration_ms(max_ttft)}") - # Context for the number, so the whole cell recedes. return f"[secondary]{' '.join(parts)}[/]" +def _breaches_constraints(preset: Preset) -> bool: + metrics = preset.validations[0].benchmark.metrics + if preset.max_ttft is not None and metrics.ttft_ms.p50 > preset.max_ttft: + return True + return preset.min_context_length is not None and ( + preset.context_length < preset.min_context_length + ) + + def format_preset_benchmark(preset: Preset, *, verbose: bool = False) -> str: benchmark = preset.validations[0].benchmark metrics = benchmark.metrics - # The workload and context define the number, so they are always shown next - # to it: two presets are comparable only when all three match. parts = [ f"tok/s/user={_format_number(benchmark.effective_per_user_tok_per_s)}", ] @@ -344,12 +338,13 @@ def format_preset_benchmark(preset: Preset, *, verbose: bool = False) -> str: f"ttft={_format_duration_ms(metrics.ttft_ms.p50)}", f"ctx={_format_token_count(preset.context_length)}", ] - return " ".join(parts) + text = " ".join(parts) + if _breaches_constraints(preset): + return f"[secondary]*{text}[/]" + return text def _format_duration_ms(value: float) -> str: - """Milliseconds below a second, seconds above it. A bare `4152` reads as small - until you notice the unit; `4.15s` does not.""" # 999.6 rounds to 1000, which must read as 1s rather than 1000ms. if value < 999.5: return f"{_format_number(value)}ms" @@ -357,6 +352,7 @@ def _format_duration_ms(value: float) -> str: def _format_token_count(value: int) -> str: + """Abbreviates only exact multiples of 1024/1024², so 2048 becomes "2K" but 2050 stays "2050".""" for divisor, suffix in ((1024 * 1024, "M"), (1024, "K")): if value >= divisor and value % divisor == 0: return f"{value // divisor}{suffix}" diff --git a/src/dstack/_internal/cli/services/presets/presets.py b/src/dstack/_internal/cli/services/presets/presets.py index 983c8ba82..2801e8e24 100644 --- a/src/dstack/_internal/cli/services/presets/presets.py +++ b/src/dstack/_internal/cli/services/presets/presets.py @@ -26,6 +26,9 @@ def build_preset( model: str, context_length: int, benchmark: PresetBenchmark, + trial: Optional[int] = None, + min_context_length: Optional[int] = None, + max_ttft: Optional[int] = None, preset_id: Optional[str] = None, name: Optional[str] = None, ) -> Preset: @@ -45,6 +48,9 @@ def build_preset( id=preset_id or make_preset_id(service, context_length=context_length), model=model, context_length=context_length, + trial=trial, + min_context_length=min_context_length, + max_ttft=max_ttft, created_at=get_current_datetime(), service=service, validations=[validation], @@ -73,6 +79,13 @@ def preset_to_data(preset: Preset) -> dict[str, Any]: **({"name": preset.name} if preset.name else {}), "model": preset.model, "context_length": preset.context_length, + **({"trial": preset.trial} if preset.trial is not None else {}), + **( + {"min_context_length": preset.min_context_length} + if preset.min_context_length is not None + else {} + ), + **({"max_ttft": preset.max_ttft} if preset.max_ttft is not None else {}), "created_at": preset.created_at.isoformat(), "service": service_configuration_to_preset_data(preset.service), "validations": [ @@ -85,6 +98,9 @@ def preset_to_data(preset: Preset) -> dict[str, Any]: def service_configuration_to_preset_data( configuration: ServiceConfiguration, ) -> dict[str, Any]: + """The canonical service form used for preset identity and hashing: drops + type/name/gateway/profile fields, serializes env as sorted `key=value` + strings, and removes empty collections.""" service_data = json.loads(configuration.model_dump_json(exclude_none=True)) service_data.pop("type", None) service_data.pop("name", None) diff --git a/src/dstack/_internal/cli/services/presets/prompt.py b/src/dstack/_internal/cli/services/presets/prompt.py index e565fc2fd..4ede67763 100644 --- a/src/dstack/_internal/cli/services/presets/prompt.py +++ b/src/dstack/_internal/cli/services/presets/prompt.py @@ -1,47 +1,191 @@ import re +from dataclasses import dataclass, field from pathlib import Path -from typing import Optional +from typing import Optional, Union 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) +# `` ... `` ... `` renders one branch: +# the first when the variable NAME is set (with `{NAME}` interpolated in it), +# the `else` branch (optional) otherwise. Blocks nest, and markers may sit +# inline within a line. A marker alone on its line disappears with the whole +# line, and the branch body is dedented by the indentation shared by every +# line of it; a body whose lines do not share one exact indentation is kept +# as written. +_MARKER_PATTERN = re.compile(r"") +_IF_PATTERN = re.compile(r"if\s+(\w+)") -# `` is a note for maintainers and is dropped before the agent sees -# the document. Any other comment is left alone, so that a plain `` or a -# malformed directive stays visible instead of disappearing silently. +# `` is a maintainer note, dropped before the agent sees the +# document. The `!` is required, so ordinary `` comments are preserved. _NOTE_PATTERN = re.compile(r"\n?", re.DOTALL) +@dataclass +class _Text: + value: str + at_line_start: bool + + +@dataclass +class _IfNode: + name: str + # Indents of this node's own-line markers; they sit at the enclosing + # body's level, so the enclosing branch counts them as its lines. + marker_indents: list[str] = field(default_factory=list) + # An inline-opened block has no line structure of its own to dedent. + opened_on_own_line: bool = False + then: list[Union[_Text, "_IfNode"]] = field(default_factory=list) + otherwise: list[Union[_Text, "_IfNode"]] = field(default_factory=list) + + +def _parse_directives( + text: str, variables: dict[str, Optional[str]] +) -> list[Union[_Text, _IfNode]]: + """Both branches are always parsed, so an unknown variable cannot hide + behind a flag combination.""" + root: list[Union[_Text, _IfNode]] = [] + stack: list[_IfNode] = [] + in_else: list[bool] = [] + + def current() -> list[Union[_Text, _IfNode]]: + if not stack: + return root + return stack[-1].otherwise if in_else[-1] else stack[-1].then + + position = 0 + at_line_start = True + for match in _MARKER_PATTERN.finditer(text): + line_start = text.rfind("\n", 0, match.start()) + 1 + line_end = text.find("\n", match.end()) + line_end = len(text) if line_end == -1 else line_end + indent = text[line_start : match.start()] + own_line = not indent.strip() and not text[match.end() : line_end].strip() + if own_line: + # The marker's whole line goes away, indentation and newline included. + run_end, next_position, next_at_line_start = ( + line_start, + min(line_end + 1, len(text)), + True, + ) + else: + run_end, next_position, next_at_line_start = match.start(), match.end(), False + if run_end > position: + current().append(_Text(text[position:run_end], at_line_start)) + position, at_line_start = next_position, next_at_line_start + directive = match.group(1).strip() + if directive == "else": + if not stack or in_else[-1]: + raise CLIError("`else` without a matching `if` in the agent system prompt") + in_else[-1] = True + if own_line: + stack[-1].marker_indents.append(indent) + elif directive == "end": + if not stack: + raise CLIError("`end` without a matching `if` in the agent system prompt") + if own_line: + stack[-1].marker_indents.append(indent) + stack.pop() + in_else.pop() + else: + if (if_match := _IF_PATTERN.fullmatch(directive)) is None: + raise CLIError(f"Invalid directive {directive!r} in the agent system prompt") + name = if_match.group(1) + if name not in variables: + raise CLIError(f"Unknown variable {name!r} in the agent system prompt") + node = _IfNode(name=name, opened_on_own_line=own_line) + if own_line: + node.marker_indents.append(indent) + current().append(node) + stack.append(node) + in_else.append(False) + if stack: + raise CLIError("Unclosed `if` in the agent system prompt") + if position < len(text): + root.append(_Text(text[position:], at_line_start)) + return root + + +def _branch_indent(parts: list[Union[_Text, _IfNode]]) -> str: + """The one exact indentation shared by every line of the branch, or `""` + when there is none. A nested block's interior belongs to that block, but + its own-line markers sit at this branch's level and count.""" + indents = [] + for index, part in enumerate(parts): + if isinstance(part, _IfNode): + indents += part.marker_indents + continue + lines = part.value.split("\n") + for line_index, line in enumerate(lines): + if line_index == 0 and not part.at_line_start: + continue # a fragment continuing a line already counted + is_last = line_index == len(lines) - 1 + continues_inline = ( + is_last + and index + 1 < len(parts) + and isinstance(parts[index + 1], _IfNode) + and not parts[index + 1].opened_on_own_line + ) + if line.strip(): + indents.append(line[: len(line) - len(line.lstrip())]) + elif line and continues_inline: + # A whitespace-only fragment whose line is an inline block. + indents.append(line) + if indents and all(indent == indents[0] for indent in indents): + return indents[0] + return "" + + +def _render_branch( + parts: list[Union[_Text, _IfNode]], + variables: dict[str, Optional[str]], + applied: set[str], + dedent: bool, +) -> str: + indent = _branch_indent(parts) if dedent else "" + rendered: list[str] = [] + for part in parts: + if isinstance(part, _Text): + value = part.value + if " - `shared_prefix_tokens`: how many of `input_tokens` are identical in every request. `0` means every request is fully unique. - `baseline`: whether the first trial must be a baseline rather than an - optimization attempt; see `# Trials`. + optimization attempt; see `# Trials (Main Section)`. - `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 @@ -53,21 +53,26 @@ concurrencies instead of one.--> 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, patching serving framework source code, or P/D -disaggregation setups. - +Docker image and dependencies, the serving framework parameters, patch the +serving framework source code, generate custom kernels, and patch drivers. + + + Do not use P/D disaggregation setups, + unless `## Additional instructions` explicitly allows it. + + Do not use P/D disaggregation setups. + + + + + ## Additional instructions + + ``` + {prompt} + ``` + - ## CLI And Skills All trials and the final verification are done using `dstack`. This includes @@ -88,6 +93,10 @@ follow it the same way. Files provided to you in the workspace root (read them; never edit them): - `constraints.json`: the effective constraints; see `# Constraints`. + +- `previous/`: results of previous sessions; + see `## Previous Sessions`. + Files you are expected to maintain in the workspace root: @@ -95,17 +104,13 @@ Files you are expected to maintain in the workspace root: `# Runs`. - `progress.jsonl`: progress messages, written through the `progress` helper; see `# Progress`. -- `trials.jsonl`: the append-only record of completed trials; see `# Trials`. -- `verifications.jsonl`: the append-only record of final service attempts; - see `# Final Service`. +- `trials/`: one directory per trial; see `# Trials (Main Section)`. +- `service/`: one directory per final service attempt; see `# Final Service`. - `final_report.json`: the final report; see `# Final Report`. -You may create any other working files (run YAML files, benchmark output, -notes) in the workspace, and only there: do not deliberately save files -elsewhere on this machine. Incidental writes made by the tools you run -(caches, temporary files, SSH configuration) are fine wherever those tools -keep them. Files inside running `dstack` tasks or services are not subject -to this rule. +On this machine, do not deliberately create, change, or delete files +outside the workspace. Inside running `dstack` tasks and services, write +whatever the work needs. # Runs @@ -165,15 +170,30 @@ how to get better performance than the previous trials. Sometimes it is worth continuing to improve a previous trial's idea, but when that risks settling into a local optimum, search for a substantially different approach rather than tweaking parameters further. - - - + + + Before starting trials, analyze the previous sessions' results provided in + `previous/` (sessions: {previous}; see `## Previous Sessions`). The + objective for this session is to significantly improve them. + + + Note, the first trial is a baseline rather than an optimization attempt: + reproduce the best previous trial (only if it's comparable, e.g. shares + the same constraints). If no previous trial is comparable, report it via + `progress` (see `# Progress`) and serve the model the way the chosen + serving framework recommends for this model and hardware. Change only + what is necessary to make it run, and report each such change via + `progress` (see `# Progress`). + + + + Note, the first trial is a baseline rather than an optimization attempt: + serve the model the way the chosen serving framework recommends for this + model and hardware. Change only what is necessary to make it run, and + report each such change via `progress` (see `# Progress`). + + + Trial ideas must not rely only on what you already know. Research what limits performance and how to improve it for the chosen model, serving framework, and hardware. Actively seek credible and recent sources: benchmarks, newly @@ -196,6 +216,9 @@ trial and whenever a benchmark exposes a bottleneck. Don't skip profiling the serving engine, especially if it could help you find an idea for significantly improving the current numbers. +When a trial starts, create its directory `trials//`, where `` is the +trial number: trials are numbered from 1 in the order they start. + 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. @@ -208,25 +231,36 @@ 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. A trial is also failed when its benchmark does not meet -the constraints (see `# Constraints`). When a trial that changed several things -fails, be mindful of which specific change was the root cause. - -Each `trials.jsonl` record is one JSON line with exactly four fields: +the final performance into a complete `dstack` task configuration with +exact commands, and write it to `trials//task.dstack.yml` (do not +create this file until the trial is completed). 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 when `trials//task.dstack.yml` is applied +(not `sleep infinity`). If the trial required patching the serving +framework source code, generating custom kernels, or patching drivers, +make sure to include the required exact patches into +`trials//task.dstack.yml` via `files` (see `## Patching Framework`). + + +Once `trials//task.dstack.yml` is written, write the corresponding +benchmark results (see `## Benchmark` for the structure) to +`trials//trial.json`: the presence of `trials//trial.json` is what +marks trial `` completed. The benchmark +may be skipped in one case only: you failed to make the configuration run +at all — a failed trial. `trials//task.dstack.yml` may be skipped in +one case only: you failed to get benchmark results at all. A trial is +also failed when its benchmark does not meet the constraints (see +`# Constraints`). When a trial that changed several things fails, be +mindful of which specific change was the root cause. + +`trials//trial.json` is one JSON object with these fields and no others: ``` -{"task": {...}, "resources": {...}, "context_length": ..., "benchmark": {...}, "learned": ..., "failed": ...} +{"resources": {...}, "context_length": ..., "benchmark": {...}, "learned": ..., "failed": ...} ``` -- `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 @@ -242,8 +276,8 @@ Each `trials.jsonl` record is one JSON line with exactly four fields: `null` only when the benchmark couldn't be done at all. - `learned`: the major things this trial taught you that you did not know before it ran. Required for every trial, including a failed one. -- `failed`: `true` if the benchmark broke a constraint such as `max_ttft` or - `min_context_length`, absent otherwise. +- `failed`: `true` if the configuration never ran or the benchmark broke a + constraint such as `max_ttft` or `min_context_length`, absent otherwise. You're expected to do exactly `trials_num` trials (see `# Constraints`). @@ -290,7 +324,7 @@ never part of the measured metrics. All verification and benchmark requests must succeed. Record every benchmark using the following structure and field names — -trial benchmarks in `trials.jsonl`, the final benchmark as +trial benchmarks in `trials//trial.json`, the final benchmark as `final_report.json.benchmark` (values are illustrative): ```json @@ -323,11 +357,56 @@ warmup requests. Never invent missing values. After each benchmark, find the largest context the configuration handles by sending real requests, and record it: for a trial, as the `context_length` -field of its `trials.jsonl` record; for the final benchmark, as +field of `trials//trial.json`; for the final benchmark, as `final_report.json.context_length`. Stopping at the required minimum is not enough. + + ## Previous Sessions + + Results of sessions that ran before this one are provided in + `previous/`, exclusively so you can see what was already tried and how it + worked. Each `previous//` holds one session's results, in the same + format you write yours: `constraints.json`, `trials//` with + `trial.json`, `task.dstack.yml`, and `patches/`, `service//`, and + `final_report.json`. + + +## Patching Framework + +If the trial required patching the serving framework source code, +generating custom kernels, or patching drivers (see `# Constraints`), +when you save `trials//task.dstack.yml`, you must replicate these +exact patches. + +To do this, you must save the required patches in the +`trials//patches` directory (next to `trials//task.dstack.yml`), +and refer to them from `trials//task.dstack.yml` in the `files` +property. This will mount them inside the container. + +A patch is a unified diff against the file it changes. + +Example: + +```yaml +files: + - patches/vllm/model_executor/layers/fused_moe/fused_moe.py.patch:/patches/vllm/model_executor/layers/fused_moe/fused_moe.py.patch +``` + +Then, patches can be applied with `patch` (exactly when needed) from the +`commands`. + +Make sure to include only the patches that are required, and avoid +including unnecessary ones. + +This "patching framework" will allow you to replicate the fixes made +during the interactive SSH session via `trials//task.dstack.yml`. + +And, since writing `trials//task.dstack.yml` is done after the +interactive trial is completed, it's especially important to review that +patches are correct (and will exactly replicate the result). + # Task Usage Trials are done entirely using `dstack` tasks. For maximum efficiency, it is a @@ -371,7 +450,9 @@ SSH fleets can be treated as VM-based backends as they support both idle instanc # Final Service Once the trials are over, pick the best trial that has not been verified yet -and submit its configuration as a `dstack` service. Make it work with only +and submit its configuration as a `dstack` service. If there is no remaining +non-failed trial, pick the best failed trial that has a benchmark and has not +been verified yet. Make it work with only minor tweaks if needed; do not change the important decisions made during the trial. Set the service `model` name to the client-facing model name from `constraints.json` (see `# Constraints`). `model` is required: it also enables @@ -379,19 +460,36 @@ the trial. Set the service `model` name to the client-facing model name from serve requests. If the service never passes the probe, treat that as a real failure of the configuration, not something to work around by removing `model`. -Record every attempt in `verifications.jsonl`, append-only: one line -immediately after submitting the service, one when the attempt ends (values -are illustrative): +Record every attempt in its own directory `service//`, where `` is +the attempt number: attempts are numbered from 1 in the order they are +submitted. Immediately after submitting the service, create `service//` +and write the submitted service YAML to `service//service.dstack.yml`. +If the service is based on a trial that required patches (see +`## Patching Framework`), save the required patches in +`service//patches` and refer to them from +`service//service.dstack.yml` in the `files` property. +When the attempt ends, write `service//verification.json`, one JSON +object with these fields and no others (values are illustrative): + +```json +{"trial": 3, "run_name": "qwen-preset-2", "status": "verified"} +``` + +or ```json -{"trial": 3, "run_name": "qwen-preset-2", "status": "verifying"} -{"trial": 3, "run_name": "qwen-preset-2", "status": "failed", "reason": "..."} -{"trial": 2, "run_name": "qwen-preset-3", "status": "verifying"} -{"trial": 2, "run_name": "qwen-preset-3", "status": "verified"} +{"trial": 2, "run_name": "qwen-preset-3", "status": "failed", "reason": "..."} ``` -`trial` is the 1-based line number of that trial in `trials.jsonl`. Keep -`reason` to one sentence. +- `trial`: the `` of the trial being verified. +- `run_name`: the run name of the attempt's `dstack` service. +- `status`: `verified`, or `failed` when the attempt ended without a + verified service. +- `reason`: required when `status` is `failed`, absent otherwise: the + reason of the failure. + +The presence of `service//verification.json` is what marks attempt +`` finished. 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 @@ -432,15 +530,17 @@ references in `final_report.json.service_yaml`; use environment variable names o # Final Report `final_report.json` may contain only `success`, `run_id`, `run_name`, -`service_yaml`, `base`, `model`, `context_length`, `benchmark`, and -`failure_summary`. +`service_yaml`, `trial`, `base`, `model`, `context_length`, `benchmark`, +and `failure_summary`. -On success, include exactly: +On success (even if you had to pick a failed trial because no non-failed +trial remained), include exactly: - `success`: `true` - `run_id`: the final verified service run ID - `run_name`: the final verified service run name - `service_yaml`: the full YAML of the verified final service +- `trial`: the `` of the verified trial - `base`: the base model repo, determined by the rules below - `model`: the exact repo/path loaded by the final service command - `context_length`: the largest context verified for the final service, as @@ -457,7 +557,7 @@ Set `final_report.json.base` as follows: to `model.repo`. - Do not infer `final_report.json.base` only from the repo name. -On failure, include exactly: +On failure (no trial verification was successful), include exactly: - `success`: `false` - `failure_summary`: the reason a preset could not be created and any change @@ -471,7 +571,7 @@ Verify that `final_report.json` is correct and matches the required schema. Stop only after `final_report.json` is written, and the report submitted: either one final `dstack` service was verified and benchmarked, or the trials and unverified candidates were exhausted (see -`# Trials` and `# Final Service`). +`# Trials (Main Section)` 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, diff --git a/src/dstack/_internal/cli/services/presets/session.py b/src/dstack/_internal/cli/services/presets/session.py index eb66f8112..aa93f7fad 100644 --- a/src/dstack/_internal/cli/services/presets/session.py +++ b/src/dstack/_internal/cli/services/presets/session.py @@ -29,8 +29,10 @@ _PROGRESS_FILENAME = "progress.jsonl" _RUNS_FILENAME = "runs.jsonl" -_TRIALS_FILENAME = "trials.jsonl" -_VERIFICATIONS_FILENAME = "verifications.jsonl" +_TRIALS_DIRNAME = "trials" +_SERVICE_DIRNAME = "service" +_TRIAL_RESULT_FILENAME = "trial.json" +_VERIFICATION_RESULT_FILENAME = "verification.json" _CONSTRAINTS_FILENAME = "constraints.json" _FINAL_REPORT_FILENAME = "final_report.json" _SESSION_FILENAME = "session.json" @@ -38,8 +40,8 @@ class SessionBusyError(CLIError): - """Another live process owns the session — it is following or finalizing it. - Callers that only want to view can fall back to a read-only follow.""" + """Raised when another live process owns the session; view-only callers can + fall back to a read-only follow.""" @dataclass @@ -47,9 +49,8 @@ class PresetAgentSession: path: Path debug: bool preset_id: str = "" - # Whether progress lines echo to this process's console (a live attach), on - # top of always being recorded to agent.log. Background reconcile sets it - # False so finalizing a detached session stays silent on the read command. + # Background reconcile sets this False so finalizing a detached session stays + # silent on the read command; agent.log is written regardless. echo: bool = field(default=True, repr=False) _log_enabled: bool = field(default=True, init=False, repr=False) @@ -66,12 +67,12 @@ def runs_path(self) -> Path: return self.path / _RUNS_FILENAME @property - def trials_path(self) -> Path: - return self.path / _TRIALS_FILENAME + def trials_dir(self) -> Path: + return self.path / _TRIALS_DIRNAME @property - def verifications_path(self) -> Path: - return self.path / _VERIFICATIONS_FILENAME + def service_dir(self) -> Path: + return self.path / _SERVICE_DIRNAME def write_prompt(self, prompt: str) -> None: _write_private_text(self.path / "prompt.md", prompt + "\n") @@ -236,16 +237,15 @@ def _pid_alive(pid: Any, started_at: Any = None) -> bool: 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).""" + """True if either a live agent (possibly detached) or a live CLI (possibly + between agent retries) still owns the session.""" if _pid_alive(manifest.get("agent_pid"), manifest.get("agent_started_at")): return True pid = manifest.get("pid") if not isinstance(pid, int) or pid <= 0 or pid == os.getpid(): return False - # Guard the CLI pid with its start time too: after an ungraceful CLI death - # the OS can recycle the pid, and a bare pid_exists() would read a dead - # session as still owned (falsely blocking reconcile / follow). + # Guard the CLI pid with its start time: a recycled pid would otherwise read + # a dead session as still owned, falsely blocking reconcile / follow. return _pid_alive(pid, manifest.get("pid_started_at")) @@ -281,7 +281,6 @@ def load_attachable_agent_session(preset_id: str) -> PresetAgentSession: def load_agent_session(preset_id: str) -> PresetAgentSession: - """Loads a session of any status for read-only inspection (its log).""" path = get_presets_dir() / preset_id session = PresetAgentSession(path=path, debug=False, preset_id=preset_id) if not path.is_dir() or not session.read_manifest(): @@ -290,7 +289,6 @@ def load_agent_session(preset_id: str) -> PresetAgentSession: def print_session_log(session: PresetAgentSession) -> None: - """Prints the session's redacted progress log verbatim, no markup.""" try: content = session.log_path.read_text(encoding="utf-8") except OSError: @@ -308,9 +306,8 @@ def mark_session_owner( keep_service: Optional[bool] = None, claude_model: Optional[str] = None, ) -> None: - """Records this process as the session's owner (pid + start time) and, when - given, the finalize context a later detached reconcile needs (project and - keep-service intent). `None` fields are left untouched.""" + """Beyond recording ownership, stores the finalize context a later detached + reconcile needs; `None` fields are left untouched.""" fields: dict[str, Any] = { "status": "running", "pid": os.getpid(), @@ -326,8 +323,8 @@ def mark_session_owner( def session_report_exists(manifest: dict[str, Any]) -> bool: - """Whether the agent left a final report on disk — the durable completion - signal a detached session is finalized from.""" + """True once the agent has written final_report.json, marking a detached + session ready to finalize.""" workspace = manifest.get("workspace") if not isinstance(workspace, str) or not workspace: return False @@ -335,11 +332,10 @@ def session_report_exists(manifest: dict[str, Any]) -> bool: def try_claim_session(session: PresetAgentSession) -> Optional[int]: - """Takes an exclusive, kernel-held lock for the duration of a session's - finalization, so concurrent readers can't both finalize it. Returns an open - file descriptor to release via `release_session_claim`, or None if another - process holds it. The kernel drops the lock if the holder dies, so there are - no stale locks to reason about.""" + """Takes an exclusive kernel lock so two readers can't both finalize the + session; returns an fd to release via `release_session_claim`, or None if + another process holds it. The kernel drops the lock if the holder dies, so + there are no stale locks.""" try: fd = os.open(session.path / ".reconcile.lock", os.O_CREAT | os.O_RDWR, 0o600) except OSError: @@ -359,8 +355,6 @@ def release_session_claim(fd: Optional[int]) -> None: def _try_lock_fd(fd: int) -> bool: - """Non-blocking exclusive lock on an open fd; True if acquired, False if - another process holds it.""" if IS_WINDOWS: import msvcrt @@ -384,13 +378,13 @@ def _try_lock_fd(fd: int) -> bool: def claimed_session_name(manifest: dict[str, Any]) -> Optional[str]: - """The name this session holds.""" value = manifest.get("name") return value if isinstance(value, str) and value else None def iter_agent_sessions() -> Iterator[PresetAgentSession]: - """Yields a handle for every session directory under the presets dir.""" + """Skips dotfiles and `models--*` HuggingFace cache dirs that share the + presets directory but aren't sessions.""" root = get_presets_dir() if not root.is_dir(): return @@ -432,18 +426,17 @@ def list_agent_sessions() -> list[dict[str, Any]]: entry["id"] = path.name entry["name"] = claimed_session_name(manifest) entry["status"] = status - entry["trials"] = _summarize_session_trials(path / _TRIALS_FILENAME) - entry["verification"] = _read_last_session_verification(path / _VERIFICATIONS_FILENAME) + entry["trials"] = _summarize_session_trials(path / _TRIALS_DIRNAME) + entry["verification"] = _read_last_session_verification(path / _SERVICE_DIRNAME) entry["constraints"] = _read_session_constraints(path) entries.append(entry) return entries def _read_session_constraints(path: Path) -> dict[str, Any]: - """The objective the session was given. The session's own copy is read first: it - is written at creation and outlives the agent workspace, which is removed once - the session finishes. The workspace copy is the fallback, for sessions recorded - before the session-level copy existed.""" + """Reads the session's own copy first: it outlives the agent workspace (removed + once the session finishes). The workspace copy is a backward-compat fallback for + sessions recorded before the session-level copy existed.""" for candidate in ( path / _CONSTRAINTS_FILENAME, path / "workspace" / "w" / _CONSTRAINTS_FILENAME, @@ -457,49 +450,56 @@ def _read_session_constraints(path: Path) -> dict[str, Any]: return {} -def _read_last_session_verification(path: Path) -> Optional[dict[str, Any]]: - """The final service attempt in flight or last finished, from the session's - mirrored verification records. The last line wins: the agent appends one when - an attempt starts and another when it ends.""" +def _numbered_subdirs(path: Path) -> list[Path]: try: - lines = path.read_text(encoding="utf-8").splitlines() + entries = [entry for entry in path.iterdir() if entry.is_dir() and entry.name.isdigit()] except OSError: + return [] + return sorted(entries, key=lambda entry: int(entry.name)) + + +def _read_record(path: Path) -> Optional[dict[str, Any]]: + """None if the record file is missing or caught half-written; the copy is + retried, so treat None as transient, not final.""" + try: + record = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): return None - for line in reversed(lines): - try: - record = json.loads(line) - except json.JSONDecodeError: - continue - if isinstance(record, dict) and isinstance(record.get("status"), str): - return record - return None + return record if isinstance(record, dict) else None + + +def _read_last_session_verification(path: Path) -> Optional[dict[str, Any]]: + """An attempt whose result file has not appeared yet is still in flight + (reported as verifying).""" + attempts = _numbered_subdirs(path) + if not attempts: + return None + last = attempts[-1] + record = _read_record(last / _VERIFICATION_RESULT_FILENAME) + if record is not None and isinstance(record.get("status"), str): + return record + return {"status": "verifying"} 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 = [] + """A trial directory without `trial.json` is still in flight and is not + counted.""" + records = [] + for trial_dir in _numbered_subdirs(path): + record = _read_record(trial_dir / _TRIAL_RESULT_FILENAME) + if record is not None: + records.append(record) count = 0 best: Optional[dict[str, Any]] = None # The fastest trial that broke a constraint, shown only when nothing passed. best_failed: Optional[dict[str, Any]] = None # One entry per trial in order, `None` for a trial that produced no benchmark. series: list[Optional[float]] = [] - # Parallel to `series`: a trial that measured but broke a constraint. + # Parallel to `series`: True where the trial broke a constraint. failed: list[bool] = [] # Kept outside `best` so a run where nothing passed still shows what it ran on. gpu: Optional[str] = None - for line in lines: - try: - record = json.loads(line) - except json.JSONDecodeError: - continue - if not isinstance(record, dict): - continue - # One record per trial (the agent contract); trials may share a task, - # so task names must not be deduplicated. + for record in records: count += 1 benchmark = record.get("benchmark") failed.append(bool(record.get("failed"))) @@ -600,12 +600,13 @@ def _process_started_at(pid: int) -> Optional[float]: return None -def _write_private_text(path: Path, content: str) -> None: +def _write_private_bytes(path: Path, content: bytes) -> None: # Atomic tmp + fsync + replace (mkstemp already creates the file 0600), so # a crash mid-write cannot leave a truncated manifest or offsets file. + # Binary mode: mirrored files are copies, and text mode rewrites newlines. fd, temporary = tempfile.mkstemp(dir=path.parent, prefix=f".{path.name}.", suffix=".tmp") try: - with os.fdopen(fd, "w", encoding="utf-8") as f: + with os.fdopen(fd, "wb") as f: f.write(content) f.flush() os.fsync(f.fileno()) @@ -622,7 +623,11 @@ def _write_private_text(path: Path, content: str) -> None: with suppress(PermissionError): os.replace(temporary, path) return - path.write_text(content, encoding="utf-8") + path.write_bytes(content) finally: with suppress(FileNotFoundError): os.unlink(temporary) + + +def _write_private_text(path: Path, content: str) -> None: + _write_private_bytes(path, content.encode("utf-8")) diff --git a/src/dstack/_internal/cli/services/presets/store.py b/src/dstack/_internal/cli/services/presets/store.py index b48b350e4..a36dbd6a3 100644 --- a/src/dstack/_internal/cli/services/presets/store.py +++ b/src/dstack/_internal/cli/services/presets/store.py @@ -22,9 +22,8 @@ class PresetStore: - """Presets live at `//` — one directory per preset holding - the artifact (`preset.yaml`) next to the creation session internals. - Deleted presets are archived under `/.archive/`.""" + """One `//preset.yaml` per preset; delete archives the directory + under `/.archive/` instead of removing it.""" def __init__(self, root: Path | None = None) -> None: self.root = root or get_dstack_dir() / "presets" @@ -88,11 +87,9 @@ def find_by_name(self, name: str) -> Preset | None: return None def find_by_id_or_name(self, ref: str) -> Preset | None: - """Resolves a preset reference that may be an ID or a claimed name.""" return self.get(ref) or self.find_by_name(ref) def release_name(self, name: str) -> Preset | None: - """Releases `name` from the preset holding it, keeping the preset.""" preset = self.find_by_name(name) if preset is None: return None @@ -185,7 +182,7 @@ def _parse_preset_configuration(stream: TextIO) -> PresetConfiguration: def resolve_preset_prompt( configuration: PresetConfiguration, configuration_path: str ) -> str | None: - """The resolved user prompt text; file paths are relative to the configuration file.""" + """Prompt-file paths resolve relative to the configuration file's directory (cwd for stdin).""" if configuration.prompt is None: return None if isinstance(configuration.prompt, str): diff --git a/src/dstack/_internal/cli/services/presets/tail.py b/src/dstack/_internal/cli/services/presets/tail.py index 2bafbba1d..682e4caa2 100644 --- a/src/dstack/_internal/cli/services/presets/tail.py +++ b/src/dstack/_internal/cli/services/presets/tail.py @@ -7,9 +7,10 @@ from pathlib import Path from typing import Any, Callable, Optional, Sequence -from dstack._internal.cli.services.presets.redaction import redact +from dstack._internal.cli.services.presets.redaction import redact, redact_bytes from dstack._internal.cli.services.presets.session import ( PresetAgentSession, + _write_private_bytes, _write_private_text, print_preset_progress, ) @@ -17,9 +18,7 @@ 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.""" + """`readline()` over a growing file whose persisted offset lets a later attach resume exactly where the previous reader stopped.""" _POLL_SECONDS = 0.2 _MAX_CHUNK = 1024 * 1024 @@ -73,10 +72,9 @@ async def readline(self) -> bytes: class _OffsetStore: - """Persists tailer/mirror byte offsets so resumed sessions do not repeat - output. One instance serves the whole session — every reader and mirror - shares it with disjoint keys, and the exclusive session claim guarantees no - other process writes the file.""" + """Persists per-stream read offsets under disjoint keys; a thread lock + suffices because the session claim guarantees no other process writes this + file.""" def __init__(self, path: Path) -> None: self._path = path @@ -100,7 +98,6 @@ def set(self, key: str, value: int) -> None: def open_session_offsets(session: PresetAgentSession) -> _OffsetStore: - """The session's single offset store, shared by all its tailers.""" return _OffsetStore(session.path / ".offsets.json") @@ -146,8 +143,6 @@ def flush(self) -> None: class _RecordMirror: - """Mirrors a workspace record file into the persistent session directory, redacted.""" - def __init__( self, *, @@ -202,6 +197,68 @@ def flush(self) -> None: console.print(f"[warning]Could not mirror {self._target.name}: {e}[/]") +class _DirectoryMirror: + """Mirrors a directory by re-copying each whole file whose size or mtime + changed; a half-written file is simply recopied complete on a later flush.""" + + _MAX_FILE_BYTES = 8 * 1024 * 1024 + + def __init__( + self, + *, + source: Path, + target: Path, + redacted_values: Sequence[str], + echo: bool = True, + ) -> None: + self._source = source + self._target = target + self._redacted_values = redacted_values + self._echo = echo + self._copied: dict[Path, tuple[int, int]] = {} + self._warned: set[Path] = set() + + async def run(self) -> None: + while True: + # File IO runs off the event loop; see _FileLineReader.readline. + await asyncio.to_thread(self.flush) + await asyncio.sleep(1) + + def flush(self) -> None: + if not self._source.is_dir(): + return + for source_file in sorted(self._source.rglob("*")): + try: + if not source_file.is_file(): + continue + stat = source_file.stat() + except OSError: + continue + relative = source_file.relative_to(self._source) + if stat.st_size > self._MAX_FILE_BYTES: + if self._echo and relative not in self._warned: + self._warned.add(relative) + console.print( + f"[warning]Not mirroring {relative}: " + f"larger than {self._MAX_FILE_BYTES // (1024 * 1024)} MiB[/]" + ) + continue + signature = (stat.st_size, stat.st_mtime_ns) + if self._copied.get(relative) == signature: + continue + try: + content = source_file.read_bytes() + target = self._target / relative + target.parent.mkdir(mode=0o700, parents=True, exist_ok=True) + _write_private_bytes(target, redact_bytes(content, self._redacted_values)) + except OSError as e: + if self._echo and relative not in self._warned: + self._warned.add(relative) + console.print(f"[warning]Could not mirror {relative}: {e}[/]") + continue + self._copied[relative] = signature + + def _parse_progress(line: str) -> Optional[str]: try: value = json.loads(line) diff --git a/src/dstack/_internal/cli/services/presets/verify.py b/src/dstack/_internal/cli/services/presets/verify.py index 13a87990b..42a2b54c6 100644 --- a/src/dstack/_internal/cli/services/presets/verify.py +++ b/src/dstack/_internal/cli/services/presets/verify.py @@ -47,9 +47,8 @@ def load_preset_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. + # Redact known secrets; an unknown leaked token is still caught downstream by + # the command bearer-token check. report_data = redact_structure(report_data, redacted_values) try: report = AgentFinalReport.model_validate(report_data) @@ -65,14 +64,40 @@ def load_preset_agent_report( return report +def _rewrite_workspace_file_paths( + service: ServiceConfiguration, *, workspace_path: Path, session_path: Path +) -> None: + """Re-roots `files` onto the session's mirrored copies because the submission + workspace is deleted when the session ends; only `trials/` and `service/` are mirrored.""" + workspace_root = workspace_path.resolve() + for mapping in service.files: + try: + relative = Path(mapping.local_path).resolve().relative_to(workspace_root) + except ValueError: + raise CLIError( + f"Claude final service file '{mapping.local_path}' is outside the agent workspace" + ) + target = session_path / relative + if relative.parts[:1] not in (("trials",), ("service",)) or not target.exists(): + raise CLIError( + f"Claude final service file '{mapping.local_path}' has no mirrored copy" + f" at '{target}'" + ) + mapping.local_path = str(target) + + def build_verified_preset( *, run: Run, preset_configuration: PresetConfiguration, report: AgentFinalReport, + workspace_path: Optional[Path] = None, + session_path: Optional[Path] = None, preset_id: Optional[str] = None, name: Optional[str] = None, ) -> Preset: + """Cross-checks the agent's self-reported final report against the actual run + and service state before trusting it to build a preset.""" if run.id != report.run_id or run.run_spec.run_name != report.run_name: raise CLIError("Claude final report identifies a different service run") if run.status != RunStatus.RUNNING or run.service is None: @@ -91,11 +116,6 @@ def build_verified_preset( raise CLIError("Claude final report base does not match the requested model") elif report.model != preset_configuration.model.exact_repo: raise CLIError("Claude changed an exact model request") - if ( - preset_configuration.min_context_length is not None - and report.context_length < preset_configuration.min_context_length - ): - raise CLIError("Claude final service does not meet the requested context length") target_type = ( "gateway" if urlparse(run.service.url).scheme in {"http", "https"} else "server-proxy" @@ -111,6 +131,12 @@ def build_verified_preset( for key, value in preset_configuration.env.items(): if isinstance(value, EnvSentinel) and key in portable_service.env: portable_service.env[key] = value + if portable_service.files: + if workspace_path is None or session_path is None: + raise CLIError("Claude final service uses files but no workspace is attached") + _rewrite_workspace_file_paths( + portable_service, workspace_path=workspace_path, session_path=session_path + ) return build_preset( name=name, service=portable_service, @@ -119,6 +145,9 @@ def build_verified_preset( model=report.model, context_length=report.context_length, benchmark=benchmark, + trial=report.trial, + min_context_length=preset_configuration.min_context_length, + max_ttft=preset_configuration.max_ttft, preset_id=preset_id, ) diff --git a/src/dstack/_internal/cli/services/presets/workspace.py b/src/dstack/_internal/cli/services/presets/workspace.py index 54b24abf5..ae1a0bfee 100644 --- a/src/dstack/_internal/cli/services/presets/workspace.py +++ b/src/dstack/_internal/cli/services/presets/workspace.py @@ -10,17 +10,18 @@ from contextlib import suppress from dataclasses import dataclass from pathlib import Path -from typing import Optional +from typing import Optional, Sequence from dstack._internal.cli.services.presets.session import ( _CONSTRAINTS_FILENAME, _FINAL_REPORT_FILENAME, _PROGRESS_FILENAME, _RUNS_FILENAME, - _TRIALS_FILENAME, - _VERIFICATIONS_FILENAME, + _SERVICE_DIRNAME, + _TRIALS_DIRNAME, PresetAgentSession, ) +from dstack._internal.cli.utils.common import warn from dstack._internal.compat import IS_WINDOWS from dstack._internal.core.errors import CLIError @@ -52,12 +53,12 @@ def runs_path(self) -> Path: return self.path / _RUNS_FILENAME @property - def trials_path(self) -> Path: - return self.path / _TRIALS_FILENAME + def trials_dir(self) -> Path: + return self.path / _TRIALS_DIRNAME @property - def verifications_path(self) -> Path: - return self.path / _VERIFICATIONS_FILENAME + def service_dir(self) -> Path: + return self.path / _SERVICE_DIRNAME @property def constraints_path(self) -> Path: @@ -123,9 +124,9 @@ def remove_agent_workspace(session: PresetAgentSession) -> None: def scrub_workspace_token(session: PresetAgentSession) -> None: - """Removes the agent's dstack config (a live token) from a workspace kept - for resume, so an interrupted session leaves no credential on disk. Resume - re-mints it via `build_preset_agent_env`.""" + """The dstack config holds a live token; scrubbing it leaves an interrupted + session with no on-disk credential, and resume re-mints it via + `build_preset_agent_env`.""" workspace = session.read_manifest().get("workspace") if not workspace: return @@ -145,6 +146,8 @@ def _create_workspace_alias(real: Path) -> Path: def _ensure_workspace_alias(alias: Path, real: Path) -> None: + """Recreates the alias symlink idempotently, refusing an existing path unless + it is a symlink to `real` owned by the current user.""" if os.path.lexists(alias): if ( alias.is_symlink() @@ -160,6 +163,9 @@ def _ensure_workspace_alias(alias: Path, real: Path) -> None: def _validate_control_socket_path(build_root: Path) -> None: + """Rejects the workspace if the longest possible run-name SSH control-socket + path would exceed the Unix-socket length limit (`'x' * _MAX_RUN_NAME_LENGTH` + stands in for the worst-case run name).""" if IS_WINDOWS: return path = build_root / "h" / ".dstack" / "ssh" / f"{'x' * _MAX_RUN_NAME_LENGTH}.control.sock" @@ -171,13 +177,10 @@ def _prepare_workspace(workspace: PresetAgentWorkspace) -> None: workspace.path.mkdir(mode=0o700, parents=True, exist_ok=False) workspace.dstack_home.mkdir(mode=0o700) workspace.temp_path.mkdir(mode=0o700) - for path in [ - workspace.progress_path, - workspace.runs_path, - workspace.trials_path, - workspace.verifications_path, - ]: + for path in [workspace.progress_path, workspace.runs_path]: path.touch() + workspace.trials_dir.mkdir() + workspace.service_dir.mkdir() workspace.bin_path.mkdir() _install_python_command(workspace.bin_path, "progress", _get_progress_script()) (workspace.dstack_home / ".ssh").mkdir(mode=0o700) @@ -268,6 +271,47 @@ def _get_progress_script() -> str: """ +def install_previous_records( + workspace: PresetAgentWorkspace, previous_sessions: Sequence[PresetAgentSession] +) -> None: + """Remove-then-recopy, so a crashed partial copy heals on the next run.""" + for session in previous_sessions: + target_root = workspace.path / "previous" / session.preset_id + shutil.rmtree(target_root, ignore_errors=True) + if not _copy_session_records(session.path, target_root): + warn(f"Previous session {session.preset_id} has no records") + + +def _copy_session_records(source_root: Path, target_root: Path) -> bool: + copied = False + for name in (_CONSTRAINTS_FILENAME, _FINAL_REPORT_FILENAME): + if (source_root / name).is_file(): + target_root.mkdir(parents=True, exist_ok=True) + shutil.copyfile(source_root / name, target_root / name) + copied = True + for group, filenames in ( + (_TRIALS_DIRNAME, ("trial.json", "task.dstack.yml")), + (_SERVICE_DIRNAME, ("service.dstack.yml", "verification.json")), + ): + source_group = source_root / group + if not source_group.is_dir(): + continue + for record_dir in sorted(source_group.iterdir()): + if not record_dir.is_dir() or not record_dir.name.isdigit(): + continue + target_dir = target_root / group / record_dir.name + for name in filenames: + if (record_dir / name).is_file(): + target_dir.mkdir(parents=True, exist_ok=True) + shutil.copyfile(record_dir / name, target_dir / name) + copied = True + patches = record_dir / "patches" + if group == _TRIALS_DIRNAME and patches.is_dir(): + shutil.copytree(patches, target_dir / "patches", dirs_exist_ok=True) + copied = True + return copied + + def _install_skills(workspace: Path) -> None: source_dir = _get_skills_dir() target_dir = workspace / ".claude" / "skills" @@ -280,6 +324,9 @@ def _install_skills(workspace: Path) -> None: def _get_skills_dir() -> Path: + """Returns the bundled skills dir, preferring the pip-packaged + `resources/skills` copy and falling back to the repo checkout (`parents[6]` + is the repo root).""" source_path = Path(__file__).resolve() candidates = ( source_path.parent / "resources" / "skills", diff --git a/src/tests/_internal/cli/preset_factories.py b/src/tests/_internal/cli/preset_factories.py index 6d5f6cad1..8a799bc8d 100644 --- a/src/tests/_internal/cli/preset_factories.py +++ b/src/tests/_internal/cli/preset_factories.py @@ -142,6 +142,7 @@ def get_successful_preset_report(run: Run) -> AgentFinalReport: run_id=run.id, run_name=run.run_spec.run_name, service_yaml="type: service", + trial=1, base="Qwen/Qwen3.5-27B", model="community/Qwen3.5-27B-GPTQ-Int4", context_length=32768, diff --git a/src/tests/_internal/cli/services/presets/test_agent.py b/src/tests/_internal/cli/services/presets/test_agent.py index c3e9a44ca..22fa2bfb0 100644 --- a/src/tests/_internal/cli/services/presets/test_agent.py +++ b/src/tests/_internal/cli/services/presets/test_agent.py @@ -333,6 +333,57 @@ async def test_sends_prompt_and_redacts_raw_output(self, tmp_path, monkeypatch, assert "[redacted]" in trace[0]["event"] assert capsys.readouterr().out == "" + @pytest.mark.asyncio + async def test_mirrors_trial_and_service_records_into_the_session(self, tmp_path, monkeypatch): + script = tmp_path / "fake_claude.py" + script.write_text( + """import json +import os +import sys + +sys.stdin.read() +os.makedirs("trials/1") +with open("trials/1/task.dstack.yml", "w") as f: + f.write("env:\\n - TOKEN=secret-token\\n") +with open("trials/1/trial.json", "w") as f: + json.dump({"learned": "uses secret-token"}, f) +os.makedirs("service/1") +with open("service/1/verification.json", "w") as f: + json.dump({"trial": 1, "status": "verified"}, f) +print(json.dumps({"type": "result", "structured_output": {"ok": True}})) +""" + ) + (tmp_path / "progress.jsonl").touch() + (tmp_path / "trials").mkdir() + (tmp_path / "service").mkdir() + monkeypatch.setattr( + "dstack._internal.cli.services.presets.agent._build_claude_command", + lambda **_: [sys.executable, str(script)], + ) + workspace = PresetAgentWorkspace(path=tmp_path, dstack_home=tmp_path / "home") + session_path = tmp_path / "session" + session_path.mkdir() + (session_path / "agent.log").touch() + agent_session = PresetAgentSession(path=session_path, debug=False) + + output = await run_preset_agent( + prompt="p", + env=os.environ.copy(), + workspace=workspace, + auth=_claude_auth(), + redacted_values=("secret-token",), + agent_session=agent_session, + ) + + assert output.report_data == {"ok": True} + trial_dir = session_path / "trials" / "1" + assert (trial_dir / "task.dstack.yml").read_text() == "env:\n - TOKEN=[redacted]\n" + assert json.loads((trial_dir / "trial.json").read_text()) == {"learned": "uses [redacted]"} + assert json.loads((session_path / "service" / "1" / "verification.json").read_text()) == { + "trial": 1, + "status": "verified", + } + @pytest.mark.asyncio async def test_accepts_stream_event_larger_than_64_kib(self, tmp_path, monkeypatch): script = tmp_path / "fake_claude.py" @@ -452,6 +503,94 @@ def test_missing_source_is_no_op(self, tmp_path): assert not (tmp_path / "target.jsonl").exists() +class TestDirectoryMirror: + def _mirror(self, tmp_path, **kwargs): + from dstack._internal.cli.services.presets.tail import _DirectoryMirror + + return _DirectoryMirror( + source=tmp_path / "w" / "trials", + target=tmp_path / "session" / "trials", + **{"redacted_values": ["dstack-secret"], **kwargs}, + ) + + def test_copies_the_tree_redacted(self, tmp_path): + source = tmp_path / "w" / "trials" / "1" + source.mkdir(parents=True) + (source / "task.dstack.yml").write_text("env:\n - TOKEN=dstack-secret\n") + (source / "trial.json").write_text('{"learned": "x"}') + mirror = self._mirror(tmp_path) + + mirror.flush() + + target = tmp_path / "session" / "trials" / "1" + assert (target / "task.dstack.yml").read_text() == "env:\n - TOKEN=[redacted]\n" + assert (target / "trial.json").read_text() == '{"learned": "x"}' + + def test_copies_bytes_verbatim(self, tmp_path): + # Patches must replicate the trial exactly, and text mode rewrites + # newlines and replaces bytes that are not valid UTF-8. + source = tmp_path / "w" / "trials" / "1" / "patches" + source.mkdir(parents=True) + (source / "tuned.csv").write_bytes(b"m,n,k\r\n8,1536,4096\r\n") + (source / "weights.bin").write_bytes(b"\x00\xff\xfe binary") + mirror = self._mirror(tmp_path) + + mirror.flush() + + target = tmp_path / "session" / "trials" / "1" / "patches" + assert (target / "tuned.csv").read_bytes() == b"m,n,k\r\n8,1536,4096\r\n" + assert (target / "weights.bin").read_bytes() == b"\x00\xff\xfe binary" + + def test_a_rewritten_source_converges_instead_of_corrupting(self, tmp_path): + # The failure mode this mirror exists to remove: a source rewritten + # under a byte-offset tailer used to commit a torn record forever. + source = tmp_path / "w" / "trials" / "1" + source.mkdir(parents=True) + (source / "trial.json").write_text('{"learned": "first"}') + mirror = self._mirror(tmp_path) + mirror.flush() + + (source / "trial.json").write_text('{"learned": "rewritten"}') + os.utime(source / "trial.json", ns=(1, 1)) # force a distinct signature + mirror.flush() + + target = tmp_path / "session" / "trials" / "1" / "trial.json" + assert json.loads(target.read_text()) == {"learned": "rewritten"} + + def test_an_unchanged_file_is_not_rewritten(self, tmp_path): + source = tmp_path / "w" / "trials" / "1" + source.mkdir(parents=True) + (source / "trial.json").write_text('{"learned": "x"}') + mirror = self._mirror(tmp_path) + mirror.flush() + target = tmp_path / "session" / "trials" / "1" / "trial.json" + first_stat = target.stat().st_mtime_ns + + mirror.flush() + + assert target.stat().st_mtime_ns == first_stat + + def test_missing_source_is_no_op(self, tmp_path): + mirror = self._mirror(tmp_path) + + mirror.flush() + + assert not (tmp_path / "session").exists() + + def test_skips_files_above_the_size_limit(self, tmp_path, monkeypatch): + from dstack._internal.cli.services.presets.tail import _DirectoryMirror + + monkeypatch.setattr(_DirectoryMirror, "_MAX_FILE_BYTES", 8) + source = tmp_path / "w" / "trials" / "1" + source.mkdir(parents=True) + (source / "trial.json").write_text('{"learned": "far larger than eight bytes"}') + mirror = self._mirror(tmp_path, echo=False) + + mirror.flush() + + assert not (tmp_path / "session" / "trials" / "1" / "trial.json").exists() + + class TestWriteAgentInfo: def test_writes_model_params_and_auth(self, tmp_path, monkeypatch): monkeypatch.setattr( @@ -848,10 +987,17 @@ def test_refusals(self, tmp_path, monkeypatch, manifest, match): load_resumable_agent_session(preset_id) +def _write_trials(tmp_path, records): + trials_dir = tmp_path / "trials" + for number, record in enumerate(records, 1): + (trials_dir / str(number)).mkdir(parents=True) + (trials_dir / str(number) / "trial.json").write_text(json.dumps(record)) + return trials_dir + + class TestSummarizeSessionTrials: def test_counts_records_even_when_trials_share_a_task(self, tmp_path): record = { - "task": {"name": "qwen-ab12cd34-1"}, "resources": {"gpu": {"name": "A40", "memory": "48GB", "count": 1}}, "benchmark": { "workload": {"concurrency": 8}, @@ -860,18 +1006,13 @@ def test_counts_records_even_when_trials_share_a_task(self, tmp_path): } 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]) - ) + trials_dir = _write_trials(tmp_path, [record, rerun, record, nameless]) - summary = _summarize_session_trials(path) + summary = _summarize_session_trials(trials_dir) # 4 records = 4 trials: one long-lived task commonly hosts several - # trials, so shared task names must not collapse the count. + # trials, so identical records must not collapse the count. assert summary["count"] == 4 assert summary["best"] == { "tok_s": 2300.0, @@ -882,41 +1023,91 @@ def test_counts_records_even_when_trials_share_a_task(self, tmp_path): "gpu": "A40:48GB:1", } + def test_trials_are_ordered_numerically_beyond_nine(self, tmp_path): + # Twelve trials: a string sort would chart 1, 10, 11, 12, 2, ... + records = [ + { + "benchmark": { + "metrics": {"total_output_tokens": float(i * 100), "duration_seconds": 1.0}, + "workload": {"concurrency": 1}, + } + } + for i in range(1, 13) + ] + trials_dir = _write_trials(tmp_path, records) + + summary = _summarize_session_trials(trials_dir) + + assert summary["series"] == [float(i * 100) for i in range(1, 13)] + + def test_a_trial_still_in_flight_is_not_counted(self, tmp_path): + record = { + "resources": {}, + "benchmark": { + "workload": {"concurrency": 8}, + "metrics": {"duration_seconds": 10.0, "total_output_tokens": 20000}, + }, + } + trials_dir = _write_trials(tmp_path, [record]) + # Trial 2 has started but has no result yet. + (trials_dir / "2").mkdir() + (trials_dir / "2" / "task.dstack.yml").write_text("type: task\n") + + summary = _summarize_session_trials(trials_dir) + + assert summary["count"] == 1 + assert summary["series"] == [2000.0] + + def test_a_torn_result_copy_reads_as_still_in_flight(self, tmp_path): + trials_dir = _write_trials(tmp_path, []) + (trials_dir / "1").mkdir(parents=True) + (trials_dir / "1" / "trial.json").write_text('{"benchmark": {"met') + + summary = _summarize_session_trials(trials_dir) + + assert summary["count"] == 0 + class TestReadLastSessionVerification: - def test_last_record_wins_and_a_missing_file_is_not_verifying(self, tmp_path): - path = tmp_path / "verifications.jsonl" - - assert _read_last_session_verification(path) is None - - path.write_text( - "\n".join( - json.dumps(entry) - for entry in [ - {"trial": 3, "run_name": "p-2", "status": "verifying"}, - {"trial": 3, "run_name": "p-2", "status": "failed", "reason": "probe"}, - {"trial": 2, "run_name": "p-3", "status": "verifying"}, - ] - ) - + "\n" + def test_newest_attempt_wins_and_a_missing_directory_is_not_verifying(self, tmp_path): + service_dir = tmp_path / "service" + + assert _read_last_session_verification(service_dir) is None + + (service_dir / "1").mkdir(parents=True) + (service_dir / "1" / "service.dstack.yml").write_text("type: service\n") + (service_dir / "1" / "verification.json").write_text( + json.dumps({"trial": 3, "run_name": "p-2", "status": "failed", "reason": "probe"}) ) + (service_dir / "2").mkdir() + (service_dir / "2" / "service.dstack.yml").write_text("type: service\n") + + # Attempt 2 has no result yet, so it is the one in flight. + assert _read_last_session_verification(service_dir) == {"status": "verifying"} + + def test_the_newest_result_is_returned(self, tmp_path): + # Attempts 9 and 10: a string sort would pick "9" as the newest. + service_dir = tmp_path / "service" + for number, status in ((1, "failed"), (9, "failed"), (10, "verified")): + (service_dir / str(number)).mkdir(parents=True) + (service_dir / str(number) / "verification.json").write_text( + json.dumps({"trial": number, "run_name": f"p-{number}", "status": status}) + ) - assert _read_last_session_verification(path) == { - "trial": 2, - "run_name": "p-3", - "status": "verifying", + assert _read_last_session_verification(service_dir) == { + "trial": 10, + "run_name": "p-10", + "status": "verified", } - def test_skips_partial_trailing_lines(self, tmp_path): - # The mirror appends as the agent writes, so the file can be read - # mid-line. - path = tmp_path / "verifications.jsonl" - path.write_text( - json.dumps({"trial": 1, "run_name": "p-2", "status": "verifying"}) - + '\n{"trial": 1, "run_na' - ) + def test_a_torn_result_copy_reads_as_verifying(self, tmp_path): + # The mirror copies whole files, but a reader can still catch one + # mid-replace; the next pass converges. + service_dir = tmp_path / "service" + (service_dir / "1").mkdir(parents=True) + (service_dir / "1" / "verification.json").write_text('{"trial": 1, "run_na') - assert _read_last_session_verification(path)["status"] == "verifying" + assert _read_last_session_verification(service_dir) == {"status": "verifying"} class TestFileLineReader: diff --git a/src/tests/_internal/cli/services/presets/test_apply.py b/src/tests/_internal/cli/services/presets/test_apply.py index 0597d3eab..d6c628042 100644 --- a/src/tests/_internal/cli/services/presets/test_apply.py +++ b/src/tests/_internal/cli/services/presets/test_apply.py @@ -27,7 +27,9 @@ def test_accepts_matching_base_model_and_context(self): _validate_preset_matches(preset, configuration=configuration) - def test_rejects_insufficient_context(self): + def test_warns_on_insufficient_context_instead_of_failing(self, capsys): + # The preset is chosen by ID and may be the best a session could verify; + # the shortfall is stated and the plan confirmation decides. preset = get_preset(preset_id="small", context_length=4096) configuration = PresetConfiguration( name="qwen", @@ -35,8 +37,11 @@ def test_rejects_insufficient_context(self): min_context_length=8192, ) - with pytest.raises(CLIError, match="context length"): - _validate_preset_matches(preset, configuration=configuration) + _validate_preset_matches(preset, configuration=configuration) + + output = capsys.readouterr().out + assert "verified for context length 4096" in output + assert "8192" in output def test_exact_request_matches_repo_and_client_facing_name(self): matching = get_preset(preset_id="matching") diff --git a/src/tests/_internal/cli/services/presets/test_create.py b/src/tests/_internal/cli/services/presets/test_create.py index 3297f21d9..d4b2b9cbc 100644 --- a/src/tests/_internal/cli/services/presets/test_create.py +++ b/src/tests/_internal/cli/services/presets/test_create.py @@ -26,6 +26,7 @@ create_preset, follow_preset, reconcile_detached_sessions, + resolve_previous_sessions, stop_preset_session, ) from dstack._internal.cli.services.presets.session import ( @@ -347,6 +348,157 @@ async def run_agent(**kwargs): assert creation_context.run_apis.stopped_names == stopped_names +class TestResolvePreviousSessions: + def _store(self, tmp_path, monkeypatch, *ids): + store = tmp_path / "presets-store" + for preset_id in ids: + root = store / preset_id + (root / "trials" / "1").mkdir(parents=True) + (root / "trials" / "1" / "trial.json").write_text("{}") + (root / "session.json").write_text(json.dumps({"status": "failed"})) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.session.get_presets_dir", + lambda: store, + ) + return store + + def test_resolves_and_dedupes_in_order(self, tmp_path, monkeypatch): + self._store(tmp_path, monkeypatch, "a1b2c3d4", "e5f6a7b8") + + sessions = resolve_previous_sessions(["e5f6a7b8", "a1b2c3d4", "e5f6a7b8"]) + + assert [session.preset_id for session in sessions] == ["e5f6a7b8", "a1b2c3d4"] + + def test_rejects_an_unknown_reference(self, tmp_path, monkeypatch): + self._store(tmp_path, monkeypatch, "a1b2c3d4") + + with pytest.raises(CLIError, match="'nope' does not exist"): + resolve_previous_sessions(["a1b2c3d4", "nope"]) + + def test_warns_when_a_chained_session_is_not_included(self, tmp_path, monkeypatch, capsys): + store = self._store(tmp_path, monkeypatch, "a1b2c3d4", "e5f6a7b8") + (store / "e5f6a7b8" / "session.json").write_text( + json.dumps({"status": "failed", "previous": ["a1b2c3d4", "00000000"]}) + ) + + resolve_previous_sessions(["e5f6a7b8", "a1b2c3d4"]) + + output = capsys.readouterr().out + assert "e5f6a7b8 was created with --previous 00000000" in output + # The included parent must not be warned about. + assert output.count("was created with") == 1 + + def test_rejects_a_previous_session_that_is_still_running(self, tmp_path, monkeypatch): + store = self._store(tmp_path, monkeypatch, "a1b2c3d4") + (store / "a1b2c3d4" / "session.json").write_text(json.dumps({"status": "running"})) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.session_process_alive", + lambda manifest: True, + ) + + with pytest.raises(CLIError, match="still running"): + resolve_previous_sessions(["a1b2c3d4"]) + + def test_accepts_a_stale_running_session_whose_process_died(self, tmp_path, monkeypatch): + store = self._store(tmp_path, monkeypatch, "a1b2c3d4") + (store / "a1b2c3d4" / "session.json").write_text(json.dumps({"status": "running"})) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.session_process_alive", + lambda manifest: False, + ) + + sessions = resolve_previous_sessions(["a1b2c3d4"]) + + assert [session.preset_id for session in sessions] == ["a1b2c3d4"] + + +class TestEffectivePrevious: + def _args(self, previous): + # The real parser builds the namespace, so profile attributes stay in + # sync with `register_profile_args` instead of being hand-listed. + import argparse + + from dstack._internal.cli.services.profile import register_profile_args + + parser = argparse.ArgumentParser() + register_profile_args(parser) + args = parser.parse_args([]) + args.name = None + args.trials = None + args.previous = previous + args.no_profile = True + return args + + def test_flag_overrides_and_property_stands_without_it(self): + from dstack._internal.cli.commands.preset import _get_effective_configuration + + def configuration(): + # A fresh object per call: the merger mutates its input. + return PresetConfiguration( + name="qwen", model={"base": "Qwen/Qwen3.5-27B"}, previous=["from-config"] + ) + + overridden = _get_effective_configuration( + configuration(), self._args(["from-flag"]), require_name=False + ) + kept = _get_effective_configuration(configuration(), self._args(None), require_name=False) + + assert overridden.previous == ["from-flag"] + assert kept.previous == ["from-config"] + + +class TestCreateWithPrevious: + @pytest.mark.asyncio + async def test_installs_records_pins_manifest_and_extends_the_prompt( + self, creation_context, monkeypatch, tmp_path + ): + store = tmp_path / "presets-store" + root = store / "8d3b01aa" + (root / "trials" / "1").mkdir(parents=True) + (root / "trials" / "1" / "trial.json").write_text('{"learned": "x"}') + (root / "session.json").write_text(json.dumps({"status": "failed"})) + monkeypatch.setattr( + "dstack._internal.cli.services.presets.session.get_presets_dir", + lambda: store, + ) + session_path = tmp_path / "fresh" + session_path.mkdir() + agent_session = PresetAgentSession(path=session_path, debug=False) + seen = {} + + async def run_agent(**kwargs): + seen["prompt"] = kwargs["prompt"] + seen["record"] = ( + kwargs["workspace"].path / "previous" / "8d3b01aa" / "trials" / "1" / "trial.json" + ).is_file() + return PresetAgentProcessOutput( + report_data=json.loads( + get_successful_preset_report(creation_context.run).model_dump_json() + ) + ) + + monkeypatch.setattr( + "dstack._internal.cli.services.presets.create.run_preset_agent", + run_agent, + ) + await _create_preset( + api=creation_context.api, + configuration=creation_context.configuration, + source_configuration=creation_context.source_configuration, + store=creation_context.store, + build_name="qwen-build", + agent_session=agent_session, + previous=resolve_previous_sessions(["8d3b01aa"]), + ) + + assert seen["record"] is True + assert "## Previous Sessions" in seen["prompt"] + assert "8d3b01aa" in seen["prompt"] + assert agent_session.read_manifest()["previous"] == ["8d3b01aa"] + # constraints.json is a session record even without --debug. + assert (session_path / "constraints.json").is_file() + + class TestBuildName: def test_derives_slug_for_nameless_and_keeps_prefix_bounded(self): assert _get_build_name(None, "Qwen/Qwen3.5-27B", "a1b2c3d4") == "qwen3-5-27b-a1b2c3d4" diff --git a/src/tests/_internal/cli/services/presets/test_output.py b/src/tests/_internal/cli/services/presets/test_output.py index 93976e690..6d7079913 100644 --- a/src/tests/_internal/cli/services/presets/test_output.py +++ b/src/tests/_internal/cli/services/presets/test_output.py @@ -62,13 +62,17 @@ def test_a_preset_saved_before_the_field_existed_still_loads(self): class TestPrintPresets: - def test_preserves_benchmark_concurrency_at_narrow_width(self, monkeypatch): + def test_preserves_constraints_and_benchmark_at_narrow_width(self, monkeypatch): output = StringIO() monkeypatch.setattr(output_module, "console", plain_console(output, width=79)) output_module.print_presets([get_preset()]) - assert "conc=1" in "".join(output.getvalue().split()) + # Both columns wrap rather than clip, so their full content survives even + # when a long model name would otherwise squeeze them out. + joined = "".join(output.getvalue().split()) + assert "conc=1" in joined + assert "ttft=108ms" in joined def test_prints_submitted_column(self, monkeypatch): output = StringIO() @@ -110,7 +114,8 @@ def test_shows_progress_after_status_and_best_benchmark(self): } ) - assert row["STATUS"] == "[bold sea_green3]trialing[/] [secondary](2/3)[/]" + # 2 completed, so the 3rd is the one being trialed. + assert row["STATUS"] == "[bold sea_green3]trialing[/] [secondary](3/3)[/]" # Per-user speed is 1/TPOT, the same definition the preset row uses — not # the aggregate over concurrency, which would read 292 here. assert row["BENCHMARK"].startswith("tok/s/user=292") @@ -155,7 +160,20 @@ def test_shows_zero_progress_without_benchmark(self): {"id": "ab12cd34", "status": "running", "trials_num": 3, "trials": {"count": 0}} ) - assert row["STATUS"] == "[bold sea_green3]trialing[/] [secondary](0/3)[/]" + assert row["STATUS"] == "[bold sea_green3]trialing[/] [secondary](1/3)[/]" + + def test_verifying_keeps_the_completed_count(self): + row = _session_row( + { + "id": "ab12cd34", + "status": "running", + "trials_num": 3, + "trials": {"count": 3}, + "verification": {"status": "verifying"}, + } + ) + + assert row["STATUS"] == "[bold deep_sky_blue1]verifying[/] [secondary](3/3)[/]" def test_omits_progress_without_trials_data(self): row = _session_row({"id": "ab12cd34", "status": "interrupted"}) @@ -269,29 +287,35 @@ def test_a_spent_trial_budget_alone_stays_trialing(self): assert row["STATUS"].startswith("[bold sea_green3]trialing[/]") +def _write_trials(tmp_path, records): + trials_dir = tmp_path / "trials" + for number, record in enumerate(records, 1): + (trials_dir / str(number)).mkdir(parents=True) + (trials_dir / str(number) / "trial.json").write_text(json.dumps(record)) + return trials_dir + + class TestFailedTrials: def test_a_failed_trial_keeps_its_benchmark_but_never_becomes_best(self, tmp_path): from dstack._internal.cli.services.presets.session import _summarize_session_trials # The failed trial is the fastest. It broke `max_ttft`, so promoting it # would put a non-compliant configuration at the top of the listing. - path = tmp_path / "trials.jsonl" - path.write_text( - "\n".join( - json.dumps( - { - "benchmark": { - "metrics": {"total_output_tokens": tokens, "duration_seconds": 1.0}, - "workload": {"concurrency": 8}, - }, - **({"failed": True} if failed else {}), - } - ) + trials_dir = _write_trials( + tmp_path, + [ + { + "benchmark": { + "metrics": {"total_output_tokens": tokens, "duration_seconds": 1.0}, + "workload": {"concurrency": 8}, + }, + **({"failed": True} if failed else {}), + } for tokens, failed in ((100.0, False), (900.0, True), (300.0, False)) - ) + ], ) - summary = _summarize_session_trials(path) + summary = _summarize_session_trials(trials_dir) assert summary["count"] == 3 # Still charted: the trial happened and its number is real. @@ -304,12 +328,12 @@ def test_reports_the_gpu_when_no_trial_produced_a_benchmark(self, tmp_path): # Neither `best` nor `best_failed` can carry the hardware here, so this is # the only thing left that knows what the run was on. - path = tmp_path / "trials.jsonl" - path.write_text( - json.dumps({"resources": {"gpu": {"name": "MI300X", "memory": "192GB", "count": 1}}}) + trials_dir = _write_trials( + tmp_path, + [{"resources": {"gpu": {"name": "MI300X", "memory": "192GB", "count": 1}}}], ) - summary = _summarize_session_trials(path) + summary = _summarize_session_trials(trials_dir) assert summary["best"] is None assert summary["best_failed"] is None @@ -318,28 +342,26 @@ def test_reports_the_gpu_when_no_trial_produced_a_benchmark(self, tmp_path): def test_the_fastest_failed_trial_is_kept_when_nothing_passed(self, tmp_path): from dstack._internal.cli.services.presets.session import _summarize_session_trials - path = tmp_path / "trials.jsonl" - path.write_text( - "\n".join( - json.dumps( - { - "benchmark": { - "metrics": { - "total_output_tokens": tokens, - "duration_seconds": 1.0, - "tpot_ms": {"p50": 34.4}, - "ttft_ms": {"p50": 4300.0}, - }, - "workload": {"concurrency": 4}, + trials_dir = _write_trials( + tmp_path, + [ + { + "benchmark": { + "metrics": { + "total_output_tokens": tokens, + "duration_seconds": 1.0, + "tpot_ms": {"p50": 34.4}, + "ttft_ms": {"p50": 4300.0}, }, - "failed": True, - } - ) + "workload": {"concurrency": 4}, + }, + "failed": True, + } for tokens in (100.0, 300.0, 200.0) - ) + ], ) - summary = _summarize_session_trials(path) + summary = _summarize_session_trials(trials_dir) assert summary["best"] is None assert summary["best_failed"]["tok_s"] == 300.0 @@ -375,9 +397,9 @@ def test_a_run_that_met_nothing_still_shows_what_it_measured(self): class TestFailedTrialSpark: - def test_a_trial_that_broke_a_constraint_is_yellow_but_still_charted(self): - # Its number is real, so it earns a bar; the breach makes it yellow, not - # red — red is reserved for a trial that produced nothing. + def test_a_trial_that_broke_a_constraint_is_charted_but_not_the_best(self): + # Its number is real, so it earns a bar rather than a `·`, and green goes + # to the best trial that meets the constraints even on a lower number. session = { "id": "ab12cd34", "status": "running", @@ -390,8 +412,26 @@ def test_a_trial_that_broke_a_constraint_is_yellow_but_still_charted(self): spark = output_module._format_trial_spark(session) - assert spark.count("gold1") == 1 - assert "indian_red1" not in spark + assert spark.count("indian_red1") == 1 + assert spark.count("sea_green3") == 1 + assert "gold1" not in spark assert "·" not in spark - # The failed trial is the highest number and must not be styled as best. + + def test_the_best_failed_trial_is_gold_while_none_passes(self): + # With nothing meeting the constraints, the best result so far is still + # what the run has to show; the rest are context. + session = { + "id": "ab12cd34", + "status": "running", + "trials": { + "count": 3, + "series": [100.0, 900.0, 300.0], + "failed": [True, True, True], + }, + } + + spark = output_module._format_trial_spark(session) + + assert spark.count("gold1") == 1 + assert spark.count("indian_red1") == 2 assert "sea_green3" not in spark diff --git a/src/tests/_internal/cli/services/presets/test_prompt.py b/src/tests/_internal/cli/services/presets/test_prompt.py index fc1d8cf2c..0e81afac9 100644 --- a/src/tests/_internal/cli/services/presets/test_prompt.py +++ b/src/tests/_internal/cli/services/presets/test_prompt.py @@ -42,20 +42,110 @@ def test_fails_loudly_when_the_prompt_has_no_directives(self, tmp_path, monkeypa def test_drops_maintainer_notes_but_keeps_every_other_comment(self, tmp_path, monkeypatch): noted = tmp_path / "system_prompt.md" noted.write_text( - "Kept.\n\n" - "Plain and stay.\n" + "Kept.\n\nPlain stays.\n" ) monkeypatch.setattr(prompt_module, "_SYSTEM_PROMPT_PATH", noted) text = get_preset_agent_system_prompt() assert "TODO" not in text - assert text == "Kept.\nPlain and stay." + assert text == "Kept.\nPlain stays." - def test_rejects_unknown_directive_variables(self, tmp_path, monkeypatch): + def test_rejects_unknown_variables_even_in_a_dropped_branch(self, tmp_path, monkeypatch): broken = tmp_path / "system_prompt.md" - broken.write_text("Text more.\n") + broken.write_text("\n\nX\n\n\n") monkeypatch.setattr(prompt_module, "_SYSTEM_PROMPT_PATH", broken) + # `previous` is unset, so the branch would be dropped; the typo inside + # it must not hide behind that. with pytest.raises(CLIError, match="Unknown variable"): get_preset_agent_system_prompt() + + def test_rejects_malformed_directives(self, tmp_path, monkeypatch): + broken = tmp_path / "system_prompt.md" + broken.write_text("Text more.\n") + monkeypatch.setattr(prompt_module, "_SYSTEM_PROMPT_PATH", broken) + + with pytest.raises(CLIError, match="Invalid directive"): + get_preset_agent_system_prompt() + + broken.write_text("An opener that never closes its comment: BCD\n") + monkeypatch.setattr(prompt_module, "_SYSTEM_PROMPT_PATH", doc) + + assert get_preset_agent_system_prompt(baseline=True) == "ABD" + assert get_preset_agent_system_prompt() == "ACD" + + def test_dedents_only_an_exactly_indented_body(self, tmp_path, monkeypatch): + doc = tmp_path / "system_prompt.md" + monkeypatch.setattr(prompt_module, "_SYSTEM_PROMPT_PATH", doc) + # The reference examples: `baseline` plays , `previous` plays . + cases = [ + ("asaga\n", "asaga"), + ("\n asaga\n\n", "asaga"), + ( + "\n asaga\n\n c\n\n", + "asaga\n\nc", + ), + ( + "\n asaga\n\n" + " \n c\n \n\n", + "asaga\n\nc", + ), + # Not exact: the body keeps its indentation as written. + ( + "\nasaga\n\n c\n\n", + "asaga\n\n c", + ), + # An inline-opened block is never dedented. + ( + "asaga\n\n c\n\n", + "asaga\n\n c", + ), + ] + for content, expected in cases: + doc.write_text(content) + previous = "x" if "previous" in content else None + rendered = get_preset_agent_system_prompt(baseline=True, previous=previous) + assert rendered.strip("\n") == expected, content + + def test_nested_blocks_render_one_branch(self, tmp_path, monkeypatch): + nested = tmp_path / "system_prompt.md" + nested.write_text( + "\n" + " IDS={previous}\n" + "\n" + " \n" + " SEEDED\n" + " \n" + "\n" + " \n" + " SOLO\n" + " \n" + "\n" + ) + monkeypatch.setattr(prompt_module, "_SYSTEM_PROMPT_PATH", nested) + + assert get_preset_agent_system_prompt().strip() == "" + assert get_preset_agent_system_prompt(baseline=True).strip() == "SOLO" + assert get_preset_agent_system_prompt(previous="a1, b2").strip() == "IDS=a1, b2" + both = get_preset_agent_system_prompt(baseline=True, previous="a1") + assert both.strip() == "IDS=a1\n\nSEEDED" + + def test_unbalanced_blocks_fail_loudly(self, tmp_path, monkeypatch): + for content, error in [ + ("\nnever closed\n", "Unclosed"), + ("text\n\n", "`end` without"), + ("text\n\n", "`else` without"), + ("\n\n\n\n", "`else` without"), + ]: + broken = tmp_path / "system_prompt.md" + broken.write_text(content) + monkeypatch.setattr(prompt_module, "_SYSTEM_PROMPT_PATH", broken) + with pytest.raises(CLIError, match=error): + get_preset_agent_system_prompt() diff --git a/src/tests/_internal/cli/services/presets/test_verify.py b/src/tests/_internal/cli/services/presets/test_verify.py index b3b27e954..335394f9b 100644 --- a/src/tests/_internal/cli/services/presets/test_verify.py +++ b/src/tests/_internal/cli/services/presets/test_verify.py @@ -18,6 +18,7 @@ ) from dstack._internal.core.errors import CLIError from dstack._internal.core.models.envs import EnvSentinel +from dstack._internal.core.models.files import FilePathMapping from dstack._internal.core.models.profiles import ProfileParams from tests._internal.cli.preset_factories import ( get_running_service_run, @@ -71,6 +72,74 @@ def test_builds_portable_self_contained_preset(self): assert validation.benchmark.target.type == "server-proxy" assert validation.benchmark.client.type == "local" + def test_rewrites_file_paths_onto_the_mirrored_session_copies(self, tmp_path): + # `files` local paths resolve into the agent workspace at submission, and + # the workspace is deleted when the session ends; the preset must point at + # the session's mirrored copies or it cannot be applied later. + workspace = tmp_path / "session" / "workspace" / "w" + (workspace / "service" / "1" / "patches").mkdir(parents=True) + (workspace / "service" / "1" / "patches" / "moe.py.patch").write_text("--- a\n+++ b\n") + session = tmp_path / "session" + (session / "service" / "1" / "patches").mkdir(parents=True) + (session / "service" / "1" / "patches" / "moe.py.patch").write_text("--- a\n+++ b\n") + run = get_running_service_run() + run.run_spec.configuration.files = [ + FilePathMapping( + local_path=str(workspace / "service" / "1" / "patches"), path="/patches" + ) + ] + + preset = build_verified_preset( + run=run, + preset_configuration=PresetConfiguration( + name="qwen-build", model={"base": "Qwen/Qwen3.5-27B"} + ), + report=get_successful_preset_report(run), + workspace_path=workspace, + session_path=session, + ) + + assert preset.service.files[0].local_path == str(session / "service" / "1" / "patches") + # The run spec itself is untouched: only the preset copy is re-rooted. + assert run.run_spec.configuration.files[0].local_path == str( + workspace / "service" / "1" / "patches" + ) + + def test_rejects_a_file_without_a_mirrored_copy(self, tmp_path): + workspace = tmp_path / "session" / "workspace" / "w" + (workspace / "patches").mkdir(parents=True) # workspace root: not mirrored + session = tmp_path / "session" + run = get_running_service_run() + run.run_spec.configuration.files = [ + FilePathMapping(local_path=str(workspace / "patches"), path="/patches") + ] + + with pytest.raises(CLIError, match="no mirrored copy"): + build_verified_preset( + run=run, + preset_configuration=PresetConfiguration( + name="qwen-build", model={"base": "Qwen/Qwen3.5-27B"} + ), + report=get_successful_preset_report(run), + workspace_path=workspace, + session_path=session, + ) + + def test_rejects_files_when_no_workspace_is_attached(self, tmp_path): + run = get_running_service_run() + run.run_spec.configuration.files = [ + FilePathMapping(local_path=str(tmp_path / "patches"), path="/patches") + ] + + with pytest.raises(CLIError, match="no workspace is attached"): + build_verified_preset( + run=run, + preset_configuration=PresetConfiguration( + name="qwen-build", model={"base": "Qwen/Qwen3.5-27B"} + ), + report=get_successful_preset_report(run), + ) + def test_rejects_variant_for_exact_model_request(self): run = get_running_service_run() report = get_successful_preset_report(run).model_copy(update={"model": "other/model"}) diff --git a/src/tests/_internal/cli/services/presets/test_workspace.py b/src/tests/_internal/cli/services/presets/test_workspace.py new file mode 100644 index 000000000..06bd521e6 --- /dev/null +++ b/src/tests/_internal/cli/services/presets/test_workspace.py @@ -0,0 +1,83 @@ +import pytest + +from dstack._internal.cli.services.presets.session import PresetAgentSession +from dstack._internal.cli.services.presets.workspace import ( + PresetAgentWorkspace, + install_previous_records, +) + +pytestmark = pytest.mark.windows + + +def _previous_session(tmp_path, preset_id="8d3b01aa"): + root = tmp_path / "store" / preset_id + (root / "trials" / "1" / "patches").mkdir(parents=True) + (root / "trials" / "1" / "trial.json").write_text('{"learned": "x"}') + (root / "trials" / "1" / "task.dstack.yml").write_text("type: task\n") + (root / "trials" / "1" / "patches" / "moe.py.patch").write_text("--- a\n+++ b\n") + (root / "service" / "1").mkdir(parents=True) + (root / "service" / "1" / "service.dstack.yml").write_text("type: service\n") + (root / "service" / "1" / "verification.json").write_text('{"status": "verified"}') + (root / "constraints.json").write_text("{}") + (root / "final_report.json").write_text('{"success": true}') + # Everything below must stay out of the copy. + (root / "session.json").write_text("{}") + (root / "agent.log").write_text("log") + (root / "trace.jsonl").write_text("{}") + (root / "runs.jsonl").write_text("{}") + (root / "trials" / "not-a-trial").mkdir() + (root / "trials" / "not-a-trial" / "trial.json").write_text("{}") + return PresetAgentSession(path=root, debug=False, preset_id=preset_id) + + +def _workspace(tmp_path): + path = tmp_path / "workspace" / "w" + path.mkdir(parents=True) + return PresetAgentWorkspace(path=path, dstack_home=tmp_path / "workspace" / "h") + + +class TestInstallPreviousRecords: + def test_copies_exactly_the_record_subset(self, tmp_path): + session = _previous_session(tmp_path) + workspace = _workspace(tmp_path) + + install_previous_records(workspace, [session]) + + target = workspace.path / "previous" / "8d3b01aa" + copied = sorted( + file.relative_to(target).as_posix() for file in target.rglob("*") if file.is_file() + ) + assert copied == [ + "constraints.json", + "final_report.json", + "service/1/service.dstack.yml", + "service/1/verification.json", + "trials/1/patches/moe.py.patch", + "trials/1/task.dstack.yml", + "trials/1/trial.json", + ] + + def test_recopy_removes_stale_files(self, tmp_path): + session = _previous_session(tmp_path) + workspace = _workspace(tmp_path) + install_previous_records(workspace, [session]) + stale = workspace.path / "previous" / "8d3b01aa" / "trials" / "9" / "trial.json" + stale.parent.mkdir(parents=True) + stale.write_text("{}") + + install_previous_records(workspace, [session]) + + assert not stale.exists() + assert (workspace.path / "previous" / "8d3b01aa" / "trials" / "1" / "trial.json").exists() + + def test_a_session_without_records_warns(self, tmp_path, capsys): + root = tmp_path / "store" / "empty000" + root.mkdir(parents=True) + (root / "session.json").write_text("{}") + session = PresetAgentSession(path=root, debug=False, preset_id="empty000") + workspace = _workspace(tmp_path) + + install_previous_records(workspace, [session]) + + assert "empty000 has no records" in capsys.readouterr().out + assert not (workspace.path / "previous" / "empty000").exists()