Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 42 additions & 9 deletions src/dstack/_internal/cli/commands/preset.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
import argparse
import io
import os
import sys
import time
from contextlib import suppress
from contextlib import redirect_stderr, suppress
from pathlib import Path

from argcomplete import FilesCompleter # type: ignore[attr-defined]
Expand All @@ -14,6 +16,7 @@
PresetListOutput,
)
from dstack._internal.cli.services.completion import ProjectNameCompleter
from dstack._internal.cli.services.configurators import APPLY_STDIN_NAME
from dstack._internal.cli.services.presets.apply import apply_preset
from dstack._internal.cli.services.presets.create import (
CreationStopped,
Expand All @@ -28,13 +31,14 @@
)
from dstack._internal.cli.services.presets.output import get_presets_table, print_presets
from dstack._internal.cli.services.presets.session import (
list_agent_sessions,
load_resumable_agent_session,
list_preset_sessions,
load_resumable_session,
resolve_session_ref,
)
from dstack._internal.cli.services.presets.store import (
PresetStore,
load_preset_configuration,
parse_preset_configuration,
resolve_preset_prompt,
)
from dstack._internal.cli.services.profile import (
Expand Down Expand Up @@ -251,9 +255,13 @@ def _list(self, args: argparse.Namespace) -> None:
limit=args.limit,
)
return
# The store warns about unreadable presets on stderr once per read;
# inside Live that would tear the render on every refresh. The first
# read happens before Live starts so warnings print once, above the
# table; refreshes read with stderr suppressed.
presets, sessions = self._list_presets_and_sessions(base=base, repo=repo)
with Live(console=console, refresh_per_second=LIVE_TABLE_REFRESH_RATE_PER_SEC) as live:
while True:
presets, sessions = self._list_presets_and_sessions(base=base, repo=repo)
live.update(
get_presets_table(
presets,
Expand All @@ -264,13 +272,15 @@ def _list(self, args: argparse.Namespace) -> None:
)
)
time.sleep(LIVE_TABLE_PROVISION_INTERVAL_SECS)
with redirect_stderr(io.StringIO()):
presets, sessions = self._list_presets_and_sessions(base=base, repo=repo)

def _list_presets_and_sessions(
self, *, base: str | None, repo: str | None
) -> tuple[list[Preset], list[dict]]:
self._reconcile()
presets = PresetStore().list()
sessions = list_agent_sessions()
sessions = list_preset_sessions()
if base or repo:
repo_to_base = {preset.model: preset.base for preset in presets}
presets = _filter_presets(presets, base=base, repo=repo)
Expand All @@ -282,13 +292,14 @@ def _list_presets_and_sessions(
return presets, sessions

def _create(self, args: argparse.Namespace) -> None:
configuration_path, configuration = load_preset_configuration(args.configuration_file)
_check_stdin_configuration_confirmable(args)
_, configuration = _read_configuration_arg(args.configuration_file)
configuration = _get_effective_configuration(configuration, args, require_name=False)
user_prompt = resolve_preset_prompt(configuration, configuration_path)
user_prompt = resolve_preset_prompt(configuration, _prompt_base(args.configuration_file))
store = PresetStore()
resume_session = None
if getattr(args, "resume", None):
resume_session = load_resumable_agent_session(args.resume)
resume_session = load_resumable_session(args.resume)
if getattr(args, "trials", None) is not None:
console.print(
"[warning]--trials is ignored when resuming: "
Expand Down Expand Up @@ -368,7 +379,7 @@ def _get(self, args: argparse.Namespace) -> None:

def _apply(self, args: argparse.Namespace) -> None:
self._reconcile()
configuration_path, configuration = load_preset_configuration(args.configuration_file)
configuration_path, configuration = _read_configuration_arg(args.configuration_file)
configuration = _get_effective_configuration(configuration, args)
apply_preset(
api=Client.from_config(project_name=args.project),
Expand Down Expand Up @@ -532,6 +543,28 @@ def _confirm_preset_creation(store: PresetStore, name: str | None, *, assume_yes
return True


def _read_configuration_arg(configuration_file: str) -> tuple[str, PresetConfiguration]:
"""`-f <path>`, or `-f -` for stdin — the same convention as `dstack apply`."""
if configuration_file == APPLY_STDIN_NAME:
return APPLY_STDIN_NAME, parse_preset_configuration(sys.stdin)
path = Path(configuration_file)
return str(path.resolve()), load_preset_configuration(path)


def _prompt_base(configuration_file: str) -> Path:
"""Prompt files resolve relative to the configuration file; cwd for stdin."""
if configuration_file == APPLY_STDIN_NAME:
return Path.cwd()
return Path(configuration_file).resolve().parent


def _check_stdin_configuration_confirmable(args: argparse.Namespace) -> None:
# Same rule as `dstack apply`: the confirmation prompt cannot read from a
# stdin that is the configuration itself.
if not args.yes and args.configuration_file == APPLY_STDIN_NAME:
raise CLIError("Cannot read configuration from stdin if -y/--yes is not specified")


def _get_effective_configuration(
configuration: PresetConfiguration,
args: argparse.Namespace,
Expand Down
38 changes: 25 additions & 13 deletions src/dstack/_internal/cli/models/configurations.py
Original file line number Diff line number Diff line change
Expand Up @@ -135,14 +135,15 @@ class PresetConfiguration(
),
] = None
min_context_length: Annotated[
Optional[PositiveInt], Field(description="The minimum required context length")
Optional[PositiveInt],
Field(description="The minimum required context length. Required for creation"),
] = None
max_ttft: Annotated[
Optional[PositiveInt],
Field(
description=(
"The maximum p50 time to first token, in milliseconds, that any benchmark"
" may report"
" may report. Required for creation"
)
),
] = None
Expand All @@ -151,7 +152,7 @@ class PresetConfiguration(
Field(
description=(
"The number of benchmarked trials during preset creation"
" before the best one is promoted"
" before the best one is promoted. Required for creation"
)
),
] = None
Expand All @@ -168,7 +169,8 @@ class PresetConfiguration(
Optional[PositiveInt],
Field(
description=(
"The number of simultaneous requests used for benchmarks during preset creation"
"The number of simultaneous requests used for benchmarks during preset"
" creation. Required for creation"
)
),
] = None
Expand All @@ -182,7 +184,7 @@ class PresetConfiguration(
),
] = None
output_tokens: Annotated[
Optional[PositiveInt],
Optional[Annotated[int, Field(ge=2)]],
Field(
description=(
"The number of output tokens per request used for benchmarks during"
Expand All @@ -191,7 +193,7 @@ class PresetConfiguration(
),
] = None
shared_prefix_tokens: Annotated[
Optional[PositiveInt],
Optional[Annotated[int, Field(ge=0)]],
Field(
description=(
"How many of `input_tokens` are a prefix identical in every benchmark request,"
Expand Down Expand Up @@ -338,13 +340,23 @@ class PresetConstraints(CoreModel):
max_ttft: PositiveInt
trials_num: PositiveInt
concurrency: PositiveInt
input_tokens: Optional[PositiveInt] = None
output_tokens: Optional[PositiveInt] = None
shared_prefix_tokens: Optional[int] = None
dataset: Optional[str] = None
baseline: bool = False
fleets: list[str] = Field(min_length=1)
env: list[str] = []
baseline: bool
fleets: Annotated[list[str], Field(min_length=1)]
env: list[str]


class PresetRandomConstraints(PresetConstraints):
"""Constraints for the synthetic `random` dataset, which the request shape defines."""

input_tokens: PositiveInt
output_tokens: Annotated[int, Field(ge=2)]
shared_prefix_tokens: Annotated[int, Field(ge=0)]


class PresetDatasetConstraints(PresetConstraints):
"""Constraints for a named dataset, which defines its own request shape."""

dataset: str


def _validate_model(value: Any, *, field: str) -> str:
Expand Down
Loading
Loading