diff --git a/docker/server/release/Dockerfile b/docker/server/release/Dockerfile index a88439d61e..df5de20fbe 100644 --- a/docker/server/release/Dockerfile +++ b/docker/server/release/Dockerfile @@ -22,6 +22,8 @@ ADD https://astral.sh/uv/install.sh /uv-installer.sh RUN sh /uv-installer.sh && rm /uv-installer.sh ENV PATH="/root/.local/bin/:$PATH" +RUN uv run --no-project --with claude-agent-sdk==0.2.110 python -c "from pathlib import Path; import claude_agent_sdk, shutil; shutil.copy2(Path(claude_agent_sdk.__file__).parent / '_bundled' / 'claude', '/usr/local/bin/claude')" && \ + chmod 755 /usr/local/bin/claude RUN uv tool install "dstack[all]==$VERSION" COPY entrypoint.sh entrypoint.sh diff --git a/docker/server/stgn/Dockerfile b/docker/server/stgn/Dockerfile index b2c4f96370..c8757808da 100644 --- a/docker/server/stgn/Dockerfile +++ b/docker/server/stgn/Dockerfile @@ -20,6 +20,8 @@ ADD https://astral.sh/uv/install.sh /uv-installer.sh RUN sh /uv-installer.sh && rm /uv-installer.sh ENV PATH="/root/.local/bin/:$PATH" +RUN uv run --no-project --with claude-agent-sdk==0.2.110 python -c "from pathlib import Path; import claude_agent_sdk, shutil; shutil.copy2(Path(claude_agent_sdk.__file__).parent / '_bundled' / 'claude', '/usr/local/bin/claude')" && \ + chmod 755 /usr/local/bin/claude COPY pyproject.toml uv.lock README.md ./ COPY src src RUN uv sync --extra all diff --git a/mkdocs/docs/reference/env.md b/mkdocs/docs/reference/env.md index 86a7dd051d..4a3ab58f85 100644 --- a/mkdocs/docs/reference/env.md +++ b/mkdocs/docs/reference/env.md @@ -162,6 +162,18 @@ slows down processing and may cause CPU spikes due to frequent SSH-connection es `https://dstack.example.com/{arch}/{version}/dstack-runner`. * `DSTACK_SHIM_DOWNLOAD_URL` – Overrides `dstack-shim` binary download URL. The URL can contain `{version}` and/or `{arch}` placeholders, see `DSTACK_RUNNER_DOWNLOAD_URL` for the details. + * `DSTACK_AGENT_CLAUDE_USE_EXISTING_AUTH` – Allows the endpoint agent to use + the server user's existing Claude CLI authentication when + `DSTACK_AGENT_ANTHROPIC_API_KEY` is not set. This is only for local + development. Do not set it together with `DSTACK_AGENT_ANTHROPIC_API_KEY`. + Production servers must use `DSTACK_AGENT_ANTHROPIC_API_KEY`. + This mode runs Claude without `--bare`, so Claude may read the server + user's Claude CLI auth and settings. It also passes the server process + `USER` to Claude because existing Claude CLI auth requires it. + * `DSTACK_AGENT_CLAUDE_EFFORT` – Passes Claude CLI `--effort` for endpoint + agent sessions. Supported values are `low`, `medium`, `high`, `xhigh`, + and `max`. If unset, dstack does not pass `--effort` and Claude CLI uses + its default. * `DSTACK_DEFAULT_CREDS_DISABLED` – Disables default credentials detection if set. Defaults to `None`. ## CLI diff --git a/pyproject.toml b/pyproject.toml index ca08438931..1081770403 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,11 +62,17 @@ artifacts = [ "src/dstack/_internal/server/statics/**", ] +[tool.hatch.build.targets.sdist.force-include] +"skills" = "skills" + [tool.hatch.build.targets.wheel] artifacts = [ "src/dstack/_internal/server/statics/**", ] +[tool.hatch.build.targets.wheel.force-include] +"skills" = "skills" + [tool.hatch.metadata.hooks.fancy-pypi-readme] content-type = "text/markdown" diff --git a/skills/dstack-prototyping/SKILL.md b/skills/dstack-prototyping/SKILL.md new file mode 100644 index 0000000000..516a912846 --- /dev/null +++ b/skills/dstack-prototyping/SKILL.md @@ -0,0 +1,92 @@ +--- +name: dstack-prototyping +description: | + Use with the dstack skill for model-serving work when the image, serving command, resources, backend/fleet choice, or service behavior is not proven. Guides task-first prototyping on real hardware, choosing fleets/backends that can reuse idle instances and caches, checking vLLM/SGLang sources, and verifying the final dstack service with a model request. +--- + +# dstack Prototyping + +Use `/dstack` for CLI commands, YAML fields, apply/attach behavior, service URLs, +and other dstack syntax. This skill explains how to use dstack runs while the +model-serving recipe is still unknown. + +## Goal + +Find a working dstack service recipe for the requested model. + +Before submitting a service, use a task on real hardware to test the serving +image, install/runtime assumptions, model download, cache path, command, port, +launch flags, resources, env vars, backend/fleet choice, and local model +request. Then submit the same recipe as a service and verify the model through +the dstack service URL. + +## Choose Where To Run + +Choose only VM-based backends, SSH fleets, or Kubernetes fleets because they support idle instances and/or instance volumes. That lets later runs reuse the provisioned/idle instance or instance volumes used by runs for caching model weights (and possibly other writes). You must follow this rule even if there are fleets/backends/offers that are cheaper. The only exception from this rule is when the required GPU class (regardless of the price) is not available through VM-based backend, SSH fleet, or Kubernetes fleet. + +Read `https://dstack.ai/docs/concepts/backends.md` to know exactly which +backends are VM-based. + +## Check Serving Sources + +Check serving-framework sources early enough to choose the image, command, +launch flags, resources, cache paths, request format, and expected model +behavior. + +For vLLM and SGLang, use these as credible sources: + +- vLLM recipes and model index: `https://recipes.vllm.ai/` and + `https://recipes.vllm.ai/models.json` +- vLLM recipe docs: `https://docs.vllm.ai/projects/recipes/en/stable/` +- SGLang docs and cookbook: `https://docs.sglang.ai/` and + `https://lmsysorg.mintlify.app/cookbook/intro` + +Use deeper serving-engine writeups, such as +`https://www.lmsys.org/blog/2026-07-02-agent-assisted-sglang-development`, when +recipes and docs do not explain the model, hardware, or serving failure. + +## Use A Task Before Service + +Before submitting a service, start a long-lived task: + +```yaml +commands: + - sleep infinity +``` + +or an equivalent idle command. + +Submit the task detached, attach or SSH into it when available, and run commands +inside the live environment. Test the image, installs, model download and cache +path, serving command, port, launch flags, local model request, and expected +model behavior. + +If the image, hardware choice, or major install path changes, submit another +task so the changed setup is tested before service verification. + +Do not move to a service after checking only GPU visibility, imports, logs, or a +health endpoint. Start the server inside the task and send a request that uses +the requested model. For a chat or reasoning model, check the response behavior +the endpoint is expected to support, such as reasoning output when that model is +supposed to expose it. + +If you stop a run to submit another one, whether it is a task or a service, +ensure the previous run has fully stopped, for example by reaching `stopped`, +`terminated`, or another terminal status. This can prevent provisioning extra +instances by reusing idle instances when possible, and can also preserve +instance volumes. + +## Verify As A Service + +Submit the service after the task has verified the recipe: image, command, port, +resources, env vars, cache mounts if used, backend/fleet choice, and model +request. + +Use the service as a duplicate check of the same recipe under dstack service +runtime. The model request that worked locally in the task must also work +through the dstack service URL. + +If service verification fails because the image, install, model download, +command, resources, cache, or model behavior needs to change, go back to a task. +If the recipe is still right and only the service config is wrong, fix the +service config and submit the service again. diff --git a/src/dstack/_internal/cli/commands/endpoint.py b/src/dstack/_internal/cli/commands/endpoint.py new file mode 100644 index 0000000000..89635f1cb1 --- /dev/null +++ b/src/dstack/_internal/cli/commands/endpoint.py @@ -0,0 +1,319 @@ +import argparse +import time +from typing import Iterable + +from rich.live import Live + +from dstack._internal.cli.commands import APIBaseCommand +from dstack._internal.cli.services.completion import ( + EndpointNameCompleter, + EndpointPresetNameCompleter, +) +from dstack._internal.cli.services.endpoint_logs import ( + EndpointLogPoller, + EndpointLogRecord, + print_endpoint_log, +) +from dstack._internal.cli.utils.common import ( + LIVE_TABLE_PROVISION_INTERVAL_SECS, + LIVE_TABLE_REFRESH_RATE_PER_SEC, + confirm_ask, + console, + get_start_time, +) +from dstack._internal.cli.utils.endpoint import ( + filter_endpoints_for_listing, + get_endpoints_table, + print_endpoint, + print_endpoints_table, +) +from dstack._internal.cli.utils.preset import print_endpoint_presets_table +from dstack._internal.core.errors import ResourceNotExistsError +from dstack._internal.core.models.endpoints import Endpoint +from dstack._internal.utils.json_utils import pydantic_orjson_dumps_with_indent + + +class EndpointCommand(APIBaseCommand): + NAME = "endpoint" + DESCRIPTION = "Manage endpoints" + + def _register(self): + super()._register() + self._parser.set_defaults(subfunc=self._list) + subparsers = self._parser.add_subparsers(dest="action") + + list_parser = subparsers.add_parser( + "list", help="List endpoints", formatter_class=self._parser.formatter_class + ) + list_parser.set_defaults(subfunc=self._list) + + for parser in [self._parser, list_parser]: + parser.add_argument( + "-a", + "--all", + help=( + "Show all endpoints. By default, it shows unfinished endpoints, " + "or the last finished endpoint if there are no unfinished endpoints." + ), + action="store_true", + ) + parser.add_argument( + "-w", + "--watch", + help="Update listing in realtime", + action="store_true", + ) + parser.add_argument( + "-v", "--verbose", action="store_true", help="Show more information" + ) + parser.add_argument( + "-n", + "--last", + help="Show only the last N endpoints. Implies --all", + type=int, + default=None, + ) + + logs_parser = subparsers.add_parser( + "logs", + help="Show endpoint logs", + formatter_class=self._parser.formatter_class, + ) + logs_parser.add_argument( + "name", + help="The name of the endpoint", + ).completer = EndpointNameCompleter() # type: ignore[attr-defined] + logs_parser.add_argument( + "-w", + "--watch", + help="Watch endpoint logs in realtime", + action="store_true", + ) + logs_parser.add_argument( + "--since", + help=( + "Show only logs newer than the specified date." + " Can be a duration (e.g. 10s, 5m, 1d) or an RFC 3339 string (e.g. 2023-09-24T15:30:00Z)." + ), + type=str, + ) + logs_parser.set_defaults(subfunc=self._logs) + + stop_parser = subparsers.add_parser( + "stop", + help="Stop an endpoint", + formatter_class=self._parser.formatter_class, + ) + stop_parser.add_argument( + "name", + help="The name of the endpoint", + ).completer = EndpointNameCompleter() # type: ignore[attr-defined] + stop_parser.add_argument( + "-y", "--yes", help="Don't ask for confirmation", action="store_true" + ) + stop_parser.set_defaults(subfunc=self._stop) + + get_parser = subparsers.add_parser( + "get", help="Get an endpoint", formatter_class=self._parser.formatter_class + ) + get_parser.add_argument( + "name", + metavar="NAME", + help="The name of the endpoint", + ).completer = EndpointNameCompleter() # type: ignore[attr-defined] + get_parser.add_argument( + "--json", + action="store_true", + help="Output in JSON format", + ) + get_parser.set_defaults(subfunc=self._get) + + preset_parser = subparsers.add_parser( + "preset", + help="Manage endpoint presets", + formatter_class=self._parser.formatter_class, + ) + preset_parser.add_argument( + "-v", "--verbose", action="store_true", help="Show more information" + ) + preset_parser.set_defaults(subfunc=self._preset_list) + preset_subparsers = preset_parser.add_subparsers(dest="preset_action") + + preset_list_parser = preset_subparsers.add_parser( + "list", + help="List endpoint presets", + formatter_class=self._parser.formatter_class, + ) + preset_list_parser.add_argument( + "-v", + "--verbose", + action="store_true", + default=argparse.SUPPRESS, + help="Show more information", + ) + preset_list_parser.set_defaults(subfunc=self._preset_list) + + preset_get_parser = preset_subparsers.add_parser( + "get", + help="Get an endpoint preset", + formatter_class=self._parser.formatter_class, + ) + preset_get_parser.add_argument( + "model", + metavar="MODEL", + help="The model of the endpoint preset", + ).completer = EndpointPresetNameCompleter() # type: ignore[attr-defined] + preset_get_parser.add_argument( + "--json", + action="store_true", + help="Output in JSON format", + ) + preset_get_parser.set_defaults(subfunc=self._preset_get) + + preset_delete_parser = preset_subparsers.add_parser( + "delete", + help="Delete endpoint presets", + formatter_class=self._parser.formatter_class, + ) + preset_delete_parser.add_argument( + "model", + metavar="MODEL", + help="The model of the endpoint preset", + ).completer = EndpointPresetNameCompleter() # type: ignore[attr-defined] + preset_delete_parser.add_argument( + "-y", "--yes", help="Don't ask for confirmation", action="store_true" + ) + preset_delete_parser.set_defaults(subfunc=self._preset_delete) + + def _command(self, args: argparse.Namespace): + super()._command(args) + args.subfunc(args) + + def _list(self, args: argparse.Namespace): + endpoints = self._get_endpoints_for_listing(args) + if not args.watch: + print_endpoints_table(endpoints, verbose=args.verbose) + return + + try: + with Live(console=console, refresh_per_second=LIVE_TABLE_REFRESH_RATE_PER_SEC) as live: + while True: + live.update(get_endpoints_table(endpoints, verbose=args.verbose)) + time.sleep(LIVE_TABLE_PROVISION_INTERVAL_SECS) + endpoints = self._get_endpoints_for_listing(args) + except KeyboardInterrupt: + pass + + def _get_endpoints_for_listing(self, args: argparse.Namespace): + endpoints = self.api.client.endpoints.list(self.api.project) + return filter_endpoints_for_listing( + endpoints, + show_all=args.all, + limit=args.last, + ) + + def _logs(self, args: argparse.Namespace): + try: + endpoint = self.api.client.endpoints.get( + project_name=self.api.project, + name=args.name, + ) + except ResourceNotExistsError: + console.print("Endpoint not found") + exit(1) + + start_time = get_start_time(args.since) + try: + for log in self._get_endpoint_logs( + endpoint=endpoint, start_time=start_time, watch=args.watch + ): + print_endpoint_log(log) + except KeyboardInterrupt: + pass + + def _get_endpoint_logs( + self, endpoint: Endpoint, start_time, watch: bool = False + ) -> Iterable[EndpointLogRecord]: + poller = EndpointLogPoller(api=self.api, endpoint=endpoint, start_time=start_time) + while True: + yield from poller.poll() + if not watch: + break + time.sleep(LIVE_TABLE_PROVISION_INTERVAL_SECS) + + def _stop(self, args: argparse.Namespace): + try: + endpoint = self.api.client.endpoints.get(project_name=self.api.project, name=args.name) + except ResourceNotExistsError: + console.print(f"Endpoint [code]{args.name}[/] does not exist") + exit(1) + + if endpoint.status.is_finished(): + console.print(f"Endpoint [code]{args.name}[/] is already {endpoint.status.value}") + return + + if not args.yes and not confirm_ask(f"Stop the endpoint [code]{args.name}[/]?"): + console.print("\nExiting...") + return + + with console.status("Stopping endpoint..."): + self.api.client.endpoints.stop(project_name=self.api.project, names=[args.name]) + + console.print(f"Endpoint [code]{args.name}[/] stopping") + + def _get(self, args: argparse.Namespace): + try: + endpoint = self.api.client.endpoints.get( + project_name=self.api.project, + name=args.name, + ) + except ResourceNotExistsError: + console.print("Endpoint not found") + exit(1) + + if args.json: + print(pydantic_orjson_dumps_with_indent(endpoint.dict(), default=None)) + return + print_endpoint(endpoint) + + def _preset_list(self, args: argparse.Namespace): + presets = self.api.client.endpoint_presets.list(self.api.project) + print_endpoint_presets_table(presets, verbose=args.verbose) + + def _preset_get(self, args: argparse.Namespace): + if not args.json: + console.print("Use --json to output the endpoint preset.") + exit(1) + try: + preset = self.api.client.endpoint_presets.get( + project_name=self.api.project, + model=args.model, + ) + except ResourceNotExistsError: + console.print(f"Endpoint preset for model [code]{args.model}[/] does not exist") + exit(1) + print(pydantic_orjson_dumps_with_indent(preset.dict(), default=None)) + + def _preset_delete(self, args: argparse.Namespace): + presets = self.api.client.endpoint_presets.list(self.api.project) + if args.model not in {preset.base for preset in presets}: + console.print(f"Endpoint preset for model [code]{args.model}[/] does not exist") + exit(1) + + if not args.yes and not confirm_ask( + f"Delete the endpoint preset for model [code]{args.model}[/]?" + ): + console.print("\nExiting...") + return + + try: + with console.status("Deleting endpoint preset..."): + self.api.client.endpoint_presets.delete( + project_name=self.api.project, + models=[args.model], + ) + except ResourceNotExistsError: + console.print(f"Endpoint preset for model [code]{args.model}[/] does not exist") + exit(1) + + console.print(f"Endpoint preset for model [code]{args.model}[/] deleted") diff --git a/src/dstack/_internal/cli/commands/logs.py b/src/dstack/_internal/cli/commands/logs.py index 78cde52f49..575dfef8db 100644 --- a/src/dstack/_internal/cli/commands/logs.py +++ b/src/dstack/_internal/cli/commands/logs.py @@ -1,5 +1,6 @@ import argparse import sys +from typing import Iterable from dstack._internal.cli.commands import APIBaseCommand from dstack._internal.cli.services.completion import RunNameCompleter @@ -39,24 +40,31 @@ def _register(self): ), type=str, ) - self._parser.add_argument("run_name").completer = RunNameCompleter(all=True) # type: ignore[attr-defined] + self._parser.add_argument("run_name").completer = RunNameCompleter() # type: ignore[attr-defined] def _command(self, args: argparse.Namespace): super()._command(args) - run = self.api.runs.get(args.run_name) - if run is None: - raise CLIError(f"Run {args.run_name} not found") - start_time = get_start_time(args.since) - logs = run.logs( - start_time=start_time, - diagnose=args.diagnose, - replica_num=args.replica, - job_num=args.job, - ) + logs = self._get_logs(args=args, start_time=start_time) try: for log in logs: sys.stdout.buffer.write(log) sys.stdout.buffer.flush() except KeyboardInterrupt: pass + + def _get_logs( + self, + args: argparse.Namespace, + start_time, + ) -> Iterable[bytes]: + run = self.api.runs.get(args.run_name) + if run is not None: + return run.logs( + start_time=start_time, + diagnose=args.diagnose, + replica_num=args.replica, + job_num=args.job, + ) + + raise CLIError(f"Run {args.run_name} not found") diff --git a/src/dstack/_internal/cli/main.py b/src/dstack/_internal/cli/main.py index 32f15a95f8..335f7693f1 100644 --- a/src/dstack/_internal/cli/main.py +++ b/src/dstack/_internal/cli/main.py @@ -8,6 +8,7 @@ from dstack._internal.cli.commands.attach import AttachCommand from dstack._internal.cli.commands.completion import CompletionCommand from dstack._internal.cli.commands.delete import DeleteCommand +from dstack._internal.cli.commands.endpoint import EndpointCommand from dstack._internal.cli.commands.event import EventCommand from dstack._internal.cli.commands.export import ExportCommand from dstack._internal.cli.commands.fleet import FleetCommand @@ -67,6 +68,7 @@ def main(): ApplyCommand.register(subparsers) AttachCommand.register(subparsers) DeleteCommand.register(subparsers) + EndpointCommand.register(subparsers) EventCommand.register(subparsers) ExportCommand.register(subparsers) FleetCommand.register(subparsers) diff --git a/src/dstack/_internal/cli/services/completion.py b/src/dstack/_internal/cli/services/completion.py index 4fa276945d..67083980af 100644 --- a/src/dstack/_internal/cli/services/completion.py +++ b/src/dstack/_internal/cli/services/completion.py @@ -70,6 +70,16 @@ def fetch_resource_names(self, api: Client) -> Iterable[str]: return [r.name for r in api.client.volumes.list(api.project)] +class EndpointNameCompleter(BaseAPINameCompleter): + def fetch_resource_names(self, api: Client) -> Iterable[str]: + return [r.name for r in api.client.endpoints.list(api.project)] + + +class EndpointPresetNameCompleter(BaseAPINameCompleter): + def fetch_resource_names(self, api: Client) -> Iterable[str]: + return [r.model for r in api.client.endpoint_presets.list(api.project)] + + class GatewayNameCompleter(BaseAPINameCompleter): def fetch_resource_names(self, api: Client) -> Iterable[str]: return [r.name for r in api.client.gateways.list(api.project)] diff --git a/src/dstack/_internal/cli/services/configurators/__init__.py b/src/dstack/_internal/cli/services/configurators/__init__.py index 91768bdcd3..e5cbd0f69f 100644 --- a/src/dstack/_internal/cli/services/configurators/__init__.py +++ b/src/dstack/_internal/cli/services/configurators/__init__.py @@ -5,6 +5,7 @@ import yaml from dstack._internal.cli.services.configurators.base import BaseApplyConfigurator +from dstack._internal.cli.services.configurators.endpoint import EndpointConfigurator from dstack._internal.cli.services.configurators.fleet import FleetConfigurator from dstack._internal.cli.services.configurators.gateway import GatewayConfigurator from dstack._internal.cli.services.configurators.run import ( @@ -32,6 +33,7 @@ DevEnvironmentConfigurator, TaskConfigurator, ServiceConfigurator, + EndpointConfigurator, FleetConfigurator, GatewayConfigurator, VolumeConfigurator, diff --git a/src/dstack/_internal/cli/services/configurators/endpoint.py b/src/dstack/_internal/cli/services/configurators/endpoint.py new file mode 100644 index 0000000000..149c06c47f --- /dev/null +++ b/src/dstack/_internal/cli/services/configurators/endpoint.py @@ -0,0 +1,418 @@ +import argparse +import shutil +import time +from pathlib import Path + +from rich.table import Table + +from dstack._internal.cli.services.configurators.base import ( + ApplyEnvVarsConfiguratorMixin, + BaseApplyConfigurator, +) +from dstack._internal.cli.services.endpoint_logs import ( + EndpointLogPoller, + EndpointLogRecord, + print_endpoint_log, +) +from dstack._internal.cli.services.profile import apply_profile_args, register_profile_args +from dstack._internal.cli.utils.common import ( + NO_OFFERS_WARNING, + confirm_ask, + console, + format_backend, + format_instance_availability, +) +from dstack._internal.cli.utils.endpoint import get_endpoints_table +from dstack._internal.cli.utils.rich import MultiItemStatus +from dstack._internal.core.errors import ConfigurationError, ResourceNotExistsError +from dstack._internal.core.models.configurations import ApplyConfigurationType +from dstack._internal.core.models.endpoints import ( + AnyEndpointProvisioningPlan, + Endpoint, + EndpointConfiguration, + EndpointPlan, + EndpointPresetPolicy, + EndpointProvisioningPlanAgent, + EndpointProvisioningPlanNone, + EndpointProvisioningPlanPreset, + EndpointStatus, +) +from dstack._internal.core.models.profiles import Profile, ProfileParams, SpotPolicy +from dstack._internal.utils.common import local_time, make_proxy_url +from dstack.api.utils import load_profile + +_NO_ENDPOINT_FLEETS_WARNING = ( + "[error]" + "The project has no fleets. Create one before submitting an endpoint.\n" + "See [link]https://dstack.ai/docs/guides/troubleshooting/#no-fleets[/link]" + "[/]\n" +) + + +class EndpointConfigurator( + ApplyEnvVarsConfiguratorMixin, + BaseApplyConfigurator[EndpointConfiguration], +): + TYPE = ApplyConfigurationType.ENDPOINT + + def apply_configuration( + self, + conf: EndpointConfiguration, + configuration_path: str, + command_args: argparse.Namespace, + configurator_args: argparse.Namespace, + ): + self.apply_args(conf, configurator_args) + with console.status("Getting apply plan..."): + plan = self.api.client.endpoints.get_plan( + project_name=self.api.project, + configuration=conf, + configuration_path=configuration_path, + ) + no_fleets = False + if _preset_plan_has_no_offers(plan.provisioning_plan): + if len(self.api.client.fleets.list(self.api.project, include_imported=True)) == 0: + no_fleets = True + _print_endpoint_plan(plan, no_fleets=no_fleets) + + confirm_message = _get_apply_confirm_message(plan) + current_resource = _get_non_terminal_current_resource(plan) + stop_run_name = None + if current_resource is not None: + if current_resource.configuration == plan.configuration: + console.print( + f"Endpoint [code]{current_resource.name}[/] already exists." + " Detected no changes." + ) + if command_args.yes and not command_args.force: + console.print("Use --force to apply anyway.") + return + else: + # TODO: Replace v1 stop/recreate with endpoint in-place update + # via service rolling deployment once endpoint versioning exists. + console.print( + f"Endpoint [code]{current_resource.name}[/] already exists." + " Detected changes that [error]cannot[/] be updated in-place." + ) + else: + serving_run_name = _get_endpoint_serving_run_name(plan) + same_name_run = None + if serving_run_name is not None: + same_name_run = self.api.runs.get(serving_run_name) + if same_name_run is not None and not same_name_run.status.is_finished(): + stop_run_name = same_name_run.name + console.print( + f"Active run [code]{stop_run_name}[/] already exists." + " The endpoint will use this backing service run name." + ) + confirm_message = _get_apply_confirm_message( + plan, + stop_run_name=stop_run_name, + ) + + if not command_args.yes and not confirm_ask(confirm_message): + console.print("\nExiting...") + return + + if current_resource is not None: + with console.status("Stopping existing endpoint..."): + self.api.client.endpoints.stop( + project_name=self.api.project, + names=[current_resource.name], + ) + while True: + try: + endpoint = self.api.client.endpoints.get( + project_name=self.api.project, + name=current_resource.name, + ) + except ResourceNotExistsError: + break + if endpoint.status.is_finished(): + break + time.sleep(1) + + if stop_run_name is not None: + with console.status("Stopping run..."): + self.api.client.runs.stop(self.api.project, [stop_run_name], abort=False) + while True: + run = self.api.runs.get(stop_run_name) + if run is None or run.status.is_finished(): + break + time.sleep(1) + + with console.status("Creating endpoint..."): + endpoint = self.api.client.endpoints.create( + project_name=self.api.project, + configuration=conf, + ) + if command_args.detach: + _print_submitted_endpoint_message(endpoint) + return + self._follow_endpoint_apply(endpoint) + + def delete_configuration( + self, + conf: EndpointConfiguration, + configuration_path: str, + command_args: argparse.Namespace, + ): + raise ConfigurationError( + "`dstack delete` does not support endpoint configurations. " + "Use `dstack endpoint stop `." + ) + + @classmethod + def register_args(cls, parser: argparse.ArgumentParser): + configuration_group = parser.add_argument_group(f"{cls.TYPE.value} Options") + configuration_group.add_argument( + "-n", + "--name", + dest="name", + help="The endpoint name", + ) + cls.register_env_args(configuration_group) + register_profile_args(parser) + + def apply_args(self, conf: EndpointConfiguration, args: argparse.Namespace): + profile = load_profile(Path.cwd(), args.profile) + _apply_profile(conf, profile) + apply_profile_args(args, conf) + if args.name: + conf.name = args.name + self.apply_env_vars(conf.env, args) + + def _follow_endpoint_apply(self, endpoint: Endpoint): + log_poller = EndpointLogPoller(api=self.api, endpoint=endpoint) + try: + with MultiItemStatus(_get_apply_status(endpoint), console=console) as live: + while not _is_endpoint_apply_finished(endpoint): + _print_endpoint_log_batch(log_poller.poll()) + live.update( + get_endpoints_table([endpoint], format_date=local_time), + status=_get_apply_status(endpoint), + ) + time.sleep(2) + endpoint = self.api.client.endpoints.get( + project_name=self.api.project, + name=endpoint.name, + ) + _print_endpoint_log_batch(log_poller.poll()) + except KeyboardInterrupt: + console.print("\nDetached") + return + + _make_endpoint_url_absolute(endpoint, self.api.client.base_url) + console.print( + get_endpoints_table( + [endpoint], + verbose=endpoint.status == EndpointStatus.RUNNING, + format_date=local_time, + ) + ) + console.print( + f"\nEndpoint [code]{endpoint.name}[/] provisioning completed " + f"[secondary]({endpoint.status.value})[/]" + ) + _print_finished_endpoint_message(endpoint) + if endpoint.status == EndpointStatus.FAILED: + exit(1) + + +def _print_endpoint_log_batch(logs: list[EndpointLogRecord]) -> None: + for log in logs: + print_endpoint_log(log) + + +def _apply_profile(conf: EndpointConfiguration, profile: Profile): + for field in ProfileParams.__fields__: + value = getattr(profile, field) + if value is not None: + setattr(conf, field, value) + + +def _is_endpoint_apply_finished(endpoint: Endpoint) -> bool: + return endpoint.status in ( + EndpointStatus.RUNNING, + EndpointStatus.STOPPED, + EndpointStatus.FAILED, + ) + + +def _get_apply_status(endpoint: Endpoint) -> str: + return f"Provisioning endpoint [code]{endpoint.name}[/]..." + + +def _print_finished_endpoint_message(endpoint: Endpoint) -> None: + if endpoint.status == EndpointStatus.RUNNING: + if endpoint.url is not None: + console.print(f"[code]{endpoint.url}[/code]") + return + + message = endpoint.status_message or endpoint.error + if message is None: + message = endpoint.status.value.capitalize() + console.print(f"[error]{message}[/error]") + + +def _print_endpoint_plan(plan: EndpointPlan, no_fleets: bool = False): + def th(s: str) -> str: + return f"[bold]{s}[/bold]" + + props = Table(box=None, show_header=False) + props.add_column(no_wrap=True) + props.add_column() + props.add_row(th("Project"), plan.project_name) + props.add_row(th("User"), plan.user) + if plan.configuration_path is not None: + props.add_row(th("Configuration"), plan.configuration_path) + props.add_row(th("Type"), plan.configuration.type) + props.add_row(th("Spot policy"), _format_spot_policy(plan)) + props.add_row(th("Max price"), _format_max_price(plan)) + props.add_row(th("Preset policy"), plan.preset_policy.value) + if isinstance(plan.provisioning_plan, EndpointProvisioningPlanPreset): + props.add_row(th("Preset"), plan.provisioning_plan.preset_base) + props.add_row(th("Recipe"), plan.provisioning_plan.recipe_id) + console.print(props) + console.print() + + if isinstance(plan.provisioning_plan, EndpointProvisioningPlanPreset): + _print_preset_plan_offers(plan.provisioning_plan, no_fleets=no_fleets) + elif isinstance(plan.provisioning_plan, EndpointProvisioningPlanNone): + style = "error" if plan.preset_policy == EndpointPresetPolicy.CREATE else "warning" + console.print(f"[{style}]{plan.provisioning_plan.reason}[/]") + console.print() + elif isinstance(plan.provisioning_plan, EndpointProvisioningPlanAgent): + # TODO: Consider showing initial candidate offers for agent provisioning as + # non-binding context. Do not present them as selected resources/final hardware. + if plan.provisioning_plan.reason is not None: + console.print(f"[warning]{plan.provisioning_plan.reason}[/]") + console.print() + + +def _get_apply_confirm_message( + plan: EndpointPlan, + stop_run_name: str | None = None, +) -> str: + current_resource = _get_non_terminal_current_resource(plan) + if current_resource is not None: + return f"Stop and override the endpoint [code]{current_resource.name}[/]?" + if stop_run_name is not None: + return f"Stop and override the run [code]{stop_run_name}[/]?" + return "Create the endpoint?" + + +def _get_endpoint_serving_run_name(plan: EndpointPlan) -> str | None: + if _get_non_terminal_current_resource(plan) is not None: + return None + if isinstance(plan.provisioning_plan, EndpointProvisioningPlanPreset): + return plan.provisioning_plan.service_name + return None + + +def _get_non_terminal_current_resource(plan: EndpointPlan) -> Endpoint | None: + if plan.current_resource is None or plan.current_resource.status.is_finished(): + return None + return plan.current_resource + + +def _format_spot_policy(plan: EndpointPlan) -> str: + if isinstance(plan.provisioning_plan, EndpointProvisioningPlanPreset): + values = {job_offers.spot for job_offers in plan.provisioning_plan.job_offers} + if len(values) == 1: + return _format_spot(next(iter(values))) + return "mixed" + return _format_endpoint_spot_policy(plan.configuration.spot_policy) + + +def _format_endpoint_spot_policy(spot_policy: SpotPolicy | None) -> str: + if spot_policy == SpotPolicy.SPOT: + return "spot" + if spot_policy == SpotPolicy.ONDEMAND: + return "on-demand" + return "auto" + + +def _format_spot(spot: bool | None) -> str: + if spot is None: + return "auto" + return "spot" if spot else "on-demand" + + +def _format_max_price(plan: EndpointPlan) -> str: + if isinstance(plan.provisioning_plan, EndpointProvisioningPlanPreset): + values = {job_offers.max_price for job_offers in plan.provisioning_plan.job_offers} + if len(values) == 1: + return _format_price_limit(next(iter(values))) + return "mixed" + return _format_price_limit(plan.configuration.max_price) + + +def _format_price_limit(max_price: float | None) -> str: + return f"${max_price:3f}".rstrip("0").rstrip(".") if max_price else "off" + + +def _preset_plan_has_no_offers( + provisioning_plan: AnyEndpointProvisioningPlan, +) -> bool: + if not isinstance(provisioning_plan, EndpointProvisioningPlanPreset): + return False + return not any(job_offers.offers for job_offers in provisioning_plan.job_offers) + + +def _print_preset_plan_offers( + provisioning_plan: EndpointProvisioningPlanPreset, + no_fleets: bool = False, +) -> None: + if _preset_plan_has_no_offers(provisioning_plan): + console.print(_NO_ENDPOINT_FLEETS_WARNING if no_fleets else NO_OFFERS_WARNING) + return + + show_replica_groups = len(provisioning_plan.job_offers) > 1 or any( + job_offers.replica_group != "0" for job_offers in provisioning_plan.job_offers + ) + offers = Table(box=None, expand=shutil.get_terminal_size(fallback=(120, 40)).columns <= 110) + offers.add_column("#") + if show_replica_groups: + offers.add_column("GROUP", no_wrap=True) + offers.add_column("BACKEND", style="grey58") + offers.add_column("RESOURCES") + offers.add_column("INSTANCE TYPE", style="grey58", no_wrap=True) + offers.add_column("PRICE", style="grey58") + offers.add_column() + offer_num = 0 + total_offers = 0 + for job_offers in provisioning_plan.job_offers: + total_offers += job_offers.total_offers + for offer in job_offers.offers: + offer_num += 1 + row = [ + str(offer_num), + format_backend(offer.backend, offer.region), + offer.instance.resources.pretty_format(include_spot=True), + offer.instance.name, + f"${offer.price:.4f}".rstrip("0").rstrip("."), + format_instance_availability(offer.availability), + ] + if show_replica_groups: + row.insert(1, job_offers.replica_group) + offers.add_row(*row, style=None if offer_num == 1 else "secondary") + if job_offers.total_offers > len(job_offers.offers): + row = ["", "..."] + if show_replica_groups: + row.insert(1, job_offers.replica_group) + offers.add_row(*row, style="secondary") + console.print(offers) + if total_offers > offer_num: + console.print(f"[secondary] Shown {offer_num} of {total_offers} offers[/]") + console.print() + + +def _print_submitted_endpoint_message(endpoint: Endpoint) -> None: + console.print(f"Endpoint [code]{endpoint.name}[/] submitted, detaching...") + + +def _make_endpoint_url_absolute(endpoint: Endpoint, server_url: str) -> None: + if endpoint.url is None: + return + endpoint.url = make_proxy_url(server_url=server_url, proxy_url=endpoint.url) diff --git a/src/dstack/_internal/cli/services/endpoint_logs.py b/src/dstack/_internal/cli/services/endpoint_logs.py new file mode 100644 index 0000000000..fe0623b10e --- /dev/null +++ b/src/dstack/_internal/cli/services/endpoint_logs.py @@ -0,0 +1,97 @@ +import base64 +from collections import Counter +from dataclasses import dataclass +from datetime import datetime, timedelta +from typing import Optional + +from rich.text import Text + +from dstack._internal.cli.utils.common import console +from dstack._internal.core.models.endpoints import Endpoint +from dstack._internal.core.models.logs import LogEvent +from dstack._internal.server.schemas.logs import PollLogsRequest + +_ENDPOINT_LOG_WATCH_OVERLAP = timedelta(seconds=20) + + +@dataclass(frozen=True) +class EndpointLogRecord: + timestamp: datetime + message: str + + +class EndpointLogPoller: + def __init__(self, *, api, endpoint: Endpoint, start_time: Optional[datetime] = None) -> None: + self._api = api + self._endpoint = endpoint + self._next_start_time = start_time + self._latest_printed_at: Optional[datetime] = None + self._seen_log_counts: Counter[tuple[datetime, str, str]] = Counter() + + def poll(self) -> list[EndpointLogRecord]: + logs: list[EndpointLogRecord] = [] + next_token = None + batch_log_counts: Counter[tuple[datetime, str, str]] = Counter() + while True: + resp = self._api.client.logs.poll( + project_name=self._api.project, + body=PollLogsRequest( + run_name=self._endpoint.name, + job_submission_id=self._endpoint.id, + start_time=self._next_start_time, + end_time=None, + descending=False, + limit=1000, + diagnose=False, + next_token=next_token, + ), + ) + for log in resp.logs: + if not self._should_print(log, batch_log_counts=batch_log_counts): + continue + logs.append(_decode_log(log)) + next_token = resp.next_token + if next_token is None: + break + if self._latest_printed_at is not None: + self._next_start_time = self._latest_printed_at - _ENDPOINT_LOG_WATCH_OVERLAP + self._cleanup_seen_log_counts(before=self._next_start_time) + return logs + + def _should_print( + self, log: LogEvent, *, batch_log_counts: Counter[tuple[datetime, str, str]] + ) -> bool: + key = _log_key(log) + batch_log_counts[key] += 1 + if batch_log_counts[key] <= self._seen_log_counts[key]: + return False + self._seen_log_counts[key] += 1 + if self._latest_printed_at is None or log.timestamp > self._latest_printed_at: + self._latest_printed_at = log.timestamp + return True + + def _cleanup_seen_log_counts(self, before: datetime) -> None: + for key in list(self._seen_log_counts): + timestamp, _, _ = key + if timestamp <= before: + del self._seen_log_counts[key] + + +def _log_key(log: LogEvent) -> tuple[datetime, str, str]: + return (log.timestamp, log.log_source.value, log.message) + + +def print_endpoint_log(log: EndpointLogRecord) -> None: + timestamp = log.timestamp.astimezone().strftime("%Y-%m-%d %H:%M:%S") + console.print( + Text(f"[{timestamp}]", style="log.time"), + Text(log.message.rstrip("\r\n"), style="log.message"), + soft_wrap=True, + ) + + +def _decode_log(log: LogEvent) -> EndpointLogRecord: + return EndpointLogRecord( + timestamp=log.timestamp, + message=base64.b64decode(log.message).decode(errors="replace"), + ) diff --git a/src/dstack/_internal/cli/services/profile.py b/src/dstack/_internal/cli/services/profile.py index e8086bb43b..63f2b4fbe5 100644 --- a/src/dstack/_internal/cli/services/profile.py +++ b/src/dstack/_internal/cli/services/profile.py @@ -6,6 +6,7 @@ from dstack._internal.core.models.profiles import ( CreationPolicy, Profile, + ProfileParams, ProfileRetry, SpotPolicy, parse_duration, @@ -135,7 +136,7 @@ def register_profile_args(parser: argparse.ArgumentParser): def apply_profile_args( args: argparse.Namespace, - profile_settings: Union[Profile, AnyRunConfiguration], + profile_settings: Union[Profile, ProfileParams, AnyRunConfiguration], ): """ Overrides `profile_settings` settings with arguments registered by `register_profile_args()`. diff --git a/src/dstack/_internal/cli/utils/endpoint.py b/src/dstack/_internal/cli/utils/endpoint.py new file mode 100644 index 0000000000..46f932626c --- /dev/null +++ b/src/dstack/_internal/cli/utils/endpoint.py @@ -0,0 +1,191 @@ +from typing import List, Optional + +from rich.table import Table + +from dstack._internal.cli.utils.common import add_row_from_dict, console +from dstack._internal.core.models.endpoints import Endpoint, EndpointStatus +from dstack._internal.utils.common import DateFormatter, pretty_date + + +def filter_endpoints_for_listing( + endpoints: List[Endpoint], + show_all: bool = False, + limit: int | None = None, +) -> List[Endpoint]: + endpoints = sorted( + endpoints, + key=lambda endpoint: (endpoint.created_at, str(endpoint.id)), + reverse=True, + ) + if limit is not None: + return endpoints[:limit] + if show_all: + return endpoints + + unfinished = [endpoint for endpoint in endpoints if not _is_endpoint_finished(endpoint)] + if unfinished: + return unfinished + return endpoints[:1] + + +def print_endpoints_table(endpoints: List[Endpoint], verbose: bool = False): + table = get_endpoints_table(endpoints, verbose=verbose) + console.print(table) + console.print() + + +def print_endpoint(endpoint: Endpoint): + console.print(get_endpoint_table(endpoint)) + console.print() + + +def get_endpoint_table( + endpoint: Endpoint, + format_date: DateFormatter = pretty_date, +) -> Table: + table = Table(box=None, show_header=False) + table.add_column(no_wrap=True) + table.add_column() + + def th(value: str) -> str: + return f"[bold]{value}[/bold]" + + table.add_row(th("Project"), endpoint.project_name) + table.add_row(th("User"), endpoint.user) + table.add_row(th("Endpoint"), endpoint.name) + model = endpoint.configuration.model.api_model_name + table.add_row(th("Model"), model) + if endpoint.model_repo is not None and endpoint.model_repo != model: + table.add_row(th("Repo"), endpoint.model_repo) + table.add_row( + th("Status"), + _format_endpoint_status( + endpoint.status, + endpoint.status_message, + ), + ) + table.add_row(th("Preset policy"), endpoint.configuration.preset_policy.value) + table.add_row(th("Service run"), endpoint.run_name or "-") + table.add_row(th("URL"), endpoint.url or "-") + table.add_row(th("Created"), format_date(endpoint.created_at)) + if endpoint.status_message: + table.add_row(th("Error"), endpoint.status_message) + return table + + +def get_endpoints_table( + endpoints: List[Endpoint], + verbose: bool = False, + format_date: DateFormatter = pretty_date, +) -> Table: + table = Table(box=None) + table.add_column("NAME", no_wrap=True) + table.add_column("MODEL") + table.add_column("STATUS", no_wrap=True) + table.add_column("POLICY", no_wrap=True) + if verbose: + table.add_column("SERVICE RUN") + table.add_column("URL") + table.add_column("CREATED") + if verbose: + table.add_column("ERROR") + + for endpoint in endpoints: + model = endpoint.configuration.model.api_model_name + row = { + "NAME": endpoint.name, + "MODEL": model, + "STATUS": _format_endpoint_status( + endpoint.status, + endpoint.status_message, + ), + "POLICY": endpoint.configuration.preset_policy.value, + "SERVICE RUN": endpoint.run_name or "-", + "URL": endpoint.url or "-", + "CREATED": format_date(endpoint.created_at), + "ERROR": endpoint.status_message, + } + add_row_from_dict(table, row) + if endpoint.model_repo is not None and endpoint.model_repo != model: + add_row_from_dict( + table, + {"MODEL": f" repo={endpoint.model_repo}"}, + style="secondary", + ) + return table + + +def _format_endpoint_status( + status: EndpointStatus, + status_message: Optional[str] = None, +) -> str: + status_value = _get_endpoint_status_value(status, status_message) + color_map = { + EndpointStatus.SUBMITTED: "grey", + EndpointStatus.PROVISIONING: "deep_sky_blue1", + EndpointStatus.PROTOTYPING: "medium_purple1", + EndpointStatus.RUNNING: "sea_green3", + EndpointStatus.STOPPING: "deep_sky_blue1", + EndpointStatus.STOPPED: "grey62", + EndpointStatus.FAILED: "indian_red1", + } + if status_value == "no offers": + style = "gold1" + else: + style = color_map.get(status, "white") + if not status.is_finished(): + style = f"bold {style}" + return f"[{style}]{status_value}[/]" + + +def _get_endpoint_status_value( + status: EndpointStatus, + status_message: Optional[str], +) -> str: + if status == EndpointStatus.FAILED: + failure_reason = _get_endpoint_failure_reason(status_message) + if failure_reason is not None: + return failure_reason + return status.value + + +def _is_endpoint_finished(endpoint: Endpoint) -> bool: + return endpoint.status.is_finished() + + +def _get_endpoint_failure_reason(status_message: Optional[str]) -> Optional[str]: + if status_message is None: + return None + normalized = " ".join(status_message.lower().split()) + if not normalized: + return None + if any( + phrase in normalized + for phrase in [ + "no matching instance offers", + "no matching offers", + "no offers", + "no-offer", + "zero offers", + "total_offers: 0", + ] + ): + return "no offers" + if "no fleets" in normalized: + return "no fleets" + if ( + "requires the server agent" in normalized + or "server agent runtime" in normalized + or "dstack_agent_" in normalized + or "claude executable" in normalized + ): + return "no agent" + if "no matching endpoint presets" in normalized: + return "no preset" + if "server agent" in normalized or "agent" in normalized: + return "agent failed" + if "run name" in normalized and "taken" in normalized: + return "conflict" + if "backing service run" in normalized: + return "run failed" + return None diff --git a/src/dstack/_internal/cli/utils/preset.py b/src/dstack/_internal/cli/utils/preset.py new file mode 100644 index 0000000000..a8f5a4b274 --- /dev/null +++ b/src/dstack/_internal/cli/utils/preset.py @@ -0,0 +1,126 @@ +from typing import List + +from rich.table import Table + +from dstack._internal.cli.utils.common import add_row_from_dict, console +from dstack._internal.core.models.endpoint_presets import EndpointPreset, EndpointPresetRecipe +from dstack._internal.utils.common import pretty_resources + + +def print_endpoint_presets_table(presets: List[EndpointPreset], verbose: bool = False): + table = get_endpoint_presets_table(presets, verbose=verbose) + console.print(table) + console.print() + + +def get_endpoint_presets_table(presets: List[EndpointPreset], verbose: bool = False) -> Table: + table = Table(box=None) + table.add_column("MODEL", no_wrap=True) + if verbose: + table.add_column("RESOURCES") + else: + table.add_column("GPU") + + for preset in presets: + if len(preset.recipes) == 1: + _add_recipe_rows( + table, + preset_base=preset.base, + label=f"[bold]{preset.base}[/]", + recipe=preset.recipes[0], + verbose=verbose, + ) + continue + add_row_from_dict(table, {"MODEL": f"[bold]{preset.base}[/]"}) + for recipe_num, recipe in enumerate(preset.recipes): + _add_recipe_rows( + table, + preset_base=preset.base, + label=f"[secondary] recipe={recipe_num}[/]", + recipe=recipe, + verbose=verbose, + ) + return table + + +def _add_recipe_rows( + table: Table, + preset_base: str, + label: str, + recipe: EndpointPresetRecipe, + verbose: bool, +) -> None: + groups = recipe.service.replica_groups + column = "RESOURCES" if verbose else "GPU" + if len(groups) == 1: + add_row_from_dict( + table, + { + "MODEL": label, + column: _format_resources(groups[0].resources, verbose=verbose), + }, + ) + _add_recipe_model_row( + table, + preset_base=preset_base, + recipe_model=recipe.model, + ) + return + + add_row_from_dict(table, {"MODEL": label}) + _add_recipe_model_row( + table, + preset_base=preset_base, + recipe_model=recipe.model, + ) + for group in groups: + add_row_from_dict( + table, + { + "MODEL": f"[secondary] group={group.name}[/]", + column: _format_resources(group.resources, verbose=verbose), + }, + ) + + +def _add_recipe_model_row( + table: Table, + preset_base: str, + recipe_model: str, +) -> None: + if recipe_model == preset_base: + return + add_row_from_dict(table, {"MODEL": f" repo={recipe_model}"}, style="secondary") + + +def _format_resources(resources, verbose: bool) -> str: + if not verbose: + if resources.gpu is None: + return "-" + return _format_gpu(resources) + formatted = resources.pretty_format() + if resources.gpu is not None and resources.gpu.count.min == 0: + if resources.gpu.count.max in (None, 0): + return ( + formatted.replace(" gpu=0..", "") + .replace(" gpu=0", "") + .replace("gpu=0..", "-") + .replace("gpu=0", "-") + ) + return formatted + + +def _format_gpu(resources) -> str: + gpu = resources.gpu + assert gpu is not None + if gpu.count.min == 0 and gpu.count.max in (None, 0): + return "-" + formatted = pretty_resources( + gpu_vendor=gpu.vendor, + gpu_name=",".join(gpu.name) if gpu.name else None, + gpu_count=gpu.count, + gpu_memory=gpu.memory, + total_gpu_memory=gpu.total_memory, + compute_capability=gpu.compute_capability, + ) + return formatted.removeprefix("gpu=") or "-" diff --git a/src/dstack/_internal/core/models/configurations.py b/src/dstack/_internal/core/models/configurations.py index 900dca8ce7..2240d5c23e 100644 --- a/src/dstack/_internal/core/models/configurations.py +++ b/src/dstack/_internal/core/models/configurations.py @@ -18,6 +18,7 @@ RegistryAuth, generate_dual_core_model, ) +from dstack._internal.core.models.endpoints import EndpointConfiguration from dstack._internal.core.models.envs import Env from dstack._internal.core.models.files import FilePathMapping from dstack._internal.core.models.fleets import FleetConfiguration @@ -1385,6 +1386,7 @@ class ApplyConfigurationType(str, Enum): DEV_ENVIRONMENT = "dev-environment" TASK = "task" SERVICE = "service" + ENDPOINT = "endpoint" FLEET = "fleet" GATEWAY = "gateway" VOLUME = "volume" @@ -1392,6 +1394,7 @@ class ApplyConfigurationType(str, Enum): AnyApplyConfiguration = Union[ AnyRunConfiguration, + EndpointConfiguration, FleetConfiguration, GatewayConfiguration, AnyVolumeConfiguration, @@ -1412,6 +1415,7 @@ class BaseApplyConfiguration(CoreModel): Union[ # Final configurations AnyRunConfiguration, + EndpointConfiguration, FleetConfiguration, GatewayConfiguration, # Base configurations (further parsing required to get a concrete AnyApplyConfiguration) @@ -1439,6 +1443,7 @@ def parse_apply_configuration(data: dict) -> AnyApplyConfiguration: AnyDstackConfiguration = Union[ AnyRunConfiguration, + EndpointConfiguration, FleetConfiguration, GatewayConfiguration, VolumeConfiguration, diff --git a/src/dstack/_internal/core/models/endpoint_presets.py b/src/dstack/_internal/core/models/endpoint_presets.py new file mode 100644 index 0000000000..f3ed18ff8e --- /dev/null +++ b/src/dstack/_internal/core/models/endpoint_presets.py @@ -0,0 +1,27 @@ +from dstack._internal.core.models.common import CoreModel +from dstack._internal.core.models.configurations import ServiceConfiguration +from dstack._internal.core.models.resources import ResourcesSpec + + +class EndpointPresetValidationReplica(CoreModel): + resources: list[ResourcesSpec] + """Exact resources for each running replica in this service replica group.""" + + +class EndpointPresetValidation(CoreModel): + replicas: list[EndpointPresetValidationReplica] + """Ordered to match `ServiceConfiguration.replica_groups`.""" + + +class EndpointPresetRecipe(CoreModel): + id: str + model: str + """Exact repo/path loaded by this recipe's service command.""" + service: ServiceConfiguration + validations: list[EndpointPresetValidation] + + +class EndpointPreset(CoreModel): + base: str + """Base/requested/API model name used for preset lookup and client requests.""" + recipes: list[EndpointPresetRecipe] diff --git a/src/dstack/_internal/core/models/endpoints.py b/src/dstack/_internal/core/models/endpoints.py new file mode 100644 index 0000000000..ab95fd43c5 --- /dev/null +++ b/src/dstack/_internal/core/models/endpoints.py @@ -0,0 +1,219 @@ +import uuid +from datetime import datetime +from enum import Enum +from typing import Annotated, Any, Literal, Optional, Union + +from pydantic import Field, validator + +from dstack._internal.core.models.common import ( + ApplyAction, + CoreModel, + generate_dual_core_model, +) +from dstack._internal.core.models.envs import Env +from dstack._internal.core.models.instances import InstanceOfferWithAvailability +from dstack._internal.core.models.profiles import ProfileParams, ProfileParamsConfig +from dstack._internal.core.models.resources import ResourcesSpec +from dstack._internal.utils.json_schema import add_extra_schema_types + + +class EndpointStatus(str, Enum): + SUBMITTED = "submitted" + PROVISIONING = "provisioning" + PROTOTYPING = "prototyping" + RUNNING = "running" + STOPPING = "stopping" + STOPPED = "stopped" + FAILED = "failed" + + @classmethod + def finished_statuses(cls) -> list["EndpointStatus"]: + return [cls.STOPPED, cls.FAILED] + + def is_finished(self) -> bool: + return self in self.finished_statuses() + + +class EndpointPresetPolicy(str, Enum): + REUSE = "reuse" + CREATE = "create" + REUSE_OR_CREATE = "reuse-or-create" + + +class EndpointConfigurationConfig(ProfileParamsConfig): + @staticmethod + def schema_extra(schema: dict): + ProfileParamsConfig.schema_extra(schema) + add_extra_schema_types( + schema["properties"]["model"], + extra_types=[{"type": "string"}], + ) + + +class EndpointModelRepo(CoreModel): + repo: str + """Exact repo/path to deploy.""" + name: Optional[str] = None + """API-facing model name. Defaults to `repo`.""" + + @property + def api_model_name(self) -> str: + return self.name or self.repo + + @property + def exact_repo(self) -> str: + return self.repo + + @property + def allows_variant_selection(self) -> bool: + return False + + @validator("repo") + def validate_repo(cls, v: str) -> str: + return _validate_endpoint_model_string(v, field_name="repo") + + @validator("name") + def validate_name(cls, v: Optional[str]) -> Optional[str]: + if v is None: + return v + return _validate_endpoint_model_string(v, field_name="name") + + +class EndpointModelBase(CoreModel): + base: str + """Base model name. The agent may choose a compatible repo/path to deploy.""" + + @property + def api_model_name(self) -> str: + return self.base + + @property + def exact_repo(self) -> None: + return None + + @property + def allows_variant_selection(self) -> bool: + return True + + @validator("base") + def validate_base(cls, v: str) -> str: + return _validate_endpoint_model_string(v, field_name="base") + + +EndpointModelSpec = Union[EndpointModelRepo, EndpointModelBase] + + +class EndpointConfiguration( + ProfileParams, + generate_dual_core_model(EndpointConfigurationConfig), +): + type: Literal["endpoint"] = "endpoint" + name: Annotated[ + Optional[str], + Field(description="The endpoint name. If not specified, a random name is generated"), + ] = None + + model: Annotated[ + EndpointModelSpec, + Field( + description=( + "The model to serve. Use a string or `repo` for an exact repo/path, " + "or `base` to let the endpoint agent choose a compatible variant." + ) + ), + ] + env: Annotated[ + Env, + Field(description="The mapping or the list of environment variables"), + ] = Env() + # TODO: Add endpoint-level resources only with hard scheduling semantics for both + # single-service and replica-group presets. V1 intentionally relies on preset/service + # resources plus ProfileParams constraints. + preset_policy: Annotated[ + EndpointPresetPolicy, + Field( + description=( + "The policy for endpoint presets. `reuse` uses an existing preset only, " + "`create` asks the server agent to create and save a new preset, and " + "`reuse-or-create` first tries an existing preset and falls back to `create`." + ) + ), + ] = EndpointPresetPolicy.REUSE_OR_CREATE + + @validator("model", pre=True) + def parse_model(cls, v) -> dict | EndpointModelSpec: + if isinstance(v, str): + return {"repo": _validate_endpoint_model_string(v, field_name="model")} + return v + + +class Endpoint(CoreModel): + id: uuid.UUID + name: str + project_name: str + user: str + configuration: EndpointConfiguration + created_at: datetime + last_processed_at: datetime + status: EndpointStatus + status_message: Optional[str] = None + run_name: Optional[str] = None + url: Optional[str] = None + model_base: Optional[str] = None + """Base model repo for the endpoint's current deployed model.""" + model_repo: Optional[str] = None + """Exact repo/path deployed by the endpoint's current service.""" + error: Optional[str] = None + + +class EndpointProvisioningPlanNone(CoreModel): + type: Literal["none"] = "none" + reason: str + + +class EndpointPlanJobOffers(CoreModel): + replica_group: str + resources: ResourcesSpec + spot: Optional[bool] + max_price: Optional[float] + offers: list[InstanceOfferWithAvailability] + total_offers: int + max_offer_price: Optional[float] + + +class EndpointProvisioningPlanPreset(CoreModel): + type: Literal["preset"] = "preset" + preset_base: str + recipe_id: str + service_name: str + job_offers: list[EndpointPlanJobOffers] + + +class EndpointProvisioningPlanAgent(CoreModel): + type: Literal["agent"] = "agent" + agent_model: str + reason: Optional[str] = None + + +AnyEndpointProvisioningPlan = Union[ + EndpointProvisioningPlanNone, + EndpointProvisioningPlanPreset, + EndpointProvisioningPlanAgent, +] + + +class EndpointPlan(CoreModel): + project_name: str + user: str + configuration: EndpointConfiguration + configuration_path: Optional[str] = None + current_resource: Optional[Endpoint] = None + action: ApplyAction + preset_policy: EndpointPresetPolicy + provisioning_plan: Annotated[AnyEndpointProvisioningPlan, Field(discriminator="type")] + + +def _validate_endpoint_model_string(value: Any, *, field_name: str) -> str: + if not isinstance(value, str) or not value.strip(): + raise ValueError(f"Endpoint model {field_name} must be a non-empty string") + return value diff --git a/src/dstack/_internal/core/models/events.py b/src/dstack/_internal/core/models/events.py index f2efb80d0e..11f3881725 100644 --- a/src/dstack/_internal/core/models/events.py +++ b/src/dstack/_internal/core/models/events.py @@ -16,6 +16,7 @@ class EventTargetType(str, Enum): INSTANCE = "instance" RUN = "run" JOB = "job" + ENDPOINT = "endpoint" VOLUME = "volume" GATEWAY = "gateway" SECRET = "secret" diff --git a/src/dstack/_internal/server/app.py b/src/dstack/_internal/server/app.py index 0f02806aa4..918757282f 100644 --- a/src/dstack/_internal/server/app.py +++ b/src/dstack/_internal/server/app.py @@ -29,6 +29,7 @@ from dstack._internal.server.routers import ( auth, backends, + endpoints, events, exports, files, @@ -255,6 +256,8 @@ def register_routes(app: FastAPI, ui: bool = True): app.include_router(logs.router) app.include_router(secrets.router) app.include_router(gateways.router) + app.include_router(endpoints.root_router) + app.include_router(endpoints.project_router) app.include_router(volumes.root_router) app.include_router(volumes.project_router) app.include_router(service_proxy.router, prefix="/proxy/services", tags=["proxy"]) diff --git a/src/dstack/_internal/server/background/pipeline_tasks/__init__.py b/src/dstack/_internal/server/background/pipeline_tasks/__init__.py index 2e83e780c0..d06a472481 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/__init__.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/__init__.py @@ -2,6 +2,7 @@ from dstack._internal.server.background.pipeline_tasks.base import Pipeline from dstack._internal.server.background.pipeline_tasks.compute_groups import ComputeGroupPipeline +from dstack._internal.server.background.pipeline_tasks.endpoints import EndpointPipeline from dstack._internal.server.background.pipeline_tasks.fleets import FleetPipeline from dstack._internal.server.background.pipeline_tasks.gateway_replicas import ( GatewayReplicaPipeline, @@ -44,6 +45,7 @@ def __init__(self) -> None: PlacementGroupPipeline(pipeline_hinter=self._hinter), RunPipeline(pipeline_hinter=self._hinter), ServiceRouterWorkerSyncPipeline(pipeline_hinter=self._hinter), + EndpointPipeline(pipeline_hinter=self._hinter), VolumePipeline(pipeline_hinter=self._hinter), ]: self.register_pipeline(builtin_pipeline) diff --git a/src/dstack/_internal/server/background/pipeline_tasks/endpoints.py b/src/dstack/_internal/server/background/pipeline_tasks/endpoints.py new file mode 100644 index 0000000000..3e5683c2b4 --- /dev/null +++ b/src/dstack/_internal/server/background/pipeline_tasks/endpoints.py @@ -0,0 +1,1334 @@ +import asyncio +import uuid +from dataclasses import dataclass, field +from datetime import timedelta +from typing import Optional, Sequence + +from pydantic import ValidationError +from sqlalchemy import or_, select, update +from sqlalchemy.orm import joinedload, load_only + +from dstack._internal.core.errors import ServerClientError +from dstack._internal.core.models.configurations import ServiceConfiguration +from dstack._internal.core.models.endpoints import ( + EndpointConfiguration, + EndpointPresetPolicy, + EndpointStatus, +) +from dstack._internal.core.models.runs import ( + ApplyRunPlanInput, + JobStatus, + RunSpec, + RunStatus, + ServiceSpec, +) +from dstack._internal.server.background.pipeline_tasks.base import ( + Fetcher, + Heartbeater, + ItemUpdateMap, + Pipeline, + PipelineItem, + Worker, + log_lock_token_changed_after_processing, + log_lock_token_mismatch, + resolve_now_placeholders, + set_processed_update_map_fields, + set_unlock_update_map_fields, +) +from dstack._internal.server.db import get_db, get_session_ctx +from dstack._internal.server.models import ( + EndpointModel, + EndpointRunSubmissionModel, + ProjectModel, + RunModel, + UserModel, +) +from dstack._internal.server.services import runs as runs_services +from dstack._internal.server.services.endpoints import ( + can_use_endpoint_agent, + emit_endpoint_status_change_event, + get_endpoint_agent_admin_required_message, + get_endpoint_configuration, + get_endpoint_no_fleets_message, + has_endpoint_existing_usable_fleets, + record_endpoint_run_submission, +) +from dstack._internal.server.services.endpoints.agent import ( + abort_agent_endpoint, + get_agent_service, + get_agent_unavailable_reason, +) +from dstack._internal.server.services.endpoints.names import get_endpoint_serving_run_name +from dstack._internal.server.services.endpoints.planning import find_preset_planning_result +from dstack._internal.server.services.endpoints.preset_building import ( + build_endpoint_preset_from_run, +) +from dstack._internal.server.services.endpoints.presets import get_endpoint_preset_service +from dstack._internal.server.services.locking import get_locker +from dstack._internal.server.services.pipelines import PipelineHinterProtocol +from dstack._internal.server.utils import sentry_utils +from dstack._internal.utils.common import get_current_datetime +from dstack._internal.utils.logging import get_logger + +logger = get_logger(__name__) + +_NO_MATCHING_PRESET_MESSAGE = "No matching endpoint presets found." +_MAX_AGENT_STATUS_MESSAGE_CHARS = 500 + + +@dataclass +class EndpointPipelineItem(PipelineItem): + status: EndpointStatus + + +class EndpointPipeline(Pipeline[EndpointPipelineItem]): + def __init__( + self, + workers_num: int = 4, + queue_lower_limit_factor: float = 0.5, + queue_upper_limit_factor: float = 2.0, + min_processing_interval: timedelta = timedelta(seconds=10), + lock_timeout: timedelta = timedelta(seconds=30), + heartbeat_trigger: timedelta = timedelta(seconds=15), + *, + pipeline_hinter: PipelineHinterProtocol, + ) -> None: + super().__init__( + workers_num=workers_num, + queue_lower_limit_factor=queue_lower_limit_factor, + queue_upper_limit_factor=queue_upper_limit_factor, + min_processing_interval=min_processing_interval, + lock_timeout=lock_timeout, + heartbeat_trigger=heartbeat_trigger, + ) + self.__heartbeater = Heartbeater[EndpointPipelineItem]( + model_type=EndpointModel, + lock_timeout=self._lock_timeout, + heartbeat_trigger=self._heartbeat_trigger, + ) + self.__fetcher = EndpointFetcher( + queue=self._queue, + queue_desired_minsize=self._queue_desired_minsize, + min_processing_interval=self._min_processing_interval, + lock_timeout=self._lock_timeout, + heartbeater=self._heartbeater, + ) + self.__workers = [ + EndpointWorker( + queue=self._queue, + heartbeater=self._heartbeater, + pipeline_hinter=pipeline_hinter, + ) + for _ in range(self._workers_num) + ] + + @property + def hint_fetch_model_name(self) -> str: + return EndpointModel.__name__ + + @property + def _heartbeater(self) -> Heartbeater[EndpointPipelineItem]: + return self.__heartbeater + + @property + def _fetcher(self) -> Fetcher[EndpointPipelineItem]: + return self.__fetcher + + @property + def _workers(self) -> Sequence["EndpointWorker"]: + return self.__workers + + +class EndpointFetcher(Fetcher[EndpointPipelineItem]): + def __init__( + self, + queue: asyncio.Queue[EndpointPipelineItem], + queue_desired_minsize: int, + min_processing_interval: timedelta, + lock_timeout: timedelta, + heartbeater: Heartbeater[EndpointPipelineItem], + queue_check_delay: float = 1.0, + ) -> None: + super().__init__( + queue=queue, + queue_desired_minsize=queue_desired_minsize, + min_processing_interval=min_processing_interval, + lock_timeout=lock_timeout, + heartbeater=heartbeater, + queue_check_delay=queue_check_delay, + ) + + @sentry_utils.instrument_pipeline_task("EndpointFetcher.fetch") + async def fetch(self, limit: int) -> list[EndpointPipelineItem]: + endpoint_lock, _ = get_locker(get_db().dialect_name).get_lockset( + EndpointModel.__tablename__ + ) + async with endpoint_lock: + async with get_session_ctx() as session: + now = get_current_datetime() + res = await session.execute( + select(EndpointModel) + .where( + EndpointModel.status.in_( + [ + EndpointStatus.SUBMITTED, + EndpointStatus.PROVISIONING, + EndpointStatus.PROTOTYPING, + EndpointStatus.RUNNING, + EndpointStatus.STOPPING, + ] + ), + or_( + EndpointModel.last_processed_at <= now - self._min_processing_interval, + EndpointModel.last_processed_at == EndpointModel.created_at, + ), + or_( + EndpointModel.lock_expires_at.is_(None), + EndpointModel.lock_expires_at < now, + ), + or_( + EndpointModel.lock_owner.is_(None), + EndpointModel.lock_owner == EndpointPipeline.__name__, + ), + ) + .order_by(EndpointModel.last_processed_at.asc()) + .limit(limit) + .with_for_update(skip_locked=True, key_share=True, of=EndpointModel) + .options( + load_only( + EndpointModel.id, + EndpointModel.lock_token, + EndpointModel.lock_expires_at, + EndpointModel.status, + ) + ) + ) + endpoint_models = list(res.scalars().all()) + lock_expires_at = get_current_datetime() + self._lock_timeout + lock_token = uuid.uuid4() + items = [] + for endpoint_model in endpoint_models: + prev_lock_expired = endpoint_model.lock_expires_at is not None + endpoint_model.lock_expires_at = lock_expires_at + endpoint_model.lock_token = lock_token + endpoint_model.lock_owner = EndpointPipeline.__name__ + items.append( + EndpointPipelineItem( + __tablename__=EndpointModel.__tablename__, + id=endpoint_model.id, + lock_expires_at=lock_expires_at, + lock_token=lock_token, + prev_lock_expired=prev_lock_expired, + status=endpoint_model.status, + ) + ) + await session.commit() + return items + + +class EndpointWorker(Worker[EndpointPipelineItem]): + def __init__( + self, + queue: asyncio.Queue[EndpointPipelineItem], + heartbeater: Heartbeater[EndpointPipelineItem], + pipeline_hinter: PipelineHinterProtocol, + ) -> None: + super().__init__( + queue=queue, + heartbeater=heartbeater, + pipeline_hinter=pipeline_hinter, + ) + + @sentry_utils.instrument_pipeline_task("EndpointWorker.process") + async def process(self, item: EndpointPipelineItem): + endpoint_model = await _refetch_locked_endpoint(item) + if endpoint_model is None: + log_lock_token_mismatch(logger, item) + return + + if endpoint_model.status == EndpointStatus.SUBMITTED: + result = await _process_submitted_endpoint( + endpoint_model=endpoint_model, + pipeline_hinter=self._pipeline_hinter, + ) + elif endpoint_model.status == EndpointStatus.PROVISIONING: + result = await _process_provisioning_endpoint( + endpoint_model=endpoint_model, + pipeline_hinter=self._pipeline_hinter, + ) + elif endpoint_model.status == EndpointStatus.PROTOTYPING: + result = await _process_prototyping_endpoint( + endpoint_model=endpoint_model, + pipeline_hinter=self._pipeline_hinter, + ) + elif endpoint_model.status == EndpointStatus.STOPPING: + result = await _process_stopping_endpoint( + endpoint_model=endpoint_model, + pipeline_hinter=self._pipeline_hinter, + ) + elif endpoint_model.status == EndpointStatus.RUNNING: + result = await _process_running_endpoint(endpoint_model) + else: + result = _ProcessResult() + + runs_to_stop = await _get_endpoint_runs_to_stop_after_failure(endpoint_model, result) + for run_to_stop in runs_to_stop: + logger.info( + "Stopping backing run %s after endpoint %s failed", + run_to_stop.run_name, + endpoint_model.name, + ) + await _stop_backing_run( + endpoint_model=endpoint_model, + run_name=run_to_stop.run_name, + pipeline_hinter=self._pipeline_hinter, + ) + + await _apply_process_result(item=item, endpoint_model=endpoint_model, result=result) + + +async def _refetch_locked_endpoint(item: EndpointPipelineItem) -> Optional[EndpointModel]: + async with get_session_ctx() as session: + res = await session.execute( + select(EndpointModel) + .where( + EndpointModel.id == item.id, + EndpointModel.lock_token == item.lock_token, + ) + .options(joinedload(EndpointModel.project).joinedload(ProjectModel.backends)) + .options(joinedload(EndpointModel.project).joinedload(ProjectModel.members)) + .options( + joinedload(EndpointModel.user).load_only( + UserModel.name, + UserModel.global_role, + UserModel.token, + ) + ) + .options(joinedload(EndpointModel.service_run).selectinload(RunModel.jobs)) + ) + return res.unique().scalar_one_or_none() + + +async def _apply_process_result( + item: EndpointPipelineItem, + endpoint_model: EndpointModel, + result: "_ProcessResult", +): + update_map = _EndpointUpdateMap() + update_map.update(result.update_map) + set_processed_update_map_fields(update_map) + set_unlock_update_map_fields(update_map) + + async with get_session_ctx() as session: + resolve_now_placeholders(update_map, now=get_current_datetime()) + res = await session.execute( + update(EndpointModel) + .where( + EndpointModel.id == endpoint_model.id, + EndpointModel.lock_token == endpoint_model.lock_token, + EndpointModel.status == item.status, + ) + .values(**update_map) + .returning(EndpointModel.id) + ) + updated_ids = list(res.scalars().all()) + if len(updated_ids) == 0: + if await _link_reported_service_run_to_stopping_endpoint(session, item, result): + return + logger.info( + "Endpoint %s changed while being processed; ignoring stale result", + endpoint_model.name, + ) + return + emit_endpoint_status_change_event( + session=session, + endpoint_model=endpoint_model, + old_status=endpoint_model.status, + new_status=update_map.get("status", endpoint_model.status), + status_message=update_map.get("status_message", endpoint_model.status_message), + ) + + +async def _link_reported_service_run_to_stopping_endpoint( + session, + item: EndpointPipelineItem, + result: "_ProcessResult", +) -> bool: + service_run_id = result.update_map.get("service_run_id") + if service_run_id is None: + return False + + update_map = _EndpointUpdateMap(service_run_id=service_run_id) + model_repo = result.update_map.get("model_repo") + if model_repo is not None: + update_map["model_repo"] = model_repo + model_base = result.update_map.get("model_base") + if model_base is not None: + update_map["model_base"] = model_base + set_processed_update_map_fields(update_map) + set_unlock_update_map_fields(update_map) + + resolve_now_placeholders(update_map, now=get_current_datetime()) + res = await session.execute( + update(EndpointModel) + .where( + EndpointModel.id == item.id, + EndpointModel.lock_token == item.lock_token, + EndpointModel.status == EndpointStatus.STOPPING, + EndpointModel.service_run_id.is_(None), + ) + .values(**update_map) + .returning(EndpointModel.id) + ) + updated_ids = list(res.scalars().all()) + if len(updated_ids) == 0: + log_lock_token_changed_after_processing(logger, item) + return False + logger.info( + "Linked reported service run %s to stopping endpoint %s", + service_run_id, + item.id, + ) + return True + + +class _EndpointUpdateMap(ItemUpdateMap, total=False): + status: EndpointStatus + status_message: Optional[str] + service_run_id: uuid.UUID + model_base: Optional[str] + model_repo: Optional[str] + provisioning_method: Optional[str] + + +@dataclass +class _ProcessResult: + update_map: _EndpointUpdateMap = field(default_factory=_EndpointUpdateMap) + + +@dataclass(frozen=True) +class _PresetSubmission: + run_id: uuid.UUID + preset_base: str + recipe_id: str + model_repo: str + + +@dataclass(frozen=True) +class _PresetSubmissionResult: + submission: Optional[_PresetSubmission] = None + unprovisionable_preset: Optional[str] = None + + +async def _get_endpoint_runs_to_stop_after_failure( + endpoint_model: EndpointModel, + result: _ProcessResult, +) -> list[RunModel]: + if result.update_map.get("status") != EndpointStatus.FAILED: + return [] + return [ + run + for run in await _get_endpoint_unfinished_runs(endpoint_model) + if run.status != RunStatus.TERMINATING + ] + + +async def _get_endpoint_unfinished_runs(endpoint_model: EndpointModel) -> list[RunModel]: + runs: list[RunModel] = [] + seen_run_ids: set[uuid.UUID] = set() + if endpoint_model.service_run is not None: + runs.append(endpoint_model.service_run) + seen_run_ids.add(endpoint_model.service_run.id) + async with get_session_ctx() as session: + res = await session.execute( + select(RunModel) + .join( + EndpointRunSubmissionModel, + EndpointRunSubmissionModel.run_id == RunModel.id, + ) + .where( + EndpointRunSubmissionModel.endpoint_id == endpoint_model.id, + RunModel.deleted == False, + ) + ) + runs.extend(run for run in res.unique().scalars().all() if run.id not in seen_run_ids) + return [run for run in runs if not run.deleted and not run.status.is_finished()] + + +async def _process_submitted_endpoint( + endpoint_model: EndpointModel, + pipeline_hinter: PipelineHinterProtocol, +) -> _ProcessResult: + endpoint_configuration = get_endpoint_configuration(endpoint_model) + async with get_session_ctx() as session: + has_usable_fleets = await has_endpoint_existing_usable_fleets( + session=session, + project=endpoint_model.project, + configuration=endpoint_configuration, + ) + if not has_usable_fleets: + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": get_endpoint_no_fleets_message(endpoint_configuration), + } + ) + try: + submission_result = await _submit_endpoint_from_preset( + endpoint_id=endpoint_model.id, + pipeline_hinter=pipeline_hinter, + ) + except ServerClientError as e: + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": e.msg, + } + ) + if submission_result.submission is not None: + logger.info( + "Provisioning endpoint %s from preset %s recipe %s", + endpoint_model.name, + submission_result.submission.preset_base, + submission_result.submission.recipe_id, + ) + update_map = _EndpointUpdateMap( + status=EndpointStatus.PROVISIONING, + status_message=None, + service_run_id=submission_result.submission.run_id, + model_base=submission_result.submission.preset_base, + model_repo=submission_result.submission.model_repo, + provisioning_method=( + f"preset:{submission_result.submission.preset_base}" + f"#{submission_result.submission.recipe_id}" + ), + ) + return _ProcessResult(update_map=update_map) + + if await _should_provision_with_agent(endpoint_model): + logger.info("Provisioning endpoint %s with server agent", endpoint_model.name) + update_map = _EndpointUpdateMap( + status=EndpointStatus.PROTOTYPING, + status_message=None, + provisioning_method="agent", + ) + return _ProcessResult(update_map=update_map) + + logger.info("Failing endpoint %s: no preset path is available", endpoint_model.name) + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": _get_no_provisioning_path_message( + endpoint_model, + unprovisionable_preset=submission_result.unprovisionable_preset, + ), + } + ) + + +async def _should_provision_with_agent(endpoint_model: EndpointModel) -> bool: + endpoint_configuration = get_endpoint_configuration(endpoint_model) + if endpoint_configuration.preset_policy == EndpointPresetPolicy.REUSE: + return False + if not get_agent_service().is_enabled(): + return False + if not can_use_endpoint_agent(user=endpoint_model.user, project=endpoint_model.project): + return False + async with get_session_ctx() as session: + return await has_endpoint_existing_usable_fleets( + session=session, + project=endpoint_model.project, + configuration=endpoint_configuration, + ) + + +async def _get_active_serving_run_name_conflict( + endpoint_model: EndpointModel, +) -> Optional[_ProcessResult]: + serving_run_name = get_endpoint_serving_run_name(endpoint_model.name) + assert serving_run_name is not None + async with get_session_ctx() as session: + run_model = await runs_services.get_run_model_by_name( + session=session, + project=endpoint_model.project, + run_name=serving_run_name, + ) + if run_model is None: + return None + if endpoint_model.service_run_id == run_model.id: + return None + if run_model.status.is_finished(): + return None + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": f"Run name '{serving_run_name}' is taken by an existing run", + } + ) + + +async def _process_provisioning_endpoint( + endpoint_model: EndpointModel, + pipeline_hinter: PipelineHinterProtocol, +) -> _ProcessResult: + if endpoint_model.service_run is None: + if endpoint_model.provisioning_method == "agent": + return await _process_prototyping_endpoint( + endpoint_model=endpoint_model, + pipeline_hinter=pipeline_hinter, + ) + conflict_result = await _get_active_serving_run_name_conflict(endpoint_model) + if conflict_result is not None: + return conflict_result + + readiness = _get_backing_service_readiness(endpoint_model) + if readiness.failed_message is not None: + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": readiness.failed_message, + } + ) + if readiness.model_base_url is None or readiness.model_name is None: + return _ProcessResult() + if endpoint_model.provisioning_method == "agent": + # The agent's final report is the functional signal. This server-side gate only + # confirms that the verified run still looks like a normal ready dstack service. + await _try_save_agent_endpoint_preset( + endpoint_model=endpoint_model, + ) + if endpoint_model.service_run_id is not None: + await _stop_non_final_submitted_runs( + endpoint_model=endpoint_model, + final_run_id=endpoint_model.service_run_id, + pipeline_hinter=pipeline_hinter, + ) + return _ProcessResult( + update_map={ + "status": EndpointStatus.RUNNING, + "status_message": None, + } + ) + + +async def _process_prototyping_endpoint( + endpoint_model: EndpointModel, + pipeline_hinter: PipelineHinterProtocol, +) -> _ProcessResult: + if endpoint_model.service_run is not None: + return await _process_agent_verified_endpoint( + endpoint_model=endpoint_model, + pipeline_hinter=pipeline_hinter, + ) + return await _provision_endpoint_with_agent( + endpoint_model=endpoint_model, + pipeline_hinter=pipeline_hinter, + ) + + +async def _process_agent_verified_endpoint( + endpoint_model: EndpointModel, + pipeline_hinter: PipelineHinterProtocol, +) -> _ProcessResult: + run_model = endpoint_model.service_run + if run_model is None: + return _ProcessResult() + readiness = _get_service_run_readiness(run_model, endpoint_name=endpoint_model.name) + if readiness.failed_message is not None: + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": readiness.failed_message, + } + ) + if readiness.model_base_url is None or readiness.model_name is None: + return _ProcessResult(update_map={"status": EndpointStatus.PROTOTYPING}) + await _save_agent_endpoint_preset( + endpoint_model=endpoint_model, + run_model=run_model, + base_model=endpoint_model.model_base, + recipe_model=endpoint_model.model_repo, + ) + await _stop_non_final_submitted_runs( + endpoint_model=endpoint_model, + final_run_id=run_model.id, + pipeline_hinter=pipeline_hinter, + ) + return _ProcessResult( + update_map={ + "status": EndpointStatus.RUNNING, + "status_message": None, + } + ) + + +async def _provision_endpoint_with_agent( + endpoint_model: EndpointModel, + pipeline_hinter: PipelineHinterProtocol, +) -> _ProcessResult: + agent_service = get_agent_service() + if not agent_service.is_enabled(): + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": _get_no_provisioning_path_message(endpoint_model), + } + ) + + result = await agent_service.provision_endpoint( + endpoint_model=endpoint_model, + pipeline_hinter=pipeline_hinter, + ) + await _record_agent_submitted_runs( + endpoint_model=endpoint_model, + submitted_run_ids=result.submitted_run_ids, + submitted_run_names=result.submitted_run_names, + ) + if result.in_progress: + return _ProcessResult( + update_map={ + "status": EndpointStatus.PROTOTYPING, + "status_message": None, + } + ) + if result.error is not None: + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": _format_agent_status_message(result.error), + } + ) + report = result.final_report + if report is None: + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": "Server agent did not return a final report", + } + ) + if not report.success: + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": _format_agent_status_message( + report.failure_summary, + default="Server agent did not verify the endpoint", + ), + } + ) + run_id = result.run_id or report.run_id + if result.run_id is not None and report.run_id is not None and result.run_id != report.run_id: + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": "Server agent returned inconsistent run id in final report", + } + ) + if ( + result.run_name is not None + and report.run_name is not None + and result.run_name != report.run_name + ): + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": ("Server agent returned inconsistent run name in final report"), + } + ) + if run_id is None: + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": "Server agent final report did not identify a run id", + } + ) + model_base, model_repo, model_error = _get_agent_report_models( + endpoint_configuration=get_endpoint_configuration(endpoint_model), + report=report, + ) + if model_error is not None: + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": model_error, + } + ) + + async with get_session_ctx() as session: + res = await session.execute( + select(RunModel) + .where( + RunModel.id == run_id, + RunModel.project_id == endpoint_model.project_id, + RunModel.deleted == False, + ) + .options(joinedload(RunModel.user)) + .options(joinedload(RunModel.jobs)) + ) + run_model = res.unique().scalar_one_or_none() + if run_model is None: + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": (f"Server agent reported run '{run_id}' but it was not found"), + } + ) + if report.run_name is not None and report.run_name != run_model.run_name: + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": ("Server agent returned inconsistent run name in final report"), + } + ) + if not _is_valid_agent_submission_run_name(endpoint_model.name, run_model.run_name): + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": ( + "Server agent final service run name must be " + f"'{endpoint_model.name}-'" + ), + } + ) + if run_model.user_id != endpoint_model.user_id: + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": f"Run '{run_model.run_name}' is not owned by the endpoint user", + } + ) + try: + run_spec = RunSpec.__response__.parse_raw(run_model.run_spec) + except ValidationError: + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": f"Run '{run_model.run_name}' has invalid run spec", + } + ) + if not isinstance(run_spec.configuration, ServiceConfiguration): + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": f"Run '{run_model.run_name}' is not a service", + } + ) + endpoint_configuration = get_endpoint_configuration(endpoint_model) + requested_model_name = endpoint_configuration.model.api_model_name + if ( + run_spec.configuration.model is None + or run_spec.configuration.model.name != requested_model_name + ): + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": "Server agent final service model does not match the endpoint model", + } + ) + try: + await _record_endpoint_run_submission( + endpoint_id=endpoint_model.id, + run_id=run_model.id, + ) + except ServerClientError as e: + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": e.msg, + } + ) + readiness = _get_service_run_readiness(run_model, endpoint_name=endpoint_model.name) + if readiness.failed_message is not None: + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": readiness.failed_message, + } + ) + if readiness.model_base_url is None or readiness.model_name is None: + return _ProcessResult( + update_map={ + "status": EndpointStatus.PROTOTYPING, + "status_message": None, + "service_run_id": run_model.id, + "model_base": model_base, + "model_repo": model_repo, + } + ) + await _save_agent_endpoint_preset( + endpoint_model=endpoint_model, + run_model=run_model, + base_model=model_base, + recipe_model=model_repo, + ) + await _stop_non_final_submitted_runs( + endpoint_model=endpoint_model, + final_run_id=run_model.id, + pipeline_hinter=pipeline_hinter, + ) + return _ProcessResult( + update_map={ + "status": EndpointStatus.RUNNING, + "status_message": None, + "service_run_id": run_model.id, + "model_base": model_base, + "model_repo": model_repo, + } + ) + + +def _get_agent_report_models( + *, + endpoint_configuration: EndpointConfiguration, + report, +) -> tuple[Optional[str], Optional[str], Optional[str]]: + exact_model_repo = endpoint_configuration.model.exact_repo + if report.base is None: + return None, None, "Server agent final report did not identify the base model repo" + if endpoint_configuration.model.allows_variant_selection: + requested_base = endpoint_configuration.model.api_model_name + if report.base != requested_base: + return ( + None, + None, + "Server agent final report base does not match the requested base model", + ) + if report.model is None: + return None, None, "Server agent final report did not identify the selected model repo" + if exact_model_repo is not None and report.model != exact_model_repo: + return ( + None, + None, + "Server agent final report model does not match the requested model repo", + ) + return report.base, report.model, None + + +def _format_agent_status_message( + message: Optional[str], + default: str = "Server agent failed", +) -> str: + if message is None or not message.strip(): + return default + one_line_message = " ".join(message.split()) + if len(one_line_message) <= _MAX_AGENT_STATUS_MESSAGE_CHARS: + return one_line_message + return one_line_message[: _MAX_AGENT_STATUS_MESSAGE_CHARS - 3].rstrip() + "..." + + +async def _record_endpoint_run_submission(endpoint_id: uuid.UUID, run_id: uuid.UUID) -> None: + async with get_session_ctx() as session: + await record_endpoint_run_submission( + session=session, + endpoint_id=endpoint_id, + run_id=run_id, + ) + await session.commit() + + +async def _record_agent_submitted_runs( + *, + endpoint_model: EndpointModel, + submitted_run_ids: Sequence[uuid.UUID], + submitted_run_names: Sequence[str], +) -> None: + valid_submitted_run_names = [ + run_name + for run_name in submitted_run_names + if _is_valid_agent_submission_run_name(endpoint_model.name, run_name) + ] + if len(submitted_run_ids) == 0 and len(valid_submitted_run_names) == 0: + return + async with get_session_ctx() as session: + res = await session.execute( + select(RunModel).where( + or_( + RunModel.id.in_(submitted_run_ids), + RunModel.run_name.in_(valid_submitted_run_names), + ), + RunModel.project_id == endpoint_model.project_id, + RunModel.user_id == endpoint_model.user_id, + RunModel.deleted == False, + ) + ) + runs_by_id = {run.id: run for run in res.scalars().all()} + runs_by_name = {run.run_name: run for run in runs_by_id.values()} + for run_id in submitted_run_ids: + if run_id not in runs_by_id: + logger.info( + "Ignoring endpoint %s submitted run %s because it is not a live run " + "owned by the endpoint user/project", + endpoint_model.name, + run_id, + ) + continue + try: + await record_endpoint_run_submission( + session=session, + endpoint_id=endpoint_model.id, + run_id=run_id, + ) + except ServerClientError as e: + logger.info( + "Ignoring endpoint %s submitted run %s: %s", + endpoint_model.name, + run_id, + e.msg, + ) + for run_name in valid_submitted_run_names: + run = runs_by_name.get(run_name) + if run is None: + continue + try: + await record_endpoint_run_submission( + session=session, + endpoint_id=endpoint_model.id, + run_id=run.id, + ) + except ServerClientError as e: + logger.info( + "Ignoring endpoint %s submitted run %s: %s", + endpoint_model.name, + run_name, + e.msg, + ) + await session.commit() + + +async def _stop_non_final_submitted_runs( + *, + endpoint_model: EndpointModel, + final_run_id: uuid.UUID, + pipeline_hinter: PipelineHinterProtocol, +) -> None: + for run_model in await _get_endpoint_unfinished_runs(endpoint_model): + if run_model.id == final_run_id or run_model.status == RunStatus.TERMINATING: + continue + logger.info( + "Stopping non-final endpoint run %s after endpoint %s verified run %s", + run_model.run_name, + endpoint_model.name, + final_run_id, + ) + await _stop_backing_run( + endpoint_model=endpoint_model, + run_name=run_model.run_name, + pipeline_hinter=pipeline_hinter, + ) + + +def _is_valid_agent_submission_run_name(endpoint_name: str, run_name: str) -> bool: + prefix = f"{endpoint_name}-" + if not run_name.startswith(prefix): + return False + suffix = run_name[len(prefix) :] + if not suffix.isdecimal(): + return False + submission_num = int(suffix) + return submission_num > 0 and str(submission_num) == suffix + + +async def _process_running_endpoint(endpoint_model: EndpointModel) -> _ProcessResult: + readiness = _get_backing_service_readiness(endpoint_model) + if readiness.failed_message is not None: + return _ProcessResult( + update_map={ + "status": EndpointStatus.FAILED, + "status_message": readiness.failed_message, + } + ) + if readiness.model_base_url is None or readiness.model_name is None: + return _ProcessResult() + return _ProcessResult(update_map={"status_message": None}) + + +async def _process_stopping_endpoint( + endpoint_model: EndpointModel, + pipeline_hinter: PipelineHinterProtocol, +) -> _ProcessResult: + agent_aborted = True + if endpoint_model.provisioning_method == "agent": + agent_aborted = await abort_agent_endpoint(endpoint_model) + run_models = await _get_endpoint_unfinished_runs(endpoint_model) + if not run_models: + if not agent_aborted: + return _ProcessResult() + return _get_stopped_result() + for run_model in run_models: + if run_model.status == RunStatus.TERMINATING: + continue + logger.info( + "Stopping backing run %s before stopping endpoint %s", + run_model.run_name, + endpoint_model.name, + ) + await _stop_backing_run( + endpoint_model=endpoint_model, + run_name=run_model.run_name, + pipeline_hinter=pipeline_hinter, + ) + return _ProcessResult() + + +async def _try_save_agent_endpoint_preset( + endpoint_model: EndpointModel, +) -> None: + run_model = endpoint_model.service_run + if run_model is None: + return + await _save_agent_endpoint_preset( + endpoint_model=endpoint_model, + run_model=run_model, + base_model=endpoint_model.model_base, + recipe_model=endpoint_model.model_repo, + ) + + +async def _save_agent_endpoint_preset( + endpoint_model: EndpointModel, + run_model: RunModel, + base_model: Optional[str] = None, + recipe_model: Optional[str] = None, +) -> None: + try: + preset = build_endpoint_preset_from_run( + run_model, + base_model=base_model, + recipe_model=recipe_model, + ) + saved_preset = await get_endpoint_preset_service().save_preset( + endpoint_model.project.name, + preset, + comments=[ + "Generated by dstack endpoint agent.", + f"endpoint: {endpoint_model.name}", + f"endpoint_id: {endpoint_model.id}", + f"run_id: {run_model.id}", + ], + ) + except Exception as e: + logger.warning( + "Failed to save endpoint preset for endpoint %s: %s", + endpoint_model.name, + e, + exc_info=True, + ) + return + recipe_ids = ", ".join(recipe.id for recipe in saved_preset.recipes) + logger.info( + "Saved endpoint preset for model %s recipes %s for endpoint %s", + saved_preset.base, + recipe_ids, + endpoint_model.name, + ) + + +@dataclass(frozen=True) +class _BackingServiceReadiness: + failed_message: Optional[str] = None + model_base_url: Optional[str] = None + model_name: Optional[str] = None + + +def _get_backing_service_readiness(endpoint_model: EndpointModel) -> _BackingServiceReadiness: + run_model = endpoint_model.service_run + if run_model is None: + return _BackingServiceReadiness(failed_message="Backing service run is missing") + return _get_service_run_readiness(run_model, endpoint_name=endpoint_model.name) + + +def _get_service_run_readiness( + run_model: RunModel, + *, + endpoint_name: str, +) -> _BackingServiceReadiness: + if run_model.deleted: + return _BackingServiceReadiness(failed_message="Backing service run was deleted") + if run_model.status.is_finished(): + return _BackingServiceReadiness( + failed_message=f"Backing service run finished with status {run_model.status.value}" + ) + if run_model.status != RunStatus.RUNNING: + return _BackingServiceReadiness() + if not _has_registered_running_job(run_model): + return _BackingServiceReadiness() + if run_model.service_spec is None: + return _BackingServiceReadiness() + try: + service_spec = ServiceSpec.__response__.parse_raw(run_model.service_spec) + except ValidationError: + logger.warning("Endpoint %s backing service spec is invalid", endpoint_name) + return _BackingServiceReadiness(failed_message="Backing service spec is invalid") + if service_spec.model is None: + return _BackingServiceReadiness() + return _BackingServiceReadiness( + model_base_url=service_spec.model.base_url, + model_name=service_spec.model.name, + ) + + +def _has_registered_running_job(run_model: RunModel) -> bool: + return any(job.status == JobStatus.RUNNING and job.registered for job in run_model.jobs) + + +async def _submit_endpoint_from_preset( + endpoint_id: uuid.UUID, + pipeline_hinter: PipelineHinterProtocol, +) -> _PresetSubmissionResult: + async with get_session_ctx() as session: + res = await session.execute( + select(EndpointModel) + .where( + EndpointModel.id == endpoint_id, + EndpointModel.status == EndpointStatus.SUBMITTED, + ) + .options(joinedload(EndpointModel.project).joinedload(ProjectModel.backends)) + .options(joinedload(EndpointModel.user)) + ) + endpoint_model = res.unique().scalar_one_or_none() + if endpoint_model is None: + return _PresetSubmissionResult() + endpoint_configuration = get_endpoint_configuration(endpoint_model) + if endpoint_configuration.preset_policy == EndpointPresetPolicy.CREATE: + return _PresetSubmissionResult() + preset_planning_result = await find_preset_planning_result( + session=session, + project=endpoint_model.project, + user=endpoint_model.user, + endpoint_name=endpoint_model.name, + endpoint_configuration=endpoint_configuration, + ) + preset_plan = preset_planning_result.provisionable + if preset_plan is None: + unprovisionable_preset = None + if preset_planning_result.unprovisionable is not None: + unprovisionable_preset = _format_preset_plan_label( + preset_planning_result.unprovisionable + ) + return _PresetSubmissionResult(unprovisionable_preset=unprovisionable_preset) + conflict_message = await _get_active_run_name_conflict_message( + session=session, + project=endpoint_model.project, + run_name=preset_plan.run_plan.run_spec.run_name, + linked_run_id=endpoint_model.service_run_id, + ) + if conflict_message is not None: + raise ServerClientError(conflict_message) + run = await runs_services.apply_plan( + session=session, + user=endpoint_model.user, + project=endpoint_model.project, + plan=ApplyRunPlanInput( + run_spec=preset_plan.run_plan.run_spec, + current_resource=preset_plan.run_plan.current_resource, + ), + force=False, + pipeline_hinter=pipeline_hinter, + ) + await record_endpoint_run_submission( + session=session, + endpoint_id=endpoint_model.id, + run_id=run.id, + ) + await session.commit() + return _PresetSubmissionResult( + submission=_PresetSubmission( + run_id=run.id, + preset_base=preset_plan.preset.base, + recipe_id=preset_plan.recipe.id, + model_repo=preset_plan.recipe.model, + ) + ) + + +async def _get_active_run_name_conflict_message( + *, + session, + project: ProjectModel, + run_name: Optional[str], + linked_run_id: Optional[uuid.UUID], +) -> Optional[str]: + if run_name is None: + return None + run_model = await runs_services.get_run_model_by_name( + session=session, + project=project, + run_name=run_name, + ) + if run_model is None: + return None + if linked_run_id == run_model.id: + return None + if run_model.status.is_finished(): + return None + return f"Run name '{run_name}' is taken by an existing run" + + +def _get_no_provisioning_path_message( + endpoint_model: EndpointModel, + unprovisionable_preset: Optional[str] = None, +) -> str: + endpoint_configuration = get_endpoint_configuration(endpoint_model) + if unprovisionable_preset is not None: + reason = f"Endpoint preset {unprovisionable_preset} matched but has no available offers." + if endpoint_configuration.preset_policy == EndpointPresetPolicy.REUSE: + return reason + if get_agent_service().is_enabled() and not can_use_endpoint_agent( + user=endpoint_model.user, + project=endpoint_model.project, + ): + return f"{reason} {get_endpoint_agent_admin_required_message()}" + agent_unavailable_reason = get_agent_unavailable_reason() + if endpoint_configuration.preset_policy == EndpointPresetPolicy.REUSE_OR_CREATE: + return ( + f"{reason} Creating a preset requires the server agent, " + f"but {agent_unavailable_reason}" + ) + if endpoint_configuration.preset_policy == EndpointPresetPolicy.REUSE: + return _NO_MATCHING_PRESET_MESSAGE + if get_agent_service().is_enabled() and not can_use_endpoint_agent( + user=endpoint_model.user, + project=endpoint_model.project, + ): + if endpoint_configuration.preset_policy == EndpointPresetPolicy.REUSE_OR_CREATE: + return ( + "No matching endpoint presets found. " + f"{get_endpoint_agent_admin_required_message()}" + ) + return get_endpoint_agent_admin_required_message() + agent_unavailable_reason = get_agent_unavailable_reason() + if endpoint_configuration.preset_policy == EndpointPresetPolicy.REUSE_OR_CREATE: + return ( + "No matching endpoint presets found. " + f"Creating a preset requires the server agent, but {agent_unavailable_reason}" + ) + return f"Preset policy create requires the server agent, but {agent_unavailable_reason}" + + +def _format_preset_plan_label(preset_plan) -> str: + return f"for model {preset_plan.preset.base} recipe {preset_plan.recipe.id}" + + +async def _stop_backing_run( + endpoint_model: EndpointModel, + run_name: str, + pipeline_hinter: PipelineHinterProtocol, +) -> None: + async with get_session_ctx() as session: + await runs_services.stop_runs( + session=session, + user=endpoint_model.user, + project=endpoint_model.project, + runs_names=[run_name], + abort=False, + pipeline_hinter=pipeline_hinter, + ) + + +def _get_stopped_result() -> _ProcessResult: + return _ProcessResult( + update_map={ + "status": EndpointStatus.STOPPED, + "status_message": None, + } + ) diff --git a/src/dstack/_internal/server/migrations/versions/2026/07_03_1243_03efb71a1563_add_endpoints.py b/src/dstack/_internal/server/migrations/versions/2026/07_03_1243_03efb71a1563_add_endpoints.py new file mode 100644 index 0000000000..ac1444c6f1 --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/2026/07_03_1243_03efb71a1563_add_endpoints.py @@ -0,0 +1,118 @@ +"""add endpoints + +Revision ID: 03efb71a1563 +Revises: e9c5e7e26c78 +Create Date: 2026-07-03 12:43:58.140371+00:00 + +""" + +import sqlalchemy as sa +import sqlalchemy_utils +from alembic import op + +import dstack._internal.server.models + +# revision identifiers, used by Alembic. +revision = "03efb71a1563" +down_revision = "e9c5e7e26c78" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "endpoints", + sa.Column("id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False), + sa.Column("name", sa.String(length=100), nullable=False), + sa.Column( + "project_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False + ), + sa.Column("user_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False), + sa.Column( + "service_run_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=True + ), + sa.Column("configuration", sa.Text(), nullable=False), + sa.Column("status", sa.String(length=100), nullable=False), + sa.Column("status_message", sa.Text(), nullable=True), + sa.Column("provisioning_method", sa.String(length=100), nullable=True), + sa.Column("created_at", dstack._internal.server.models.NaiveDateTime(), nullable=False), + sa.Column( + "last_processed_at", dstack._internal.server.models.NaiveDateTime(), nullable=False + ), + sa.Column( + "lock_expires_at", dstack._internal.server.models.NaiveDateTime(), nullable=True + ), + sa.Column("lock_token", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=True), + sa.Column("lock_owner", sa.String(length=100), nullable=True), + sa.ForeignKeyConstraint( + ["project_id"], + ["projects.id"], + name=op.f("fk_endpoints_project_id_projects"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["service_run_id"], ["runs.id"], name=op.f("fk_endpoints_service_run_id_runs") + ), + sa.ForeignKeyConstraint( + ["user_id"], ["users.id"], name=op.f("fk_endpoints_user_id_users"), ondelete="CASCADE" + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_endpoints")), + sa.UniqueConstraint("project_id", "name", name="uq_endpoints_project_id_name"), + ) + with op.batch_alter_table("endpoints", schema=None) as batch_op: + batch_op.create_index( + "ix_endpoints_pipeline_fetch_q", + [sa.literal_column("last_processed_at ASC")], + unique=False, + ) + batch_op.create_index(batch_op.f("ix_endpoints_status"), ["status"], unique=False) + + op.create_table( + "endpoint_run_submissions", + sa.Column( + "endpoint_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False + ), + sa.Column("run_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False), + sa.Column("submission_num", sa.Integer(), nullable=False), + sa.Column("submitted_at", dstack._internal.server.models.NaiveDateTime(), nullable=False), + sa.ForeignKeyConstraint( + ["endpoint_id"], + ["endpoints.id"], + name=op.f("fk_endpoint_run_submissions_endpoint_id_endpoints"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["run_id"], + ["runs.id"], + name=op.f("fk_endpoint_run_submissions_run_id_runs"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint( + "endpoint_id", "submission_num", name=op.f("pk_endpoint_run_submissions") + ), + sa.UniqueConstraint("run_id", name="uq_endpoint_run_submissions_run_id"), + ) + with op.batch_alter_table("endpoint_run_submissions", schema=None) as batch_op: + batch_op.create_index( + batch_op.f("ix_endpoint_run_submissions_endpoint_id"), ["endpoint_id"], unique=False + ) + + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("endpoint_run_submissions", schema=None) as batch_op: + batch_op.drop_index(batch_op.f("ix_endpoint_run_submissions_endpoint_id")) + + op.drop_table("endpoint_run_submissions") + + with op.batch_alter_table("endpoints", schema=None) as batch_op: + batch_op.drop_index(batch_op.f("ix_endpoints_status")) + batch_op.drop_index( + "ix_endpoints_pipeline_fetch_q", + ) + + op.drop_table("endpoints") + # ### end Alembic commands ### diff --git a/src/dstack/_internal/server/migrations/versions/2026/07_06_1700_4e1d6c2a9b77_add_endpoint_agent_sessions.py b/src/dstack/_internal/server/migrations/versions/2026/07_06_1700_4e1d6c2a9b77_add_endpoint_agent_sessions.py new file mode 100644 index 0000000000..fee7c5799e --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/2026/07_06_1700_4e1d6c2a9b77_add_endpoint_agent_sessions.py @@ -0,0 +1,75 @@ +"""add endpoint agent sessions + +Revision ID: 4e1d6c2a9b77 +Revises: 03efb71a1563 +Create Date: 2026-07-06 17:00:00.000000+00:00 + +""" + +import sqlalchemy as sa +import sqlalchemy_utils +from alembic import op + +import dstack._internal.server.models + +# revision identifiers, used by Alembic. +revision = "4e1d6c2a9b77" +down_revision = "03efb71a1563" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "endpoint_agent_sessions", + sa.Column( + "endpoint_id", sqlalchemy_utils.types.uuid.UUIDType(binary=False), nullable=False + ), + sa.Column("session_num", sa.Integer(), nullable=False), + sa.Column("status", sa.String(length=100), nullable=False), + sa.Column("workspace_path", sa.Text(), nullable=False), + sa.Column("pid", sa.Integer(), nullable=True), + sa.Column("process_host", sa.String(length=255), nullable=True), + sa.Column("progress_log_offset", sa.BigInteger(), nullable=False), + sa.Column("stdout_log_offset", sa.BigInteger(), nullable=False), + sa.Column("stderr_log_offset", sa.BigInteger(), nullable=False), + sa.Column("status_message", sa.Text(), nullable=True), + sa.Column("created_at", dstack._internal.server.models.NaiveDateTime(), nullable=False), + sa.Column("updated_at", dstack._internal.server.models.NaiveDateTime(), nullable=False), + sa.Column("finished_at", dstack._internal.server.models.NaiveDateTime(), nullable=True), + sa.ForeignKeyConstraint( + ["endpoint_id"], + ["endpoints.id"], + name=op.f("fk_endpoint_agent_sessions_endpoint_id_endpoints"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint( + "endpoint_id", "session_num", name=op.f("pk_endpoint_agent_sessions") + ), + ) + with op.batch_alter_table("endpoint_agent_sessions", schema=None) as batch_op: + batch_op.create_index( + batch_op.f("ix_endpoint_agent_sessions_endpoint_id"), ["endpoint_id"], unique=False + ) + batch_op.create_index( + "ix_endpoint_agent_sessions_endpoint_status", + ["endpoint_id", "status"], + unique=False, + ) + batch_op.create_index( + batch_op.f("ix_endpoint_agent_sessions_status"), ["status"], unique=False + ) + + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("endpoint_agent_sessions", schema=None) as batch_op: + batch_op.drop_index(batch_op.f("ix_endpoint_agent_sessions_status")) + batch_op.drop_index("ix_endpoint_agent_sessions_endpoint_status") + batch_op.drop_index(batch_op.f("ix_endpoint_agent_sessions_endpoint_id")) + + op.drop_table("endpoint_agent_sessions") + # ### end Alembic commands ### diff --git a/src/dstack/_internal/server/migrations/versions/2026/07_10_1200_e03d97df7c5a_add_endpoint_model_repo.py b/src/dstack/_internal/server/migrations/versions/2026/07_10_1200_e03d97df7c5a_add_endpoint_model_repo.py new file mode 100644 index 0000000000..35e36d7e26 --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/2026/07_10_1200_e03d97df7c5a_add_endpoint_model_repo.py @@ -0,0 +1,28 @@ +"""add endpoint model fields + +Revision ID: e03d97df7c5a +Revises: 4e1d6c2a9b77 +Create Date: 2026-07-10 12:00:00.000000+00:00 + +""" + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "e03d97df7c5a" +down_revision = "4e1d6c2a9b77" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + with op.batch_alter_table("endpoints", schema=None) as batch_op: + batch_op.add_column(sa.Column("model_base", sa.Text(), nullable=True)) + batch_op.add_column(sa.Column("model_repo", sa.Text(), nullable=True)) + + +def downgrade() -> None: + with op.batch_alter_table("endpoints", schema=None) as batch_op: + batch_op.drop_column("model_repo") + batch_op.drop_column("model_base") diff --git a/src/dstack/_internal/server/models.py b/src/dstack/_internal/server/models.py index fc73010263..6c403b15f0 100644 --- a/src/dstack/_internal/server/models.py +++ b/src/dstack/_internal/server/models.py @@ -26,6 +26,7 @@ from dstack._internal.core.models.backends.base import BackendType from dstack._internal.core.models.common import CoreConfig, generate_dual_core_model from dstack._internal.core.models.compute_groups import ComputeGroupStatus +from dstack._internal.core.models.endpoints import EndpointStatus from dstack._internal.core.models.events import EventTargetType from dstack._internal.core.models.fleets import FleetStatus from dstack._internal.core.models.gateways import GatewayReplicaStatus, GatewayStatus @@ -54,6 +55,12 @@ CASCADE_DEFAULT_WITH_DELETE_ORPHAN = "save-update, merge, delete-orphan, delete" +class EndpointAgentSessionStatus(enum.Enum): + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + + class NaiveDateTime(TypeDecorator): """ A custom type decorator that ensures datetime objects are offset-naive when stored in the database @@ -998,6 +1005,103 @@ class VolumeModel(PipelineModelMixin, BaseModel): ) +class EndpointModel(PipelineModelMixin, BaseModel): + __tablename__ = "endpoints" + + id: Mapped[uuid.UUID] = mapped_column( + UUIDType(binary=False), primary_key=True, default=uuid.uuid4 + ) + name: Mapped[str] = mapped_column(String(100)) + + project_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("projects.id", ondelete="CASCADE")) + project: Mapped["ProjectModel"] = relationship(foreign_keys=[project_id]) + + user_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("users.id", ondelete="CASCADE")) + user: Mapped["UserModel"] = relationship() + + service_run_id: Mapped[Optional[uuid.UUID]] = mapped_column(ForeignKey("runs.id")) + service_run: Mapped[Optional["RunModel"]] = relationship() + """The service run currently backing this endpoint. + + This is the current/latest run. Historical endpoint-submitted runs are + recorded in EndpointRunSubmissionModel. + """ + + configuration: Mapped[str] = mapped_column(Text) + status: Mapped[EndpointStatus] = mapped_column(EnumAsString(EndpointStatus, 100), index=True) + """`status` must be changed only via `switch_endpoint_status()`.""" + status_message: Mapped[Optional[str]] = mapped_column(Text) + model_base: Mapped[Optional[str]] = mapped_column(Text) + model_repo: Mapped[Optional[str]] = mapped_column(Text) + provisioning_method: Mapped[Optional[str]] = mapped_column(String(100)) + + created_at: Mapped[datetime] = mapped_column(NaiveDateTime, default=get_current_datetime) + last_processed_at: Mapped[datetime] = mapped_column( + NaiveDateTime, default=get_current_datetime + ) + + __table_args__ = ( + UniqueConstraint("project_id", "name", name="uq_endpoints_project_id_name"), + Index( + "ix_endpoints_pipeline_fetch_q", + last_processed_at.asc(), + ), + ) + + +class EndpointRunSubmissionModel(BaseModel): + __tablename__ = "endpoint_run_submissions" + + endpoint_id: Mapped[uuid.UUID] = mapped_column( + ForeignKey("endpoints.id", ondelete="CASCADE"), primary_key=True + ) + endpoint: Mapped["EndpointModel"] = relationship() + + run_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("runs.id", ondelete="CASCADE")) + run: Mapped["RunModel"] = relationship() + + submission_num: Mapped[int] = mapped_column(Integer, primary_key=True) + submitted_at: Mapped[datetime] = mapped_column(NaiveDateTime) + + __table_args__ = ( + UniqueConstraint("run_id", name="uq_endpoint_run_submissions_run_id"), + Index("ix_endpoint_run_submissions_endpoint_id", endpoint_id), + ) + + +class EndpointAgentSessionModel(BaseModel): + __tablename__ = "endpoint_agent_sessions" + + endpoint_id: Mapped[uuid.UUID] = mapped_column( + ForeignKey("endpoints.id", ondelete="CASCADE"), primary_key=True + ) + endpoint: Mapped["EndpointModel"] = relationship() + + session_num: Mapped[int] = mapped_column(Integer, primary_key=True) + status: Mapped[EndpointAgentSessionStatus] = mapped_column( + EnumAsString(EndpointAgentSessionStatus, 100), index=True + ) + workspace_path: Mapped[str] = mapped_column(Text) + + pid: Mapped[Optional[int]] = mapped_column(Integer) + process_host: Mapped[Optional[str]] = mapped_column(String(255)) + + progress_log_offset: Mapped[int] = mapped_column(BigInteger, default=0) + stdout_log_offset: Mapped[int] = mapped_column(BigInteger, default=0) + stderr_log_offset: Mapped[int] = mapped_column(BigInteger, default=0) + + status_message: Mapped[Optional[str]] = mapped_column(Text) + + created_at: Mapped[datetime] = mapped_column(NaiveDateTime, default=get_current_datetime) + updated_at: Mapped[datetime] = mapped_column(NaiveDateTime, default=get_current_datetime) + finished_at: Mapped[Optional[datetime]] = mapped_column(NaiveDateTime) + + __table_args__ = ( + Index("ix_endpoint_agent_sessions_endpoint_id", endpoint_id), + Index("ix_endpoint_agent_sessions_endpoint_status", endpoint_id, status), + ) + + class VolumeAttachmentModel(BaseModel): __tablename__ = "volumes_attachments" diff --git a/src/dstack/_internal/server/routers/endpoints.py b/src/dstack/_internal/server/routers/endpoints.py new file mode 100644 index 0000000000..c076b90100 --- /dev/null +++ b/src/dstack/_internal/server/routers/endpoints.py @@ -0,0 +1,182 @@ +from typing import List, Tuple + +from fastapi import APIRouter, Depends +from sqlalchemy.ext.asyncio import AsyncSession + +import dstack._internal.server.services.endpoints as endpoints_services +from dstack._internal.core.errors import ResourceNotExistsError +from dstack._internal.core.models.endpoint_presets import EndpointPreset +from dstack._internal.core.models.endpoints import Endpoint, EndpointPlan +from dstack._internal.server.db import get_session +from dstack._internal.server.models import ProjectModel, UserModel +from dstack._internal.server.schemas.endpoint_presets import ( + DeleteEndpointPresetsRequest, + GetEndpointPresetRequest, +) +from dstack._internal.server.schemas.endpoints import ( + CreateEndpointRequest, + GetEndpointPlanRequest, + GetEndpointRequest, + ListEndpointsRequest, + StopEndpointsRequest, +) +from dstack._internal.server.security.permissions import Authenticated, ProjectMember +from dstack._internal.server.services.endpoints.presets import ( + endpoint_preset_to_api_details, + endpoint_preset_to_api_model, + get_endpoint_preset_service, +) +from dstack._internal.server.services.pipelines import PipelineHinterProtocol, get_pipeline_hinter +from dstack._internal.server.utils.routers import ( + CustomORJSONResponse, + get_base_api_additional_responses, +) + +root_router = APIRouter( + prefix="/api/endpoints", + tags=["endpoints"], + responses=get_base_api_additional_responses(), +) +project_router = APIRouter(prefix="/api/project/{project_name}/endpoints", tags=["endpoints"]) + + +@root_router.post("/list", summary="List endpoints", response_model=List[Endpoint]) +async def list_endpoints( + body: ListEndpointsRequest, + session: AsyncSession = Depends(get_session), + user: UserModel = Depends(Authenticated()), +): + return CustomORJSONResponse( + await endpoints_services.list_endpoints( + session=session, + user=user, + project_name=body.project_name, + only_active=body.only_active, + prev_created_at=body.prev_created_at, + prev_id=body.prev_id, + limit=body.limit, + ascending=body.ascending, + ) + ) + + +@project_router.post("/list", summary="List project endpoints", response_model=List[Endpoint]) +async def list_project_endpoints( + session: AsyncSession = Depends(get_session), + user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()), +): + _, project = user_project + return CustomORJSONResponse( + await endpoints_services.list_project_endpoints( + session=session, + project=project, + ) + ) + + +@project_router.post("/get", summary="Get endpoint", response_model=Endpoint) +async def get_endpoint( + body: GetEndpointRequest, + session: AsyncSession = Depends(get_session), + user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()), +): + _, project = user_project + endpoint = await endpoints_services.get_endpoint_by_name( + session=session, + project=project, + name=body.name, + ) + if endpoint is None: + raise ResourceNotExistsError() + return CustomORJSONResponse(endpoint) + + +@project_router.post("/get_plan", summary="Get endpoint apply plan", response_model=EndpointPlan) +async def get_endpoint_plan( + body: GetEndpointPlanRequest, + session: AsyncSession = Depends(get_session), + user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()), +): + user, project = user_project + return CustomORJSONResponse( + await endpoints_services.get_endpoint_plan( + session=session, + project=project, + user=user, + configuration=body.configuration, + configuration_path=body.configuration_path, + ) + ) + + +@project_router.post("/create", summary="Create endpoint", response_model=Endpoint) +async def create_endpoint( + body: CreateEndpointRequest, + session: AsyncSession = Depends(get_session), + user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()), + pipeline_hinter: PipelineHinterProtocol = Depends(get_pipeline_hinter), +): + user, project = user_project + return CustomORJSONResponse( + await endpoints_services.create_endpoint( + session=session, + project=project, + user=user, + configuration=body.configuration, + pipeline_hinter=pipeline_hinter, + ) + ) + + +@project_router.post("/stop", summary="Stop endpoints") +async def stop_endpoints( + body: StopEndpointsRequest, + session: AsyncSession = Depends(get_session), + user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()), + pipeline_hinter: PipelineHinterProtocol = Depends(get_pipeline_hinter), +): + user, project = user_project + await endpoints_services.stop_endpoints( + session=session, + project=project, + names=body.names, + user=user, + pipeline_hinter=pipeline_hinter, + ) + + +@project_router.post( + "/presets/list", summary="List endpoint presets", response_model=List[EndpointPreset] +) +async def list_endpoint_presets( + user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()), +): + _, project = user_project + presets = await get_endpoint_preset_service().list_presets(project.name) + return CustomORJSONResponse([endpoint_preset_to_api_model(preset) for preset in presets]) + + +@project_router.post("/presets/get", summary="Get endpoint preset", response_model=EndpointPreset) +async def get_endpoint_preset( + body: GetEndpointPresetRequest, + user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()), +): + _, project = user_project + preset = await get_endpoint_preset_service().get_preset(project.name, body.model) + if preset is None: + raise ResourceNotExistsError() + return CustomORJSONResponse(endpoint_preset_to_api_details(preset)) + + +@project_router.post("/presets/delete", summary="Delete endpoint presets") +async def delete_endpoint_presets( + body: DeleteEndpointPresetsRequest, + user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()), +): + _, project = user_project + preset_service = get_endpoint_preset_service() + try: + for model in body.models: + await preset_service.delete_preset(project.name, model) + except FileNotFoundError: + raise ResourceNotExistsError() diff --git a/src/dstack/_internal/server/schemas/endpoint_presets.py b/src/dstack/_internal/server/schemas/endpoint_presets.py new file mode 100644 index 0000000000..91de176094 --- /dev/null +++ b/src/dstack/_internal/server/schemas/endpoint_presets.py @@ -0,0 +1,11 @@ +from typing import List + +from dstack._internal.core.models.common import CoreModel + + +class GetEndpointPresetRequest(CoreModel): + model: str + + +class DeleteEndpointPresetsRequest(CoreModel): + models: List[str] diff --git a/src/dstack/_internal/server/schemas/endpoints.py b/src/dstack/_internal/server/schemas/endpoints.py new file mode 100644 index 0000000000..0cc847d4ae --- /dev/null +++ b/src/dstack/_internal/server/schemas/endpoints.py @@ -0,0 +1,34 @@ +from datetime import datetime +from typing import List, Optional +from uuid import UUID + +from pydantic import Field + +from dstack._internal.core.models.common import CoreModel +from dstack._internal.core.models.endpoints import EndpointConfiguration + + +class ListEndpointsRequest(CoreModel): + project_name: Optional[str] + only_active: bool = False + prev_created_at: Optional[datetime] + prev_id: Optional[UUID] + limit: int = Field(100, ge=0, le=100) + ascending: bool = False + + +class GetEndpointRequest(CoreModel): + name: str + + +class GetEndpointPlanRequest(CoreModel): + configuration: EndpointConfiguration + configuration_path: Optional[str] = None + + +class CreateEndpointRequest(CoreModel): + configuration: EndpointConfiguration + + +class StopEndpointsRequest(CoreModel): + names: List[str] diff --git a/src/dstack/_internal/server/services/endpoints/__init__.py b/src/dstack/_internal/server/services/endpoints/__init__.py new file mode 100644 index 0000000000..6ff1912646 --- /dev/null +++ b/src/dstack/_internal/server/services/endpoints/__init__.py @@ -0,0 +1,766 @@ +import uuid +from datetime import datetime +from typing import Any, List, Optional + +from sqlalchemy import and_, exists, func, or_, select +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm import joinedload + +from dstack._internal.core.errors import ForbiddenError, ResourceExistsError, ServerClientError +from dstack._internal.core.models.common import ApplyAction, EntityReference +from dstack._internal.core.models.endpoints import ( + Endpoint, + EndpointConfiguration, + EndpointPlan, + EndpointPlanJobOffers, + EndpointPresetPolicy, + EndpointProvisioningPlanAgent, + EndpointProvisioningPlanNone, + EndpointProvisioningPlanPreset, + EndpointStatus, +) +from dstack._internal.core.models.fleets import FleetStatus +from dstack._internal.core.models.runs import ServiceSpec +from dstack._internal.core.models.users import GlobalRole, ProjectRole +from dstack._internal.core.services import validate_dstack_resource_name +from dstack._internal.server.db import get_db, is_db_postgres, is_db_sqlite +from dstack._internal.server.models import ( + EndpointModel, + EndpointRunSubmissionModel, + ExportedFleetModel, + FleetModel, + ImportModel, + ProjectModel, + UserModel, +) +from dstack._internal.server.services import events +from dstack._internal.server.services.endpoints.agent import ( + AgentPlan, + get_agent_service, + get_agent_unavailable_reason, +) +from dstack._internal.server.services.endpoints.planning import ( + EndpointPresetPlan, + find_preset_planning_result, +) +from dstack._internal.server.services.locking import get_locker, string_to_lock_id +from dstack._internal.server.services.pipelines import PipelineHinterProtocol +from dstack._internal.server.services.projects import ( + get_user_project_role, + list_user_project_models, +) +from dstack._internal.utils import common, random_names +from dstack._internal.utils.logging import get_logger + +logger = get_logger(__name__) + +_ENDPOINT_AGENT_ADMIN_REQUIRED_MESSAGE = ( + "Creating endpoint presets with the server agent requires project admin permissions." +) +_ENDPOINT_NO_FLEETS_MESSAGE = ( + "The project has no fleets. Create one before submitting an endpoint." +) +_ENDPOINT_NO_MATCHING_FLEETS_MESSAGE = ( + "No fleets match the endpoint configuration. Create a fleet or update `fleets` before " + "submitting an endpoint." +) + + +def switch_endpoint_status( + session: AsyncSession, + endpoint_model: EndpointModel, + new_status: EndpointStatus, + actor: events.AnyActor = events.SystemActor(), +): + old_status = endpoint_model.status + if old_status == new_status: + return + + endpoint_model.status = new_status + emit_endpoint_status_change_event( + session=session, + endpoint_model=endpoint_model, + old_status=old_status, + new_status=new_status, + status_message=endpoint_model.status_message, + actor=actor, + ) + + +def emit_endpoint_status_change_event( + session: AsyncSession, + endpoint_model: EndpointModel, + old_status: EndpointStatus, + new_status: EndpointStatus, + status_message: Optional[str], + actor: events.AnyActor = events.SystemActor(), +) -> None: + if old_status == new_status: + return + msg = get_endpoint_status_change_message( + old_status=old_status, + new_status=new_status, + status_message=status_message, + ) + events.emit(session, msg, actor=actor, targets=[events.Target.from_model(endpoint_model)]) + + +def get_endpoint_status_change_message( + old_status: EndpointStatus, + new_status: EndpointStatus, + status_message: Optional[str], +) -> str: + msg = f"Endpoint status changed {old_status.upper()} -> {new_status.upper()}" + if status_message is not None: + msg += f" ({status_message})" + return msg + + +async def list_endpoints( + session: AsyncSession, + user: UserModel, + project_name: Optional[str], + only_active: bool, + prev_created_at: Optional[datetime], + prev_id: Optional[uuid.UUID], + limit: int, + ascending: bool, +) -> List[Endpoint]: + projects = await list_user_project_models( + session=session, + user=user, + only_names=True, + ) + if project_name is not None: + projects = [p for p in projects if p.name == project_name] + endpoint_models = await list_projects_endpoint_models( + session=session, + projects=projects, + only_active=only_active, + prev_created_at=prev_created_at, + prev_id=prev_id, + limit=limit, + ascending=ascending, + ) + return [endpoint_model_to_endpoint(e) for e in endpoint_models] + + +async def list_projects_endpoint_models( + session: AsyncSession, + projects: List[ProjectModel], + only_active: bool, + prev_created_at: Optional[datetime], + prev_id: Optional[uuid.UUID], + limit: int, + ascending: bool, +) -> List[EndpointModel]: + if not projects: + return [] + filters: list[Any] = [EndpointModel.project_id.in_(p.id for p in projects)] + if only_active: + filters.append(EndpointModel.status.not_in(EndpointStatus.finished_statuses())) + if prev_created_at is not None: + if ascending: + if prev_id is None: + filters.append(EndpointModel.created_at > prev_created_at) + else: + filters.append( + or_( + EndpointModel.created_at > prev_created_at, + and_( + EndpointModel.created_at == prev_created_at, + EndpointModel.id < prev_id, + ), + ) + ) + else: + if prev_id is None: + filters.append(EndpointModel.created_at < prev_created_at) + else: + filters.append( + or_( + EndpointModel.created_at < prev_created_at, + and_( + EndpointModel.created_at == prev_created_at, + EndpointModel.id > prev_id, + ), + ) + ) + order_by = (EndpointModel.created_at.desc(), EndpointModel.id) + if ascending: + order_by = (EndpointModel.created_at.asc(), EndpointModel.id.desc()) + res = await session.execute( + select(EndpointModel) + .where(*filters) + .order_by(*order_by) + .limit(limit) + .options(joinedload(EndpointModel.user)) + .options(joinedload(EndpointModel.project)) + .options(joinedload(EndpointModel.service_run)) + ) + return list(res.unique().scalars().all()) + + +async def list_project_endpoints( + session: AsyncSession, + project: ProjectModel, + names: Optional[List[str]] = None, +) -> List[Endpoint]: + endpoint_models = await list_project_endpoint_models( + session=session, + project=project, + names=names, + ) + return [endpoint_model_to_endpoint(e) for e in endpoint_models] + + +async def list_project_endpoint_models( + session: AsyncSession, + project: ProjectModel, + names: Optional[List[str]] = None, +) -> List[EndpointModel]: + filters = [EndpointModel.project_id == project.id] + if names is not None: + filters.append(EndpointModel.name.in_(names)) + res = await session.execute( + select(EndpointModel) + .where(*filters) + .options(joinedload(EndpointModel.user)) + .options(joinedload(EndpointModel.project)) + .options(joinedload(EndpointModel.service_run)) + ) + return list(res.unique().scalars().all()) + + +async def get_endpoint_by_name( + session: AsyncSession, + project: ProjectModel, + name: str, +) -> Optional[Endpoint]: + endpoint_model = await get_project_endpoint_model_by_name( + session=session, + project=project, + name=name, + ) + if endpoint_model is None: + return None + return endpoint_model_to_endpoint(endpoint_model) + + +async def get_project_endpoint_model_by_name( + session: AsyncSession, + project: ProjectModel, + name: str, +) -> Optional[EndpointModel]: + filters = [ + EndpointModel.name == name, + EndpointModel.project_id == project.id, + ] + res = await session.execute( + select(EndpointModel) + .where(*filters) + .options(joinedload(EndpointModel.user)) + .options(joinedload(EndpointModel.project)) + .options(joinedload(EndpointModel.service_run)) + ) + return res.unique().scalar_one_or_none() + + +async def get_endpoint_plan( + session: AsyncSession, + project: ProjectModel, + user: UserModel, + configuration: EndpointConfiguration, + configuration_path: Optional[str], +) -> EndpointPlan: + current_resource = None + if configuration.name is not None: + current_resource = await get_endpoint_by_name( + session=session, + project=project, + name=configuration.name, + ) + if current_resource is not None and current_resource.status.is_finished(): + current_resource = None + has_usable_fleets = await has_endpoint_existing_usable_fleets( + session=session, + project=project, + configuration=configuration, + ) + if not has_usable_fleets: + return EndpointPlan( + project_name=project.name, + user=user.name, + configuration=configuration, + configuration_path=configuration_path, + current_resource=current_resource, + action=ApplyAction.CREATE if current_resource is None else ApplyAction.UPDATE, + preset_policy=configuration.preset_policy, + provisioning_plan=EndpointProvisioningPlanNone( + reason=get_endpoint_no_fleets_message(configuration) + ), + ) + preset_plan = None + unprovisionable_preset_plan = None + if configuration.preset_policy != EndpointPresetPolicy.CREATE: + preset_planning_result = await find_preset_planning_result( + session=session, + project=project, + user=user, + endpoint_name=configuration.name, + endpoint_configuration=configuration, + ) + preset_plan = preset_planning_result.provisionable + unprovisionable_preset_plan = preset_planning_result.unprovisionable + if preset_plan is None and configuration.preset_policy == EndpointPresetPolicy.REUSE: + preset_plan = unprovisionable_preset_plan + provisioning_plan = _get_no_provisioning_plan( + configuration.preset_policy, + unprovisionable_preset_plan=unprovisionable_preset_plan, + ) + if preset_plan is not None: + provisioning_plan = _endpoint_preset_plan_to_provisioning_plan(preset_plan) + elif ( + configuration.preset_policy != EndpointPresetPolicy.REUSE + and get_agent_service().is_enabled() + ): + if can_use_endpoint_agent(user=user, project=project): + provisioning_plan = _agent_plan_to_provisioning_plan( + get_agent_service().get_plan(), + reason=_get_unprovisionable_preset_reason(unprovisionable_preset_plan), + ) + else: + provisioning_plan = _get_agent_forbidden_provisioning_plan( + configuration.preset_policy, + unprovisionable_preset_plan=unprovisionable_preset_plan, + ) + return EndpointPlan( + project_name=project.name, + user=user.name, + configuration=configuration, + configuration_path=configuration_path, + current_resource=current_resource, + action=ApplyAction.CREATE if current_resource is None else ApplyAction.UPDATE, + preset_policy=configuration.preset_policy, + provisioning_plan=provisioning_plan, + ) + + +def can_use_endpoint_agent(user: UserModel, project: ProjectModel) -> bool: + if user.global_role == GlobalRole.ADMIN: + return True + return get_user_project_role(user=user, project=project) == ProjectRole.ADMIN + + +def get_endpoint_agent_admin_required_message() -> str: + return _ENDPOINT_AGENT_ADMIN_REQUIRED_MESSAGE + + +def get_endpoint_no_fleets_message(configuration: EndpointConfiguration) -> str: + if configuration.fleets: + return _ENDPOINT_NO_MATCHING_FLEETS_MESSAGE + return _ENDPOINT_NO_FLEETS_MESSAGE + + +async def has_endpoint_existing_usable_fleets( + *, + session: AsyncSession, + project: ProjectModel, + configuration: EndpointConfiguration, +) -> bool: + filters = [ + FleetModel.deleted == False, + FleetModel.status == FleetStatus.ACTIVE, + _get_endpoint_usable_fleet_ownership_filter(project), + ] + if configuration.fleets is not None: + fleet_filter = _get_endpoint_configured_fleets_filter(configuration.fleets, project) + if fleet_filter is None: + return False + filters.append(fleet_filter) + res = await session.execute( + select(FleetModel.id).join(FleetModel.project).where(*filters).limit(1) + ) + return res.scalar_one_or_none() is not None + + +def _get_endpoint_usable_fleet_ownership_filter(project: ProjectModel): + is_fleet_imported_subquery = exists().where( + ImportModel.project_id == project.id, + ImportModel.export_id == ExportedFleetModel.export_id, + ExportedFleetModel.fleet_id == FleetModel.id, + ) + return or_(FleetModel.project_id == project.id, is_fleet_imported_subquery) + + +def _get_endpoint_configured_fleets_filter(fleets, project: ProjectModel): + conditions = [] + for ref in map(EntityReference.parse, fleets): + if ref.project is None: + conditions.append( + and_( + FleetModel.name == ref.name, + FleetModel.project_id == project.id, + ) + ) + else: + conditions.append( + and_( + FleetModel.name == ref.name, + ProjectModel.name == ref.project, + ) + ) + if not conditions: + return None + return or_(*conditions) + + +def _agent_plan_to_provisioning_plan( + agent_plan: AgentPlan, + reason: Optional[str] = None, +) -> EndpointProvisioningPlanAgent: + return EndpointProvisioningPlanAgent( + agent_model=agent_plan.model, + reason=reason, + ) + + +def _get_no_provisioning_plan( + preset_policy: EndpointPresetPolicy, + unprovisionable_preset_plan: Optional[EndpointPresetPlan] = None, +) -> EndpointProvisioningPlanNone: + if preset_policy == EndpointPresetPolicy.REUSE: + reason = _get_unprovisionable_preset_reason(unprovisionable_preset_plan) + if reason is None: + reason = "No matching endpoint presets found." + return EndpointProvisioningPlanNone(reason=reason) + agent_unavailable_reason = get_agent_unavailable_reason() + preset_reason = _get_unprovisionable_preset_reason(unprovisionable_preset_plan) + if preset_policy == EndpointPresetPolicy.REUSE_OR_CREATE: + if preset_reason is not None: + return EndpointProvisioningPlanNone( + reason=( + f"{preset_reason} Creating a preset requires the server agent, " + f"but {agent_unavailable_reason}" + ) + ) + return EndpointProvisioningPlanNone( + reason=( + "No matching endpoint presets found. Creating a preset requires the server " + f"agent, but {agent_unavailable_reason}" + ) + ) + return EndpointProvisioningPlanNone( + reason=(f"Preset policy create requires the server agent, but {agent_unavailable_reason}") + ) + + +def _get_agent_forbidden_provisioning_plan( + preset_policy: EndpointPresetPolicy, + unprovisionable_preset_plan: Optional[EndpointPresetPlan] = None, +) -> EndpointProvisioningPlanNone: + preset_reason = _get_unprovisionable_preset_reason(unprovisionable_preset_plan) + if preset_reason is not None: + return EndpointProvisioningPlanNone( + reason=f"{preset_reason} {_ENDPOINT_AGENT_ADMIN_REQUIRED_MESSAGE}" + ) + if preset_policy == EndpointPresetPolicy.REUSE_OR_CREATE: + return EndpointProvisioningPlanNone( + reason=( + f"No matching endpoint presets found. {_ENDPOINT_AGENT_ADMIN_REQUIRED_MESSAGE}" + ) + ) + return EndpointProvisioningPlanNone(reason=_ENDPOINT_AGENT_ADMIN_REQUIRED_MESSAGE) + + +def _get_unprovisionable_preset_reason( + preset_plan: Optional[EndpointPresetPlan], +) -> Optional[str]: + if preset_plan is None: + return None + return ( + f"Endpoint preset for model {preset_plan.preset.base} " + f"recipe {preset_plan.recipe.id} matched but has no available offers." + ) + + +def _endpoint_preset_plan_to_provisioning_plan( + preset_plan: EndpointPresetPlan, +) -> EndpointProvisioningPlanPreset: + run_spec = preset_plan.run_plan.get_effective_run_spec() + service_name = run_spec.run_name or run_spec.configuration.name or "(generated)" + return EndpointProvisioningPlanPreset( + preset_base=preset_plan.preset.base, + recipe_id=preset_plan.recipe.id, + service_name=service_name, + job_offers=[ + EndpointPlanJobOffers( + replica_group=job_plan.job_spec.replica_group, + resources=job_plan.job_spec.requirements.resources, + spot=job_plan.job_spec.requirements.spot, + max_price=job_plan.job_spec.requirements.max_price, + offers=job_plan.offers, + total_offers=job_plan.total_offers, + max_offer_price=job_plan.max_price, + ) + for job_plan in preset_plan.run_plan.job_plans + ], + ) + + +async def create_endpoint( + session: AsyncSession, + project: ProjectModel, + user: UserModel, + configuration: EndpointConfiguration, + pipeline_hinter: PipelineHinterProtocol, +) -> Endpoint: + _validate_endpoint_configuration(configuration) + + lock_namespace = f"endpoint_names_{project.name}" + if is_db_sqlite(): + await session.commit() + elif is_db_postgres(): + await session.execute( + select(func.pg_advisory_xact_lock(string_to_lock_id(lock_namespace))) + ) + lock, _ = get_locker(get_db().dialect_name).get_lockset(lock_namespace) + async with lock: + now = common.get_current_datetime() + if configuration.name is not None: + endpoint_model = await get_project_endpoint_model_by_name( + session=session, + project=project, + name=configuration.name, + ) + if endpoint_model is not None: + if not endpoint_model.status.is_finished(): + raise ResourceExistsError() + await _check_can_submit_endpoint_configuration( + session=session, + project=project, + user=user, + configuration=configuration, + ) + endpoint_model.user = user + endpoint_model.service_run_id = None + endpoint_model.service_run = None + endpoint_model.model_base = None + endpoint_model.model_repo = None + endpoint_model.configuration = configuration.json() + endpoint_model.status = EndpointStatus.SUBMITTED + endpoint_model.status_message = None + endpoint_model.provisioning_method = None + endpoint_model.created_at = now + endpoint_model.last_processed_at = now + endpoint_model.lock_expires_at = None + endpoint_model.lock_token = None + endpoint_model.lock_owner = None + events.emit( + session, + message=f"Endpoint submitted. Status: {endpoint_model.status.upper()}", + actor=events.UserActor.from_user(user), + targets=[events.Target.from_model(endpoint_model)], + ) + await session.commit() + pipeline_hinter.hint_fetch(EndpointModel.__name__) + return endpoint_model_to_endpoint(endpoint_model) + else: + configuration.name = await generate_endpoint_name(session=session, project=project) + + await _check_can_submit_endpoint_configuration( + session=session, + project=project, + user=user, + configuration=configuration, + ) + + endpoint_model = EndpointModel( + id=uuid.uuid4(), + name=configuration.name, + project=project, + user_id=user.id, + status=EndpointStatus.SUBMITTED, + configuration=configuration.json(), + created_at=now, + last_processed_at=now, + ) + session.add(endpoint_model) + events.emit( + session, + message=f"Endpoint created. Status: {endpoint_model.status.upper()}", + actor=events.UserActor.from_user(user), + targets=[events.Target.from_model(endpoint_model)], + ) + await session.commit() + pipeline_hinter.hint_fetch(EndpointModel.__name__) + return endpoint_model_to_endpoint(endpoint_model) + + +async def _check_can_submit_endpoint_configuration( + session: AsyncSession, + project: ProjectModel, + user: UserModel, + configuration: EndpointConfiguration, +) -> None: + if not await has_endpoint_existing_usable_fleets( + session=session, + project=project, + configuration=configuration, + ): + raise ServerClientError(get_endpoint_no_fleets_message(configuration)) + if configuration.preset_policy == EndpointPresetPolicy.REUSE: + return + preset_planning_result = None + if configuration.preset_policy == EndpointPresetPolicy.REUSE_OR_CREATE: + preset_planning_result = await find_preset_planning_result( + session=session, + project=project, + user=user, + endpoint_name=configuration.name, + endpoint_configuration=configuration, + ) + if preset_planning_result.provisionable is not None: + return + if can_use_endpoint_agent(user=user, project=project): + return + raise ForbiddenError(_ENDPOINT_AGENT_ADMIN_REQUIRED_MESSAGE) + + +async def stop_endpoints( + session: AsyncSession, + project: ProjectModel, + names: List[str], + user: UserModel, + pipeline_hinter: Optional[PipelineHinterProtocol] = None, +): + endpoint_models = await list_project_endpoint_models( + session=session, + project=project, + names=names, + ) + for endpoint_model in endpoint_models: + if endpoint_model.status.is_finished(): + continue + endpoint_model.status = EndpointStatus.STOPPING + endpoint_model.status_message = None + events.emit( + session, + message="Endpoint marked for stopping", + actor=events.UserActor.from_user(user), + targets=[events.Target.from_model(endpoint_model)], + ) + await session.commit() + if pipeline_hinter is not None: + pipeline_hinter.hint_fetch(EndpointModel.__name__) + + +def endpoint_model_to_endpoint(endpoint_model: EndpointModel) -> Endpoint: + configuration = get_endpoint_configuration(endpoint_model) + run_name = None + url = None + status = endpoint_model.status + if ( + status != EndpointStatus.STOPPED + and endpoint_model.service_run is not None + and not endpoint_model.service_run.deleted + ): + run_name = endpoint_model.service_run.run_name + if ( + status == EndpointStatus.RUNNING + and endpoint_model.service_run.service_spec is not None + ): + service_spec = ServiceSpec.__response__.parse_raw( + endpoint_model.service_run.service_spec + ) + if service_spec.model is not None: + url = service_spec.model.base_url + return Endpoint( + id=endpoint_model.id, + name=endpoint_model.name, + project_name=endpoint_model.project.name, + user=endpoint_model.user.name, + configuration=configuration, + created_at=endpoint_model.created_at, + last_processed_at=endpoint_model.last_processed_at, + status=status, + status_message=endpoint_model.status_message, + run_name=run_name, + url=url, + model_base=endpoint_model.model_base, + model_repo=endpoint_model.model_repo, + error=( + endpoint_model.status_message + if endpoint_model.status == EndpointStatus.FAILED + else None + ), + ) + + +def get_endpoint_configuration(endpoint_model: EndpointModel) -> EndpointConfiguration: + return EndpointConfiguration.__response__.parse_raw(endpoint_model.configuration) + + +async def record_endpoint_run_submission( + session: AsyncSession, + endpoint_id: uuid.UUID, + run_id: uuid.UUID, +) -> EndpointRunSubmissionModel: + existing_submission = await _get_endpoint_run_submission_by_run_id( + session=session, + run_id=run_id, + ) + if existing_submission is not None: + if existing_submission.endpoint_id != endpoint_id: + raise ServerClientError("Run is already recorded for another endpoint") + return existing_submission + + res = await session.execute( + select(func.max(EndpointRunSubmissionModel.submission_num)).where( + EndpointRunSubmissionModel.endpoint_id == endpoint_id + ) + ) + submission_num = (res.scalar_one_or_none() or 0) + 1 + submission = EndpointRunSubmissionModel( + endpoint_id=endpoint_id, + run_id=run_id, + submission_num=submission_num, + submitted_at=common.get_current_datetime(), + ) + session.add(submission) + await session.flush() + return submission + + +async def _get_endpoint_run_submission_by_run_id( + session: AsyncSession, + run_id: uuid.UUID, +) -> Optional[EndpointRunSubmissionModel]: + res = await session.execute( + select(EndpointRunSubmissionModel).where(EndpointRunSubmissionModel.run_id == run_id) + ) + return res.scalar_one_or_none() + + +async def generate_endpoint_name(session: AsyncSession, project: ProjectModel) -> str: + res = await session.execute( + select(EndpointModel.name).where( + EndpointModel.project_id == project.id, + ) + ) + names = set(res.scalars().all()) + while True: + name = random_names.generate_name() + if name not in names: + return name + + +def _validate_endpoint_configuration(configuration: EndpointConfiguration): + if not configuration.model.api_model_name.strip(): + raise ServerClientError("Endpoint must specify model") + if configuration.name is not None: + validate_dstack_resource_name(configuration.name) + try: + configuration.env.as_dict() + except ValueError as e: + raise ServerClientError(f"Endpoint env is unresolved: {e}") from e diff --git a/src/dstack/_internal/server/services/endpoints/agent/__init__.py b/src/dstack/_internal/server/services/endpoints/agent/__init__.py new file mode 100644 index 0000000000..0b20911db7 --- /dev/null +++ b/src/dstack/_internal/server/services/endpoints/agent/__init__.py @@ -0,0 +1,100 @@ +import uuid +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Optional + +from dstack._internal.server import settings +from dstack._internal.server.models import EndpointModel +from dstack._internal.server.services.endpoints.agent.report import AgentFinalReport +from dstack._internal.server.services.pipelines import PipelineHinterProtocol + + +@dataclass(frozen=True) +class AgentPlan: + model: str + + +@dataclass(frozen=True) +class AgentProvisioningResult: + run_id: Optional[uuid.UUID] = None + run_name: Optional[str] = None + submitted_run_ids: tuple[uuid.UUID, ...] = () + submitted_run_names: tuple[str, ...] = () + error: Optional[str] = None + final_report: Optional[AgentFinalReport] = None + in_progress: bool = False + + +class AgentService(ABC): + @abstractmethod + def is_enabled(self) -> bool: + pass + + @abstractmethod + def get_plan(self) -> AgentPlan: + pass + + @abstractmethod + async def provision_endpoint( + self, + endpoint_model: EndpointModel, + pipeline_hinter: PipelineHinterProtocol, + ) -> AgentProvisioningResult: + pass + + async def abort_endpoint(self, endpoint_model: EndpointModel) -> bool: + return True + + +class DisabledAgentService(AgentService): + def __init__(self, reason: Optional[str] = None) -> None: + self._reason = reason + + def is_enabled(self) -> bool: + return False + + def get_plan(self) -> AgentPlan: + return AgentPlan(model=settings.AGENT_ANTHROPIC_MODEL) + + async def provision_endpoint( + self, + endpoint_model: EndpointModel, + pipeline_hinter: PipelineHinterProtocol, + ) -> AgentProvisioningResult: + return AgentProvisioningResult(error=self._reason or get_agent_unavailable_reason()) + + +async def abort_agent_endpoint(endpoint_model: EndpointModel) -> bool: + # Cancellation must work even if a new agent session cannot be started because + # the API key or Claude executable is no longer configured. + from dstack._internal.server.services.endpoints.agent.claude import ClaudeAgentService + + return await ClaudeAgentService().abort_endpoint(endpoint_model) + + +def get_agent_service() -> AgentService: + if _has_claude_agent_auth(): + from dstack._internal.server.services.endpoints.agent.claude import ( + ClaudeAgentService, + get_claude_agent_unavailable_reason, + ) + + unavailable_reason = get_claude_agent_unavailable_reason() + if unavailable_reason is None: + return ClaudeAgentService() + return DisabledAgentService(reason=unavailable_reason) + return DisabledAgentService(reason=get_agent_unavailable_reason()) + + +def get_agent_unavailable_reason() -> Optional[str]: + if not _has_claude_agent_auth(): + return "DSTACK_AGENT_ANTHROPIC_API_KEY is not set." + from dstack._internal.server.services.endpoints.agent.claude import ( + get_claude_agent_unavailable_reason, + ) + + return get_claude_agent_unavailable_reason() + + +def _has_claude_agent_auth() -> bool: + return bool(settings.AGENT_ANTHROPIC_API_KEY or settings.AGENT_CLAUDE_USE_EXISTING_AUTH) diff --git a/src/dstack/_internal/server/services/endpoints/agent/claude.py b/src/dstack/_internal/server/services/endpoints/agent/claude.py new file mode 100644 index 0000000000..a84f14d32d --- /dev/null +++ b/src/dstack/_internal/server/services/endpoints/agent/claude.py @@ -0,0 +1,2108 @@ +import asyncio +import enum +import hashlib +import json +import os +import shutil +import signal +import socket +import subprocess +import sys +from contextlib import suppress +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Awaitable, Callable, Optional, Sequence +from uuid import UUID + +import yaml +from pydantic import ValidationError +from sqlalchemy import exists, func, or_, select + +from dstack._internal.core.models.config import GlobalConfig, ProjectConfig +from dstack._internal.core.models.endpoints import EndpointConfiguration +from dstack._internal.core.models.fleets import FleetStatus +from dstack._internal.server import settings +from dstack._internal.server.db import get_session_ctx +from dstack._internal.server.models import ( + EndpointAgentSessionModel, + EndpointAgentSessionStatus, + EndpointModel, + ExportedFleetModel, + FleetModel, + ImportModel, + ProjectModel, +) +from dstack._internal.server.schemas.runner import LogEvent +from dstack._internal.server.services import logs as logs_services +from dstack._internal.server.services.endpoints.agent import ( + AgentPlan, + AgentProvisioningResult, + AgentService, +) +from dstack._internal.server.services.endpoints.agent.report import ( + AGENT_FINAL_REPORT_JSON_SCHEMA, + AgentFinalReport, +) +from dstack._internal.server.services.pipelines import PipelineHinterProtocol +from dstack._internal.utils.common import ( + get_current_datetime, + get_milliseconds_since_epoch, + run_async, +) +from dstack._internal.utils.logging import get_logger + +logger = get_logger(__name__) + +_RESOURCE_DIR = Path(__file__).parent / "resources" +_ENDPOINT_PROMPT_RESOURCE_NAME = "system_prompt.md" +_AGENT_SKILL_NAMES = ["dstack", "dstack-prototyping"] +_INHERITED_ENV_NAMES = [ + "PATH", + "SSL_CERT_FILE", + "REQUESTS_CA_BUNDLE", + "HTTP_PROXY", + "HTTPS_PROXY", + "NO_PROXY", + "http_proxy", + "https_proxy", + "no_proxy", +] +_REDACTION = "[redacted]" +_MAX_CAPTURED_OUTPUT_CHARS = 20_000 +_MAX_AGENT_LOG_MESSAGE_CHARS = 4_000 +_AGENT_LOG_BATCH_SIZE = 1 +_AGENT_PROGRESS_LOG_NAME = "progress.jsonl" +_AGENT_SUBMISSIONS_LOG_NAME = "submissions.jsonl" +_AGENT_PROCESS_STATE_NAME = "agent_process.json" +_AGENT_STDOUT_LOG_NAME = "agent_stdout.jsonl" +_AGENT_STDERR_LOG_NAME = "agent_stderr.jsonl" +_AGENT_PROGRESS_POLL_SECONDS = 1.0 +_AGENT_PROCESS_ABORT_GRACE_SECONDS = 1.0 +_CLAUDE_AGENT_TOOLS = "Bash,Read,Write,Edit,WebFetch,WebSearch,StructuredOutput" +_CLAUDE_EFFORT_LEVELS = ("low", "medium", "high", "xhigh", "max") +_AGENT_ABORT_MESSAGE = "Endpoint stop requested" +_AGENT_BIN_DIR_NAME = "bin" +_AGENT_PROGRESS_ENV = "DSTACK_ENDPOINT_PROGRESS_LOG" +_ENDPOINT_NO_FLEETS_MESSAGE = ( + "The project has no fleets. Create one before submitting an endpoint." +) +_ENDPOINT_NO_MATCHING_FLEETS_MESSAGE = ( + "No fleets match the endpoint configuration. Create a fleet or update `fleets` before " + "submitting an endpoint." +) +_AGENT_START_PROGRESS_TEMPLATE = ( + "Starting endpoint prototyping agent for {model}. Allowed fleets: {fleets}. " + "The agent will inspect offers, choose a service recipe, deploy it, and verify " + "the model API before the endpoint becomes active." +) + + +@dataclass(frozen=True) +class _AgentRunnerResult: + report: Optional[AgentFinalReport] = None + error: Optional[str] = None + + +@dataclass +class _AgentProcessOutput: + report_data: Optional[dict[str, Any]] = None + result_error: Optional[str] = None + stdout_tail: str = "" + + +@dataclass(frozen=True) +class _AgentProcessState: + pid: int + pgid: int + host: Optional[str] + + +@dataclass(frozen=True) +class _EndpointAgentConstraints: + prompt_text: str + allowed_fleets: tuple[str, ...] + + +class ClaudeAgentService(AgentService): + def __init__( + self, + runner: Optional[ + Callable[["_AgentWorkspace", dict[str, Any]], Awaitable[_AgentRunnerResult]] + ] = None, + workspace_base_dir: Path = settings.SERVER_DATA_DIR_PATH / "endpoint_agent_runs", + ) -> None: + self._runner = runner + self._workspace_base_dir = workspace_base_dir + + def is_enabled(self) -> bool: + return get_claude_agent_unavailable_reason() is None + + def get_plan(self) -> AgentPlan: + return AgentPlan(model=settings.AGENT_ANTHROPIC_MODEL) + + async def provision_endpoint( + self, + endpoint_model: EndpointModel, + pipeline_hinter: PipelineHinterProtocol, + ) -> AgentProvisioningResult: + unavailable_reason = get_claude_agent_unavailable_reason() + if unavailable_reason is not None: + return AgentProvisioningResult(error=unavailable_reason) + + try: + configuration = EndpointConfiguration.__response__.parse_raw( + endpoint_model.configuration + ) + agent_constraints = await _get_endpoint_agent_constraints( + endpoint_model=endpoint_model, + configuration=configuration, + ) + if not agent_constraints.allowed_fleets: + return AgentProvisioningResult( + error=_get_endpoint_no_fleets_message(configuration) + ) + agent_session = await _get_or_create_agent_session( + endpoint_model=endpoint_model, + workspace_base_dir=self._workspace_base_dir, + ) + workspace = _prepare_workspace( + endpoint_model=endpoint_model, + workspace_root_dir=Path(agent_session.workspace_path), + agent_constraints=agent_constraints, + ) + except Exception as e: + logger.warning("Failed to prepare endpoint agent workspace: %s", e, exc_info=True) + return AgentProvisioningResult( + error=f"Failed to prepare endpoint agent workspace: {e}" + ) + + if self._runner is not None: + return await _run_injected_agent_runner( + agent_session=agent_session, + workspace=workspace, + runner=self._runner, + ) + return await _reconcile_detached_agent_session( + agent_session=agent_session, + workspace=workspace, + ) + + async def abort_endpoint(self, endpoint_model: EndpointModel) -> bool: + agent_session = await _get_latest_agent_session(endpoint_model) + if agent_session is None or agent_session.status != EndpointAgentSessionStatus.RUNNING: + return True + try: + workspace = _prepare_workspace( + endpoint_model=endpoint_model, + workspace_root_dir=Path(agent_session.workspace_path), + install_skills=False, + ) + except Exception: + logger.warning( + "Failed to prepare endpoint agent workspace for abort: endpoint=%s", + endpoint_model.name, + exc_info=True, + ) + return False + + aborted = await _abort_agent_session_process( + agent_session=agent_session, + workspace=workspace, + ) + if not aborted: + return False + workspace.artifacts.record_error(_AGENT_ABORT_MESSAGE) + await _mark_agent_session_failed(agent_session, _AGENT_ABORT_MESSAGE) + await _flush_agent_logs(workspace) + return True + + +async def _run_injected_agent_runner( + *, + agent_session: EndpointAgentSessionModel, + workspace: "_AgentWorkspace", + runner: Callable[["_AgentWorkspace", dict[str, Any]], Awaitable[_AgentRunnerResult]], +) -> AgentProvisioningResult: + reconcile_result = await _reconcile_agent_artifacts( + agent_session=agent_session, + workspace=workspace, + ) + if reconcile_result is not None: + return reconcile_result + + runner_result = await runner(workspace, _build_agent_request(workspace)) + submissions = _load_submissions(workspace) + if runner_result.error is not None: + workspace.artifacts.record_error(runner_result.error) + await _mark_agent_session_failed(agent_session, runner_result.error) + return AgentProvisioningResult( + error=runner_result.error, + submitted_run_ids=submissions.run_ids, + submitted_run_names=submissions.run_names, + ) + report = runner_result.report + if report is None: + error = "Server agent did not return a final report" + workspace.artifacts.record_error(error) + await _mark_agent_session_failed(agent_session, error) + return AgentProvisioningResult( + error=error, + submitted_run_ids=submissions.run_ids, + submitted_run_names=submissions.run_names, + ) + workspace.artifacts.record_report(report) + await _mark_agent_session_from_report(agent_session, report) + + logger.info( + "Endpoint agent finished for endpoint %s: success=%s run_id=%s run_name=%s", + workspace.endpoint_name, + report.success, + report.run_id, + report.run_name, + ) + return AgentProvisioningResult( + run_id=report.run_id, + run_name=report.run_name, + submitted_run_ids=submissions.run_ids, + submitted_run_names=submissions.run_names, + final_report=report, + ) + + +async def _reconcile_detached_agent_session( + *, + agent_session: EndpointAgentSessionModel, + workspace: "_AgentWorkspace", +) -> AgentProvisioningResult: + reconcile_result = await _reconcile_agent_artifacts( + agent_session=agent_session, + workspace=workspace, + ) + if reconcile_result is not None: + return reconcile_result + + running_pid = _get_running_agent_process_pid(workspace, agent_session=agent_session) + if running_pid is not None: + logger.info( + "Endpoint agent process %s is already running for endpoint %s", + running_pid, + workspace.endpoint_name, + ) + submissions = _load_submissions(workspace) + return AgentProvisioningResult( + in_progress=True, + submitted_run_ids=submissions.run_ids, + submitted_run_names=submissions.run_names, + ) + + if agent_session.pid is not None: + logger.info( + "Endpoint agent process %s is no longer running for endpoint %s; " + "resuming session %s from workspace", + agent_session.pid, + workspace.endpoint_name, + agent_session.session_num, + ) + _write_trace_record( + workspace, + { + "type": "agent-process-resume", + "pid": agent_session.pid, + "process_host": agent_session.process_host, + "session_num": agent_session.session_num, + }, + ) + await _write_agent_session_progress( + agent_session, + workspace, + "Resuming endpoint prototyping agent from the existing workspace.", + ) + await _clear_agent_session_process(agent_session) + + try: + await _write_agent_session_progress( + agent_session, + workspace, + _get_agent_start_progress_message(workspace), + ) + pid = _start_agent_subprocess_detached(workspace, _build_agent_request(workspace)) + except Exception as e: + logger.warning("Failed to start endpoint agent process: %s", e, exc_info=True) + error = f"Failed to start endpoint agent process: {e}" + workspace.artifacts.record_error(error) + await _mark_agent_session_failed(agent_session, error) + return AgentProvisioningResult(error=error) + + await _mark_agent_session_running(agent_session, pid=pid, process_host=_get_process_host()) + return AgentProvisioningResult(in_progress=True) + + +async def _reconcile_agent_artifacts( + *, + agent_session: EndpointAgentSessionModel, + workspace: "_AgentWorkspace", +) -> Optional[AgentProvisioningResult]: + process_output = await _reconcile_agent_stream_files( + agent_session=agent_session, + workspace=workspace, + ) + await _reconcile_agent_progress(agent_session=agent_session, workspace=workspace) + submissions = _load_submissions(workspace) + + report = _load_final_report(workspace) + if report is None and process_output.report_data is not None: + try: + report = AgentFinalReport.parse_obj(process_output.report_data) + except ValidationError as e: + error = f"Server agent returned an invalid final report: {e}" + workspace.artifacts.record_error(error) + await _mark_agent_session_failed(agent_session, error) + return AgentProvisioningResult( + error=error, + submitted_run_ids=submissions.run_ids, + submitted_run_names=submissions.run_names, + ) + if report is not None: + if _get_running_agent_process_pid(workspace, agent_session=agent_session) is not None: + return AgentProvisioningResult( + in_progress=True, + submitted_run_ids=submissions.run_ids, + submitted_run_names=submissions.run_names, + ) + workspace.artifacts.record_report(report) + await _mark_agent_session_from_report(agent_session, report) + return AgentProvisioningResult( + run_id=report.run_id, + run_name=report.run_name, + submitted_run_ids=submissions.run_ids, + submitted_run_names=submissions.run_names, + final_report=report, + ) + + error = _load_agent_error(workspace) + if error is not None: + await _mark_agent_session_failed(agent_session, error) + return AgentProvisioningResult( + error=error, + submitted_run_ids=submissions.run_ids, + submitted_run_names=submissions.run_names, + ) + if process_output.result_error is not None: + error = "Server agent failed before returning a final report" + workspace.artifacts.record_error(error) + await _mark_agent_session_failed(agent_session, error) + return AgentProvisioningResult( + error=error, + submitted_run_ids=submissions.run_ids, + submitted_run_names=submissions.run_names, + ) + return None + + +async def _get_or_create_agent_session( + *, + endpoint_model: EndpointModel, + workspace_base_dir: Path, +) -> EndpointAgentSessionModel: + async with get_session_ctx() as db_session: + res = await db_session.execute( + select(EndpointAgentSessionModel) + .where(EndpointAgentSessionModel.endpoint_id == endpoint_model.id) + .order_by(EndpointAgentSessionModel.session_num.desc()) + .limit(1) + ) + agent_session = res.scalar_one_or_none() + if agent_session is not None and agent_session.created_at >= endpoint_model.created_at: + return agent_session + + max_session_num_res = await db_session.execute( + select(func.max(EndpointAgentSessionModel.session_num)).where( + EndpointAgentSessionModel.endpoint_id == endpoint_model.id + ) + ) + session_num = (max_session_num_res.scalar_one_or_none() or 0) + 1 + now = get_current_datetime() + legacy_workspace_path = workspace_base_dir / str(endpoint_model.id) + if session_num == 1 and _has_agent_workspace_artifacts(legacy_workspace_path): + workspace_path = legacy_workspace_path + else: + workspace_path = workspace_base_dir / str(endpoint_model.id) / str(session_num) + agent_session = EndpointAgentSessionModel( + endpoint_id=endpoint_model.id, + session_num=session_num, + status=EndpointAgentSessionStatus.RUNNING, + workspace_path=str(workspace_path), + progress_log_offset=0, + stdout_log_offset=0, + stderr_log_offset=0, + created_at=now, + updated_at=now, + ) + db_session.add(agent_session) + await db_session.commit() + return agent_session + + +async def _get_latest_agent_session( + endpoint_model: EndpointModel, +) -> Optional[EndpointAgentSessionModel]: + async with get_session_ctx() as db_session: + res = await db_session.execute( + select(EndpointAgentSessionModel) + .where(EndpointAgentSessionModel.endpoint_id == endpoint_model.id) + .order_by(EndpointAgentSessionModel.session_num.desc()) + .limit(1) + ) + return res.scalar_one_or_none() + + +def _has_agent_workspace_artifacts(root_dir: Path) -> bool: + work_dir = root_dir / "workspace" + return any( + (work_dir / filename).exists() + for filename in [ + "final_report.json", + "agent_error.json", + _AGENT_PROCESS_STATE_NAME, + _AGENT_PROGRESS_LOG_NAME, + ] + ) + + +async def _get_agent_session(db_session, agent_session: EndpointAgentSessionModel): + return await db_session.get( + EndpointAgentSessionModel, + { + "endpoint_id": agent_session.endpoint_id, + "session_num": agent_session.session_num, + }, + ) + + +async def _mark_agent_session_running( + agent_session: EndpointAgentSessionModel, + *, + pid: int, + process_host: str, +) -> None: + async with get_session_ctx() as db_session: + stored_session = await _get_agent_session(db_session, agent_session) + if stored_session is None: + return + stored_session.status = EndpointAgentSessionStatus.RUNNING + stored_session.pid = pid + stored_session.process_host = process_host + stored_session.updated_at = get_current_datetime() + await db_session.commit() + agent_session.pid = pid + agent_session.process_host = process_host + + +async def _clear_agent_session_process(agent_session: EndpointAgentSessionModel) -> None: + async with get_session_ctx() as db_session: + stored_session = await _get_agent_session(db_session, agent_session) + if stored_session is None: + return + stored_session.pid = None + stored_session.process_host = None + stored_session.updated_at = get_current_datetime() + await db_session.commit() + agent_session.pid = None + agent_session.process_host = None + + +async def _mark_agent_session_from_report( + agent_session: EndpointAgentSessionModel, + report: AgentFinalReport, +) -> None: + if report.success: + await _mark_agent_session_succeeded(agent_session) + else: + await _mark_agent_session_failed( + agent_session, + report.failure_summary or "Server agent did not verify the endpoint", + ) + + +async def _mark_agent_session_succeeded( + agent_session: EndpointAgentSessionModel, +) -> None: + now = get_current_datetime() + async with get_session_ctx() as db_session: + stored_session = await _get_agent_session(db_session, agent_session) + if stored_session is None: + return + stored_session.status = EndpointAgentSessionStatus.SUCCEEDED + stored_session.status_message = None + stored_session.finished_at = stored_session.finished_at or now + stored_session.updated_at = now + await db_session.commit() + agent_session.status = EndpointAgentSessionStatus.SUCCEEDED + agent_session.status_message = None + agent_session.finished_at = stored_session.finished_at + + +async def _mark_agent_session_failed( + agent_session: EndpointAgentSessionModel, + error: str, +) -> None: + now = get_current_datetime() + async with get_session_ctx() as db_session: + stored_session = await _get_agent_session(db_session, agent_session) + if stored_session is None: + return + stored_session.status = EndpointAgentSessionStatus.FAILED + stored_session.status_message = error + stored_session.finished_at = stored_session.finished_at or now + stored_session.updated_at = now + await db_session.commit() + agent_session.status = EndpointAgentSessionStatus.FAILED + agent_session.status_message = error + agent_session.finished_at = stored_session.finished_at + + +async def _update_agent_session_offsets( + agent_session: EndpointAgentSessionModel, + *, + progress_log_offset: Optional[int] = None, + stdout_log_offset: Optional[int] = None, + stderr_log_offset: Optional[int] = None, +) -> None: + async with get_session_ctx() as db_session: + stored_session = await _get_agent_session(db_session, agent_session) + if stored_session is None: + return + if progress_log_offset is not None: + stored_session.progress_log_offset = progress_log_offset + agent_session.progress_log_offset = progress_log_offset + if stdout_log_offset is not None: + stored_session.stdout_log_offset = stdout_log_offset + agent_session.stdout_log_offset = stdout_log_offset + if stderr_log_offset is not None: + stored_session.stderr_log_offset = stderr_log_offset + agent_session.stderr_log_offset = stderr_log_offset + stored_session.updated_at = get_current_datetime() + await db_session.commit() + + +async def _reconcile_agent_progress( + *, + agent_session: EndpointAgentSessionModel, + workspace: "_AgentWorkspace", +) -> None: + offset = await _flush_agent_progress_logs( + workspace, + offset=agent_session.progress_log_offset, + include_partial=False, + ) + if offset != agent_session.progress_log_offset: + await _update_agent_session_offsets(agent_session, progress_log_offset=offset) + + +async def _write_agent_session_progress( + agent_session: EndpointAgentSessionModel, + workspace: "_AgentWorkspace", + message: str, +) -> None: + _append_agent_progress_record(workspace, message) + offset = workspace.progress_path.stat().st_size + await _write_agent_log(workspace, message) + await _flush_agent_logs(workspace) + await _update_agent_session_offsets(agent_session, progress_log_offset=offset) + + +def _append_agent_progress_record(workspace: "_AgentWorkspace", message: str) -> None: + workspace.progress_path.parent.mkdir(parents=True, exist_ok=True) + with workspace.progress_path.open("a", encoding="utf-8") as f: + f.write(json.dumps({"message": message}, ensure_ascii=True) + "\n") + + +def _get_agent_start_progress_message(workspace: "_AgentWorkspace") -> str: + fleets = ", ".join(workspace.allowed_fleets) if workspace.allowed_fleets else "none" + return _AGENT_START_PROGRESS_TEMPLATE.format(model=workspace.model, fleets=fleets) + + +async def _reconcile_agent_stream_files( + *, + agent_session: EndpointAgentSessionModel, + workspace: "_AgentWorkspace", +) -> _AgentProcessOutput: + stdout_output, stdout_offset = await _read_agent_stream_file( + workspace=workspace, + path=workspace.work_dir / _AGENT_STDOUT_LOG_NAME, + offset=agent_session.stdout_log_offset, + stream_name="stdout", + ) + stderr_output, stderr_offset = await _read_agent_stream_file( + workspace=workspace, + path=workspace.work_dir / _AGENT_STDERR_LOG_NAME, + offset=agent_session.stderr_log_offset, + stream_name="stderr", + ) + output = _merge_process_outputs(stdout_output, stderr_output) + if ( + stdout_offset != agent_session.stdout_log_offset + or stderr_offset != agent_session.stderr_log_offset + ): + await _update_agent_session_offsets( + agent_session, + stdout_log_offset=stdout_offset, + stderr_log_offset=stderr_offset, + ) + return output + + +@dataclass(frozen=True) +class _AgentSubmissions: + run_ids: tuple[UUID, ...] = () + run_names: tuple[str, ...] = () + + +def _load_submissions(workspace: "_AgentWorkspace") -> _AgentSubmissions: + path = workspace.work_dir / _AGENT_SUBMISSIONS_LOG_NAME + if not path.exists(): + return _AgentSubmissions() + run_ids: list[UUID] = [] + run_names: list[str] = [] + seen_run_ids: set[UUID] = set() + seen_run_names: set[str] = set() + for line in path.read_text(encoding="utf-8", errors="replace").splitlines(): + if not line.strip(): + continue + try: + parsed = json.loads(line) + except json.JSONDecodeError: + _write_trace_record( + workspace, + {"type": "agent-submission-invalid", "path": str(path), "line": line}, + ) + continue + if not isinstance(parsed, dict): + continue + raw_name = parsed.get("name") + if isinstance(raw_name, str): + run_name = raw_name.strip() + if run_name and run_name not in seen_run_names: + seen_run_names.add(run_name) + run_names.append(run_name) + raw_run_id = parsed.get("run_id") + if not isinstance(raw_run_id, str) or not raw_run_id.strip(): + continue + try: + run_id = UUID(raw_run_id) + except ValueError: + _write_trace_record( + workspace, + { + "type": "agent-submission-invalid-run-id", + "path": str(path), + "run_id": raw_run_id, + }, + ) + continue + if run_id not in seen_run_ids: + seen_run_ids.add(run_id) + run_ids.append(run_id) + return _AgentSubmissions(run_ids=tuple(run_ids), run_names=tuple(run_names)) + + +def _load_agent_error(workspace: "_AgentWorkspace") -> Optional[str]: + path = workspace.work_dir / "agent_error.json" + if not path.exists(): + return None + try: + parsed = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + return "Server agent failed" + if not isinstance(parsed, dict): + return "Server agent failed" + failure_summary = parsed.get("failure_summary") + if isinstance(failure_summary, str) and failure_summary.strip(): + return failure_summary + return "Server agent failed" + + +def get_claude_agent_unavailable_reason() -> Optional[str]: + if settings.AGENT_ANTHROPIC_API_KEY and settings.AGENT_CLAUDE_USE_EXISTING_AUTH: + return ( + "DSTACK_AGENT_ANTHROPIC_API_KEY and " + "DSTACK_AGENT_CLAUDE_USE_EXISTING_AUTH cannot both be set." + ) + if not settings.AGENT_ANTHROPIC_API_KEY and not _should_use_existing_claude_auth(): + return ( + "DSTACK_AGENT_ANTHROPIC_API_KEY is not set. Set it or opt in to existing " + "Claude CLI auth with DSTACK_AGENT_CLAUDE_USE_EXISTING_AUTH=1." + ) + if ( + settings.AGENT_CLAUDE_EFFORT is not None + and settings.AGENT_CLAUDE_EFFORT not in _CLAUDE_EFFORT_LEVELS + ): + return f"DSTACK_AGENT_CLAUDE_EFFORT must be one of: {', '.join(_CLAUDE_EFFORT_LEVELS)}." + if _get_claude_executable() is None: + if settings.AGENT_CLAUDE_PATH is not None: + return ( + "DSTACK_AGENT_CLAUDE_PATH does not resolve to an executable: " + f"{settings.AGENT_CLAUDE_PATH}" + ) + return ( + "server agent runtime requires the claude executable in PATH or " + "DSTACK_AGENT_CLAUDE_PATH." + ) + return None + + +_existing_auth_warning_logged = False + + +class _AgentWorkspace: + def __init__( + self, + *, + root_dir: Path, + home_dir: Path, + work_dir: Path, + trace_path: Optional[Path], + env: dict[str, str], + redacted_values: Sequence[str], + endpoint_name: str, + model: str, + endpoint_context: Optional[str] = None, + dstack_home_dir: Optional[Path] = None, + endpoint_constraints: str = "", + max_price: Optional[float] = None, + spot_policy: Optional[str] = None, + allowed_fleets: Sequence[str] = (), + log_writer: Optional["_AgentLogWriter"] = None, + ) -> None: + self.root_dir = root_dir + self.home_dir = home_dir + self.dstack_home_dir = dstack_home_dir or home_dir + self.work_dir = work_dir + self.trace_path = trace_path + self.env = env + self.redacted_values = tuple(redacted_values) + self.endpoint_name = endpoint_name + self.model = model + self.endpoint_context = endpoint_context or ( + f"- service_model_name: {model}\n- model_repo: {model}" + ) + self.endpoint_constraints = endpoint_constraints + self.max_price = max_price + self.spot_policy = spot_policy + self.allowed_fleets = tuple(allowed_fleets) + self.log_writer = log_writer + self.artifacts = _AgentArtifactRecorder(self) + + @property + def progress_path(self) -> Path: + return self.work_dir / _AGENT_PROGRESS_LOG_NAME + + +class _AgentLogWriter: + def __init__( + self, + *, + project: ProjectModel, + endpoint_id: UUID, + endpoint_name: str, + ) -> None: + self._project = project + self._endpoint_id = endpoint_id + self._endpoint_name = endpoint_name + self._buffer: list[LogEvent] = [] + + async def write(self, message: str) -> None: + message = _truncate_log_message(message) + if not message.endswith("\n"): + message += "\n" + self._buffer.append( + LogEvent( + timestamp=get_milliseconds_since_epoch(), + message=message.encode(), + ) + ) + if len(self._buffer) >= _AGENT_LOG_BATCH_SIZE: + await self.flush() + + async def flush(self) -> None: + if not self._buffer: + return + events = self._buffer + self._buffer = [] + try: + await run_async( + logs_services.write_logs, + project=self._project, + run_name=self._endpoint_name, + job_submission_id=self._endpoint_id, + runner_logs=[], + job_logs=events, + ) + except Exception: + logger.warning( + "Failed to write endpoint agent logs for endpoint %s", + self._endpoint_name, + exc_info=True, + ) + + +def _prepare_workspace( + *, + endpoint_model: EndpointModel, + workspace_root_dir: Path, + install_skills: bool = True, + agent_constraints: Optional[_EndpointAgentConstraints] = None, +) -> _AgentWorkspace: + configuration = EndpointConfiguration.__response__.parse_raw(endpoint_model.configuration) + endpoint_env = configuration.env.as_dict() + if agent_constraints is None: + agent_constraints = _get_static_endpoint_agent_constraints(configuration, endpoint_env) + + root_dir = workspace_root_dir + home_dir = root_dir / "home" + agent_home_dir = _prepare_agent_home_dir(root_dir=root_dir, home_dir=home_dir) + work_dir = root_dir / "workspace" + dstack_dir = home_dir / ".dstack" + work_dir.mkdir(parents=True, exist_ok=True) + dstack_dir.mkdir(parents=True, exist_ok=True) + token = endpoint_model.user.token.get_plaintext_or_error() + allowed_fleets = agent_constraints.allowed_fleets + bin_dir = root_dir / _AGENT_BIN_DIR_NAME + _write_cli_config( + dstack_dir=dstack_dir, + project_name=endpoint_model.project.name, + server_url=settings.SERVER_URL, + token=token, + ) + + env = _build_agent_env( + home_dir=_get_claude_process_home_dir(agent_home_dir), + project_name=endpoint_model.project.name, + server_url=settings.SERVER_URL, + token=token, + endpoint_env=endpoint_env, + bin_dir=bin_dir, + ) + _warn_if_using_existing_claude_auth() + redacted_values = _get_redacted_values( + [ + token, + settings.AGENT_ANTHROPIC_API_KEY or "", + *endpoint_env.values(), + *_get_sensitive_inherited_env_values(env), + ] + ) + trace_path = root_dir / "trace.jsonl" if _is_debug_trace_enabled() else None + if configuration.model.allows_variant_selection: + endpoint_context = ( + f"- service_model_name: {configuration.model.api_model_name}\n" + f"- base_model: {configuration.model.api_model_name}" + ) + else: + endpoint_context = ( + f"- service_model_name: {configuration.model.api_model_name}\n" + f"- model_repo: {configuration.model.exact_repo}" + ) + workspace = _AgentWorkspace( + root_dir=root_dir, + home_dir=home_dir, + dstack_home_dir=agent_home_dir, + work_dir=work_dir, + trace_path=trace_path, + env=env, + redacted_values=redacted_values, + endpoint_name=endpoint_model.name, + model=configuration.model.api_model_name, + endpoint_context=endpoint_context, + endpoint_constraints=agent_constraints.prompt_text, + max_price=configuration.max_price, + spot_policy=configuration.spot_policy.value if configuration.spot_policy else None, + allowed_fleets=allowed_fleets, + log_writer=_AgentLogWriter( + project=endpoint_model.project, + endpoint_id=endpoint_model.id, + endpoint_name=endpoint_model.name, + ), + ) + workspace.env[_AGENT_PROGRESS_ENV] = str(workspace.progress_path) + _install_agent_helper_scripts(workspace) + if install_skills: + _install_agent_skills(workspace) + workspace.artifacts.initialize() + return workspace + + +def _prepare_agent_home_dir(*, root_dir: Path, home_dir: Path) -> Path: + home_dir.mkdir(parents=True, exist_ok=True) + ssh_dir = home_dir / ".ssh" + ssh_dir.mkdir(parents=True, exist_ok=True) + ssh_dir.chmod(0o700) + short_home_dir = _get_short_agent_home_dir(root_dir) + short_home_dir.parent.mkdir(parents=True, exist_ok=True) + if short_home_dir.is_symlink() and short_home_dir.resolve() == home_dir.resolve(): + return short_home_dir + if short_home_dir.exists() or short_home_dir.is_symlink(): + if short_home_dir.is_dir() and not short_home_dir.is_symlink(): + shutil.rmtree(short_home_dir) + else: + short_home_dir.unlink() + short_home_dir.symlink_to(home_dir, target_is_directory=True) + return short_home_dir + + +def _get_claude_process_home_dir(agent_home_dir: Path) -> Path: + if _should_use_existing_claude_auth(): + return Path.home() + return agent_home_dir + + +def _get_short_agent_home_dir(root_dir: Path) -> Path: + root_hash = hashlib.sha1(str(root_dir).encode()).hexdigest()[:8] + return Path("/tmp") / f"dah-{root_hash}" + + +class _AgentArtifactRecorder: + def __init__(self, workspace: _AgentWorkspace) -> None: + self._workspace = workspace + self._command_counter = 0 + + def initialize(self) -> None: + self._workspace.work_dir.mkdir(parents=True, exist_ok=True) + self._command_counter = self._get_last_command_output_num() + self._update_agent_state(phase="starting") + for filename in [ + _AGENT_SUBMISSIONS_LOG_NAME, + "commands.jsonl", + _AGENT_PROGRESS_LOG_NAME, + ]: + (self._workspace.work_dir / filename).touch(exist_ok=True) + + def mark_running(self) -> None: + self._update_agent_state(phase="running") + + def record_report(self, report: AgentFinalReport) -> None: + (self._workspace.work_dir / "final_report.json").write_text( + json.dumps( + _redact(report.dict(), self._workspace.redacted_values), + default=str, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + self._update_agent_state(phase="success" if report.success else "failure") + + def record_error(self, error: str) -> None: + data = { + "success": False, + "failure_summary": error, + "recorded_at": _utcnow_iso(), + } + (self._workspace.work_dir / "agent_error.json").write_text( + json.dumps(_redact(data, self._workspace.redacted_values), indent=2) + "\n", + encoding="utf-8", + ) + self._update_agent_state(phase="failure") + + def record_stream_message(self, message: dict[str, Any]) -> None: + self._record_bash_tool_uses(message) + self._record_tool_results(message) + + def _update_agent_state(self, phase: str) -> None: + path = self._workspace.work_dir / "agent_state.json" + state: dict[str, Any] = {} + if path.exists(): + try: + parsed = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + parsed = {} + if isinstance(parsed, dict): + state.update(parsed) + now = _utcnow_iso() + state.update( + { + "endpoint_name": self._workspace.endpoint_name, + "model": self._workspace.model, + "phase": phase, + "max_hourly_price": self._workspace.max_price, + "spot_policy": self._workspace.spot_policy, + "updated_at": now, + } + ) + state.setdefault("started_at", now) + path.write_text( + json.dumps(_redact(state, self._workspace.redacted_values), indent=2) + "\n", + encoding="utf-8", + ) + + def _record_bash_tool_uses(self, message: dict[str, Any]) -> None: + if message.get("type") != "assistant": + return + claude_message = message.get("message") + if not isinstance(claude_message, dict): + return + for item in claude_message.get("content", []): + if not isinstance(item, dict): + continue + if item.get("type") != "tool_use" or item.get("name") != "Bash": + continue + tool_input = item.get("input") + if not isinstance(tool_input, dict): + continue + command = tool_input.get("command") + if not isinstance(command, str) or not command.strip(): + continue + record = { + "event": "tool_use", + "timestamp": _utcnow_iso(), + "tool_use_id": item.get("id"), + "description": tool_input.get("description"), + "command": command, + } + self._append_jsonl("commands.jsonl", record) + + def _record_tool_results(self, message: dict[str, Any]) -> None: + if message.get("type") != "user": + return + claude_message = message.get("message") + if not isinstance(claude_message, dict): + return + for item in claude_message.get("content", []): + if not isinstance(item, dict) or item.get("type") != "tool_result": + continue + content = item.get("content") + output_path = None + if isinstance(content, str) and content.strip(): + output_path = self._write_command_output(content) + record = { + "event": "tool_result", + "timestamp": _utcnow_iso(), + "tool_use_id": item.get("tool_use_id") or message.get("parent_tool_use_id"), + "is_error": bool(item.get("is_error")), + "output_path": output_path, + } + self._append_jsonl("commands.jsonl", record) + + def _write_command_output(self, content: str) -> str: + self._command_counter += 1 + output_dir = self._workspace.work_dir / "command-output" + output_dir.mkdir(parents=True, exist_ok=True) + relative_path = Path("command-output") / f"{self._command_counter:04d}.txt" + output_path = self._workspace.work_dir / relative_path + output_path.write_text( + _redact(content, self._workspace.redacted_values), + encoding="utf-8", + ) + return relative_path.as_posix() + + def _append_jsonl(self, filename: str, record: dict[str, Any]) -> None: + path = self._workspace.work_dir / filename + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as f: + f.write( + json.dumps(_redact(record, self._workspace.redacted_values), default=str) + "\n" + ) + + def _get_last_command_output_num(self) -> int: + output_dir = self._workspace.work_dir / "command-output" + if not output_dir.exists(): + return 0 + nums = [] + for path in output_dir.glob("*.txt"): + try: + nums.append(int(path.stem)) + except ValueError: + continue + return max(nums, default=0) + + +def _write_cli_config( + *, + dstack_dir: Path, + project_name: str, + server_url: str, + token: str, +) -> None: + config = GlobalConfig( + projects=[ + ProjectConfig( + name=project_name, + url=server_url, + token=token, + default=True, + ) + ] + ) + (dstack_dir / "config.yml").write_text( + yaml.safe_dump(json.loads(config.json()), sort_keys=False), + encoding="utf-8", + ) + + +def _build_agent_env( + *, + home_dir: Path, + project_name: str, + server_url: str, + token: str, + endpoint_env: dict[str, str], + bin_dir: Path, +) -> dict[str, str]: + env = {name: value for name in _INHERITED_ENV_NAMES if (value := os.environ.get(name))} + path = env.get("PATH", "") + env["PATH"] = str(bin_dir) if not path else os.pathsep.join([str(bin_dir), path]) + env.update( + { + "HOME": str(home_dir), + "DSTACK_PROJECT": project_name, + "DSTACK_ENDPOINT_SERVER_URL": server_url, + "DSTACK_ENDPOINT_BEARER_TOKEN": token, + } + ) + if settings.AGENT_ANTHROPIC_API_KEY: + env["ANTHROPIC_API_KEY"] = settings.AGENT_ANTHROPIC_API_KEY + elif _should_use_existing_claude_auth() and os.environ.get("USER"): + env["USER"] = os.environ["USER"] + env.update(endpoint_env) + return env + + +async def _get_endpoint_agent_constraints( + *, + endpoint_model: EndpointModel, + configuration: EndpointConfiguration, +) -> _EndpointAgentConstraints: + endpoint_env = configuration.env.as_dict() + allowed_fleets = await _get_endpoint_allowed_fleets( + endpoint_model=endpoint_model, + configuration=configuration, + ) + return _EndpointAgentConstraints( + prompt_text=_format_endpoint_constraints( + configuration, + endpoint_env, + allowed_fleets=allowed_fleets, + ), + allowed_fleets=allowed_fleets, + ) + + +def _get_static_endpoint_agent_constraints( + configuration: EndpointConfiguration, + endpoint_env: dict[str, str], +) -> _EndpointAgentConstraints: + allowed_fleets = _get_configured_endpoint_fleet_refs(configuration) + return _EndpointAgentConstraints( + prompt_text=_format_endpoint_constraints( + configuration, + endpoint_env, + allowed_fleets=allowed_fleets, + ), + allowed_fleets=allowed_fleets, + ) + + +async def _get_endpoint_allowed_fleets( + *, + endpoint_model: EndpointModel, + configuration: EndpointConfiguration, +) -> tuple[str, ...]: + if configuration.fleets is not None and len(configuration.fleets) == 0: + return () + configured_fleets = _get_configured_endpoint_fleet_refs(configuration) + if configured_fleets: + return configured_fleets + + async with get_session_ctx() as db_session: + is_fleet_imported_subquery = exists().where( + ImportModel.project_id == endpoint_model.project_id, + ImportModel.export_id == ExportedFleetModel.export_id, + ExportedFleetModel.fleet_id == FleetModel.id, + ) + res = await db_session.execute( + select(FleetModel.name, FleetModel.project_id, ProjectModel.name) + .join(FleetModel.project) + .where( + FleetModel.deleted == False, + FleetModel.status == FleetStatus.ACTIVE, + or_( + FleetModel.project_id == endpoint_model.project_id, + is_fleet_imported_subquery, + ), + ) + .order_by(ProjectModel.name, FleetModel.name) + ) + return tuple( + name if project_id == endpoint_model.project_id else f"{project_name}/{name}" + for name, project_id, project_name in res.all() + ) + + +def _get_configured_endpoint_fleet_refs( + configuration: EndpointConfiguration, +) -> tuple[str, ...]: + if not configuration.fleets: + return () + return tuple(_format_constraint_value(fleet) for fleet in configuration.fleets) + + +def _get_endpoint_no_fleets_message(configuration: EndpointConfiguration) -> str: + if configuration.fleets: + return _ENDPOINT_NO_MATCHING_FLEETS_MESSAGE + return _ENDPOINT_NO_FLEETS_MESSAGE + + +def _install_agent_helper_scripts(workspace: _AgentWorkspace) -> None: + bin_dir = workspace.root_dir / _AGENT_BIN_DIR_NAME + bin_dir.mkdir(parents=True, exist_ok=True) + scripts = { + "progress": _get_agent_progress_script(), + } + if _should_use_existing_claude_auth(): + scripts["dstack"] = _get_agent_home_script("dstack", workspace.dstack_home_dir) + scripts["ssh"] = _get_agent_home_script("ssh", workspace.dstack_home_dir) + for name, script in scripts.items(): + script_path = bin_dir / name + script_path.write_text(script, encoding="utf-8") + script_path.chmod(0o755) + + +def _get_agent_progress_script() -> str: + return f"""#!{sys.executable} +import json +import os +from pathlib import Path +import sys + +PROGRESS_ENV = "{_AGENT_PROGRESS_ENV}" + + +def main(): + message = " ".join(sys.argv[1:]).strip() + if not message and not sys.stdin.isatty(): + message = sys.stdin.read().strip() + if not message: + print("Usage: progress ", file=sys.stderr) + return 2 + path = Path(os.environ.get(PROGRESS_ENV, "progress.jsonl")) + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as f: + f.write(json.dumps({{"message": message}}, ensure_ascii=False) + "\\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) +""" + + +def _get_agent_home_script(command: str, home_dir: Path) -> str: + real_command = shutil.which(command) + if real_command is None: + return f"""#!{sys.executable} +import sys + +print("Endpoint agent could not find the real {command} executable.", file=sys.stderr) +raise SystemExit(127) +""" + return f"""#!{sys.executable} +import os +import sys + +os.environ["HOME"] = {json.dumps(str(home_dir))} +os.execv({json.dumps(real_command)}, [{json.dumps(real_command)}, *sys.argv[1:]]) +""" + + +def _build_agent_request(workspace: _AgentWorkspace) -> dict[str, Any]: + return { + "prompt": _build_prompt(workspace), + "cwd": str(workspace.work_dir), + "env": workspace.env, + "trace_path": str(workspace.trace_path) if workspace.trace_path is not None else None, + "redacted_values": list(workspace.redacted_values), + "options": { + "tools": _CLAUDE_AGENT_TOOLS, + "allowed_tools": _CLAUDE_AGENT_TOOLS, + "disallowed_tools": "Task,NotebookEdit", + "model": settings.AGENT_ANTHROPIC_MODEL, + "max_turns": None, + "json_schema": AGENT_FINAL_REPORT_JSON_SCHEMA, + }, + } + + +def _build_prompt(workspace: _AgentWorkspace) -> str: + return f"""{_load_endpoint_prompt()} + +Endpoint context: +- endpoint_name: {workspace.endpoint_name} +{workspace.endpoint_context} + +{workspace.endpoint_constraints} +""" + + +def _load_endpoint_prompt() -> str: + return (_RESOURCE_DIR / _ENDPOINT_PROMPT_RESOURCE_NAME).read_text(encoding="utf-8").strip() + + +def _install_agent_skills(workspace: _AgentWorkspace) -> None: + source_dir = _get_packaged_skills_dir() + target_dir = workspace.work_dir / ".claude" / "skills" + if target_dir.exists(): + shutil.rmtree(target_dir) + target_dir.mkdir(parents=True, exist_ok=True) + for skill_name in _AGENT_SKILL_NAMES: + source = source_dir / skill_name + if not (source / "SKILL.md").is_file(): + raise FileNotFoundError(f"Missing endpoint agent skill: {skill_name}") + shutil.copytree(source, target_dir / skill_name) + + +def _get_packaged_skills_dir() -> Path: + for parent in Path(__file__).resolve().parents: + skills_dir = parent / "skills" + if (skills_dir / "dstack" / "SKILL.md").is_file(): + return skills_dir + raise FileNotFoundError("Could not find packaged dstack skills") + + +def _format_endpoint_constraints( + configuration: EndpointConfiguration, + endpoint_env: dict[str, str], + *, + allowed_fleets: Sequence[str], +) -> str: + lines = [ + "Fixed endpoint constraints:", + "- Do not submit any task or service that conflicts with these values.", + "- Put applicable fixed constraints into the final service YAML.", + ] + if configuration.max_price is not None: + lines.append(f"- max_price: {configuration.max_price}") + if configuration.spot_policy is not None: + lines.append(f"- spot_policy: {configuration.spot_policy.value}") + for field in ["backends", "regions", "availability_zones", "instance_types"]: + values = getattr(configuration, field) + if not values: + continue + formatted_values = [_format_constraint_value(value) for value in values] + lines.append(f"- {field}: {', '.join(formatted_values)}") + if allowed_fleets: + lines.append(f"- fleets: {', '.join(allowed_fleets)}") + if len(lines) == 3: + lines.append("- No explicit profile constraints were set.") + if endpoint_env: + lines.append( + f"- Endpoint env keys available in the agent environment: {', '.join(endpoint_env)}" + ) + else: + lines.append("- Endpoint env keys: none") + return "\n".join(lines) + + +def _format_constraint_value(value: Any) -> str: + if isinstance(value, enum.Enum): + return str(value.value) + if hasattr(value, "name") and getattr(value, "project", None) is None: + return str(value.name) + if hasattr(value, "json"): + return json.dumps(json.loads(value.json()), sort_keys=True) + if hasattr(value, "value"): + return str(value.value) + return str(value) + + +def _start_agent_subprocess_detached( + workspace: _AgentWorkspace, + request: dict[str, Any], +) -> int: + cmd = _build_claude_command(request) + workspace.artifacts.mark_running() + stdout_path = workspace.work_dir / _AGENT_STDOUT_LOG_NAME + stderr_path = workspace.work_dir / _AGENT_STDERR_LOG_NAME + stdout_path.parent.mkdir(parents=True, exist_ok=True) + with ( + stdout_path.open("ab") as stdout, + stderr_path.open("ab") as stderr, + ): + proc = subprocess.Popen( + cmd, + cwd=workspace.work_dir, + env=request["env"], + stdout=stdout, + stderr=stderr, + close_fds=True, + start_new_session=True, + ) + _write_agent_process_state(workspace, proc.pid) + _write_trace_record( + workspace, + { + "type": "agent-process-started", + "pid": proc.pid, + "host": _get_process_host(), + }, + ) + return proc.pid + + +async def _run_agent_in_subprocess( + workspace: _AgentWorkspace, + request: dict[str, Any], +) -> _AgentRunnerResult: + cmd = _build_claude_command(request) + workspace.artifacts.mark_running() + progress_tailer = _AgentProgressLogTailer(workspace) + progress_task = asyncio.create_task(progress_tailer.run()) + proc: asyncio.subprocess.Process | None = None + try: + proc = await asyncio.create_subprocess_exec( + *cmd, + cwd=workspace.work_dir, + env=request["env"], + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + _write_agent_process_state(workspace, proc.pid) + assert proc.stdout is not None + assert proc.stderr is not None + stdout_task = asyncio.create_task(_read_agent_stdout(proc.stdout, workspace)) + stderr_task = asyncio.create_task(_read_agent_stderr(proc.stderr, workspace)) + stdout_output, stderr_output, returncode = await asyncio.gather( + stdout_task, + stderr_task, + proc.wait(), + ) + finally: + progress_task.cancel() + with suppress(asyncio.CancelledError): + await progress_task + await progress_tailer.flush(include_partial=True) + if proc is not None and proc.returncode is not None: + _clear_agent_process_state(workspace, proc.pid) + process_output = _merge_process_outputs(stdout_output, stderr_output) + _write_trace_record( + workspace, + { + "type": "agent-process", + "returncode": returncode, + }, + ) + await _flush_agent_logs(workspace) + if process_output.report_data is None: + process_output.report_data = _load_final_report_artifact(workspace) + if process_output.report_data is None: + if process_output.result_error is not None: + return _AgentRunnerResult(error="Server agent failed before returning a final report") + if returncode not in (0, None): + return _AgentRunnerResult( + error=( + "Server agent process exited without a final report " + f"(return code {returncode})" + ) + ) + return _AgentRunnerResult(error="Server agent process exited without a final report") + try: + return _AgentRunnerResult(report=AgentFinalReport.parse_obj(process_output.report_data)) + except ValidationError as e: + return _AgentRunnerResult(error=f"Server agent returned an invalid final report: {e}") + + +def _build_claude_command(request: dict[str, Any]) -> list[str]: + options = request["options"] + cmd = [ + _get_claude_executable() or "claude", + "-p", + "--output-format", + "stream-json", + "--verbose", + "--tools", + options.get("tools", "default"), + "--allowedTools", + options["allowed_tools"], + "--disallowedTools", + options["disallowed_tools"], + "--permission-mode", + "bypassPermissions", + "--model", + options["model"], + "--json-schema", + json.dumps(options["json_schema"]), + request["prompt"], + ] + if not _should_use_existing_claude_auth(): + cmd[2:2] = ["--bare"] + else: + cmd[2:2] = ["--setting-sources", "project,local"] + if settings.AGENT_CLAUDE_EFFORT is not None: + cmd[2:2] = ["--effort", settings.AGENT_CLAUDE_EFFORT] + if options["max_turns"] is not None: + cmd[2:2] = ["--max-turns", str(options["max_turns"])] + return cmd + + +def _get_claude_executable() -> Optional[str]: + if settings.AGENT_CLAUDE_PATH is not None: + return shutil.which(settings.AGENT_CLAUDE_PATH) + return shutil.which("claude") + + +def _should_use_existing_claude_auth() -> bool: + return settings.AGENT_CLAUDE_USE_EXISTING_AUTH and not settings.AGENT_ANTHROPIC_API_KEY + + +def _warn_if_using_existing_claude_auth() -> None: + global _existing_auth_warning_logged + if not _should_use_existing_claude_auth() or _existing_auth_warning_logged: + return + logger.warning( + "DSTACK_AGENT_CLAUDE_USE_EXISTING_AUTH=1 is enabled. This mode is intended " + "only for local development. Production servers must set " + "DSTACK_AGENT_ANTHROPIC_API_KEY. Claude will run without --bare and may read " + "the server user's Claude CLI auth and settings." + ) + _existing_auth_warning_logged = True + + +async def _read_agent_stdout( + stream: asyncio.StreamReader, + workspace: _AgentWorkspace, +) -> _AgentProcessOutput: + return await _read_agent_stream(stream, workspace, stream_name="stdout") + + +async def _read_agent_stderr( + stream: asyncio.StreamReader, + workspace: _AgentWorkspace, +) -> _AgentProcessOutput: + return await _read_agent_stream(stream, workspace, stream_name="stderr") + + +async def _read_agent_stream( + stream: asyncio.StreamReader, + workspace: _AgentWorkspace, + stream_name: str, +) -> _AgentProcessOutput: + output = _AgentProcessOutput() + while True: + line_bytes = await stream.readline() + if not line_bytes: + return output + line = line_bytes.decode(errors="replace") + await _process_agent_stream_line( + line=line, + workspace=workspace, + stream_name=stream_name, + output=output, + ) + + +async def _read_agent_stream_file( + *, + workspace: _AgentWorkspace, + path: Path, + offset: int, + stream_name: str, +) -> tuple[_AgentProcessOutput, int]: + output = _AgentProcessOutput() + lines, new_offset = _read_complete_file_lines(path=path, offset=offset) + for line in lines: + await _process_agent_stream_line( + line=line, + workspace=workspace, + stream_name=stream_name, + output=output, + ) + return output, new_offset + + +async def _process_agent_stream_line( + *, + line: str, + workspace: _AgentWorkspace, + stream_name: str, + output: _AgentProcessOutput, +) -> None: + output.stdout_tail = _append_bounded_output(output.stdout_tail, line) + if not line.strip(): + return + try: + message = json.loads(line) + except json.JSONDecodeError: + message = {"type": "raw-output", "stream": stream_name, "line": line.rstrip("\r\n")} + _write_trace_record(workspace, message) + return + if stream_name != "stdout": + message.setdefault("stream", stream_name) + _write_trace_record(workspace, message) + workspace.artifacts.record_stream_message(message) + _update_agent_process_output(output, message) + + +def _read_complete_file_lines(*, path: Path, offset: int) -> tuple[list[str], int]: + if not path.exists(): + return [], offset + size = path.stat().st_size + if size < offset: + offset = 0 + with path.open("rb") as f: + f.seek(offset) + data = f.read() + if not data: + return [], offset + newline_pos = max(data.rfind(b"\n"), data.rfind(b"\r")) + if newline_pos == -1: + return [], offset + complete = data[: newline_pos + 1] + new_offset = offset + len(complete) + lines = [line.decode(errors="replace") for line in complete.splitlines(keepends=True)] + return lines, new_offset + + +class _AgentProgressLogTailer: + def __init__(self, workspace: _AgentWorkspace) -> None: + self._workspace = workspace + self._offset = ( + workspace.progress_path.stat().st_size if workspace.progress_path.exists() else 0 + ) + self._partial = "" + + async def run(self) -> None: + while True: + await self.flush() + await asyncio.sleep(_AGENT_PROGRESS_POLL_SECONDS) + + async def flush(self, *, include_partial: bool = False) -> None: + path = self._workspace.progress_path + if not path.exists(): + return + if path.stat().st_size < self._offset: + self._offset = 0 + self._partial = "" + with path.open("r", encoding="utf-8", errors="replace") as f: + f.seek(self._offset) + text = f.read() + self._offset = f.tell() + if not text and not (include_partial and self._partial): + return + self._partial += text + lines = self._partial.splitlines(keepends=True) + self._partial = "" + for line in lines: + if not include_partial and not line.endswith(("\n", "\r")): + self._partial = line + continue + await self._write_progress_line(line.rstrip("\r\n")) + + async def _write_progress_line(self, line: str) -> None: + if not line.strip(): + return + message = _parse_agent_progress_log_message(line) + if message is None: + _write_trace_record( + self._workspace, + {"type": "agent-progress-ignored", "line": line}, + ) + return + await _write_agent_log(self._workspace, message) + + +async def _flush_agent_progress_logs( + workspace: _AgentWorkspace, + *, + offset: int, + include_partial: bool, +) -> int: + path = workspace.progress_path + lines, new_offset = _read_complete_file_lines(path=path, offset=offset) + if include_partial and path.exists(): + size = path.stat().st_size + if size > new_offset: + with path.open("rb") as f: + f.seek(new_offset) + data = f.read() + if data: + lines.append(data.decode(errors="replace")) + new_offset = size + for line in lines: + await _write_agent_progress_line(workspace, line.rstrip("\r\n")) + await _flush_agent_logs(workspace) + return new_offset + + +async def _write_agent_progress_line(workspace: _AgentWorkspace, line: str) -> None: + if not line.strip(): + return + message = _parse_agent_progress_log_message(line) + if message is None: + _write_trace_record( + workspace, + {"type": "agent-progress-ignored", "line": line}, + ) + return + await _write_agent_log(workspace, message) + + +def _parse_agent_progress_log_message(line: str) -> Optional[str]: + try: + parsed = json.loads(line) + except json.JSONDecodeError: + message = line.strip() + return message or None + return _format_agent_progress_log_message(parsed) + + +def _format_agent_progress_log_message(record: Any) -> Optional[str]: + if isinstance(record, str): + message = record.strip() + return message or None + if not isinstance(record, dict): + return None + message = record.get("message") + if not isinstance(message, str) or not message.strip(): + return None + return message.strip() + + +def _update_agent_process_output( + output: _AgentProcessOutput, + message: dict[str, Any], +) -> None: + if message.get("type") != "result": + return + if message.get("is_error"): + output.result_error = message.get("result") or "Claude agent failed" + structured_output = message.get("structured_output") + if isinstance(structured_output, dict): + output.report_data = structured_output + return + report_data = _parse_report_json(message.get("result")) + if report_data is not None: + output.report_data = report_data + + +def _append_bounded_output(existing: str, addition: str) -> str: + combined = existing + addition + if len(combined) <= _MAX_CAPTURED_OUTPUT_CHARS: + return combined + return combined[-_MAX_CAPTURED_OUTPUT_CHARS:] + + +def _merge_process_outputs(*outputs: _AgentProcessOutput) -> _AgentProcessOutput: + merged = _AgentProcessOutput() + for output in outputs: + if output.report_data is not None: + merged.report_data = output.report_data + if output.result_error is not None: + merged.result_error = output.result_error + merged.stdout_tail = _append_bounded_output(merged.stdout_tail, output.stdout_tail) + return merged + + +async def _write_agent_log(workspace: _AgentWorkspace, message: str) -> None: + if workspace.log_writer is None: + return + message = _redact(message, workspace.redacted_values) + await workspace.log_writer.write(message) + + +async def _flush_agent_logs(workspace: _AgentWorkspace) -> None: + if workspace.log_writer is None: + return + await workspace.log_writer.flush() + + +def _truncate_log_message(message: str) -> str: + if len(message) <= _MAX_AGENT_LOG_MESSAGE_CHARS: + return message + return message[: _MAX_AGENT_LOG_MESSAGE_CHARS - 15].rstrip() + "\n... truncated" + + +def _parse_report_json(value: Any) -> Optional[dict[str, Any]]: + if not isinstance(value, str): + return None + try: + parsed = json.loads(value) + except json.JSONDecodeError: + return None + if not isinstance(parsed, dict): + return None + return parsed + + +def _load_final_report_artifact(workspace: _AgentWorkspace) -> Optional[dict[str, Any]]: + path = workspace.work_dir / "final_report.json" + if not path.exists(): + return None + try: + parsed = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + logger.warning("Endpoint agent final_report.json is not valid JSON: %s", path) + return None + if not isinstance(parsed, dict): + logger.warning("Endpoint agent final_report.json is not an object: %s", path) + return None + return parsed + + +def _load_final_report(workspace: _AgentWorkspace) -> Optional[AgentFinalReport]: + report_data = _load_final_report_artifact(workspace) + if report_data is None: + return None + try: + return AgentFinalReport.parse_obj(report_data) + except ValidationError: + logger.warning( + "Endpoint agent final_report.json does not match the report schema: %s", + workspace.work_dir / "final_report.json", + exc_info=True, + ) + return None + + +def _write_agent_process_state(workspace: _AgentWorkspace, pid: int) -> None: + path = workspace.work_dir / _AGENT_PROCESS_STATE_NAME + data = { + "pid": pid, + "pgid": pid, + "host": _get_process_host(), + "started_at": _utcnow_iso(), + } + path.write_text(json.dumps(data, indent=2) + "\n", encoding="utf-8") + + +def _clear_agent_process_state(workspace: _AgentWorkspace, pid: int) -> None: + path = workspace.work_dir / _AGENT_PROCESS_STATE_NAME + if not path.exists(): + return + try: + data = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + path.unlink(missing_ok=True) + return + if data.get("pid") == pid: + path.unlink(missing_ok=True) + + +def _get_running_agent_process_pid( + workspace: _AgentWorkspace, + agent_session: Optional[EndpointAgentSessionModel] = None, +) -> Optional[int]: + state = _load_agent_process_state(workspace, agent_session=agent_session) + if state is None: + return None + if state.host not in (None, _get_process_host()): + return state.pid + if _is_process_group_running(state.pgid): + return state.pid + _clear_agent_process_state(workspace, state.pid) + return None + + +async def _abort_agent_session_process( + *, + agent_session: EndpointAgentSessionModel, + workspace: _AgentWorkspace, +) -> bool: + state = _load_agent_process_state(workspace, agent_session=agent_session) + if state is None: + return True + if state.host not in (None, _get_process_host()): + logger.info( + "Endpoint agent process %s for endpoint %s is running on host %s; " + "waiting for that host to abort it", + state.pid, + workspace.endpoint_name, + state.host, + ) + return False + if not _is_process_group_running(state.pgid): + _clear_agent_process_state(workspace, state.pid) + return True + + logger.info( + "Stopping endpoint agent process group %s for endpoint %s", + state.pgid, + workspace.endpoint_name, + ) + _write_trace_record( + workspace, + { + "type": "agent-process-abort-requested", + "pid": state.pid, + "pgid": state.pgid, + "host": state.host, + }, + ) + if not await _terminate_process_group(state.pgid): + return False + _clear_agent_process_state(workspace, state.pid) + return True + + +async def _terminate_process_group(pgid: int) -> bool: + if not _send_process_group_signal(pgid, signal.SIGTERM): + return True + await asyncio.sleep(_AGENT_PROCESS_ABORT_GRACE_SECONDS) + if not _is_process_group_running(pgid): + return True + _send_process_group_signal(pgid, _get_kill_signal()) + await asyncio.sleep(0) + return not _is_process_group_running(pgid) + + +def _send_process_group_signal(pgid: int, sig: signal.Signals) -> bool: + killpg: Optional[Callable[[int, signal.Signals], None]] = getattr(os, "killpg", None) + try: + if killpg is not None: + killpg(pgid, sig) + else: + os.kill(pgid, sig) + except ProcessLookupError: + return False + except PermissionError: + logger.warning("Permission denied while signaling process group %s", pgid) + return True + return True + + +def _load_agent_process_state( + workspace: _AgentWorkspace, + agent_session: Optional[EndpointAgentSessionModel] = None, +) -> Optional[_AgentProcessState]: + path = workspace.work_dir / _AGENT_PROCESS_STATE_NAME + if path.exists(): + try: + data = json.loads(path.read_text(encoding="utf-8")) + except json.JSONDecodeError: + path.unlink(missing_ok=True) + return None + state = _parse_agent_process_state(data) + if state is not None: + return state + path.unlink(missing_ok=True) + if agent_session is None or agent_session.pid is None: + return None + return _AgentProcessState( + pid=agent_session.pid, + pgid=agent_session.pid, + host=agent_session.process_host, + ) + + +def _parse_agent_process_state(data: Any) -> Optional[_AgentProcessState]: + if not isinstance(data, dict): + return None + pid = data.get("pid") + if not isinstance(pid, int): + return None + pgid = data.get("pgid", pid) + if not isinstance(pgid, int): + pgid = pid + host = data.get("host") + if not isinstance(host, str): + host = None + return _AgentProcessState(pid=pid, pgid=pgid, host=host) + + +def _get_process_host() -> str: + return socket.gethostname() + + +def _is_process_running(pid: int) -> bool: + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + except PermissionError: + return True + return True + + +def _is_process_group_running(pgid: int) -> bool: + killpg: Optional[Callable[[int, int], None]] = getattr(os, "killpg", None) + if killpg is None: + return _is_process_running(pgid) + try: + killpg(pgid, 0) + except ProcessLookupError: + return False + except PermissionError: + live_process_found = _has_non_zombie_process_in_group(pgid) + return True if live_process_found is None else live_process_found + return True + + +def _get_kill_signal() -> signal.Signals: + return getattr(signal, "SIGKILL", signal.SIGTERM) + + +def _has_non_zombie_process_in_group(pgid: int) -> Optional[bool]: + try: + result = subprocess.run( + ["ps", "-axo", "pgid=,stat="], + capture_output=True, + text=True, + timeout=2, + check=False, + ) + except (OSError, subprocess.TimeoutExpired): + return None + if result.returncode != 0: + return None + found_group = False + for line in result.stdout.splitlines(): + parts = line.split(None, 1) + if len(parts) != 2: + continue + try: + line_pgid = int(parts[0]) + except ValueError: + continue + if line_pgid != pgid: + continue + found_group = True + stat = parts[1].strip() + if stat and not stat.startswith("Z"): + return True + return False if found_group else False + + +def _write_trace_record(workspace: _AgentWorkspace, data: dict[str, Any]) -> None: + if workspace.trace_path is None: + return + workspace.trace_path.parent.mkdir(parents=True, exist_ok=True) + line = json.dumps(_redact(data, workspace.redacted_values), default=str) + with workspace.trace_path.open("a", encoding="utf-8") as f: + f.write(line + "\n") + + +def _get_redacted_values(values: Sequence[str]) -> tuple[str, ...]: + return tuple(sorted({value for value in values if value}, key=len, reverse=True)) + + +def _get_sensitive_inherited_env_values(env: dict[str, str]) -> list[str]: + return [env[name] for name in _INHERITED_ENV_NAMES if name != "PATH" and name in env] + + +def _redact(value: Any, redacted_values: Sequence[str]) -> Any: + if isinstance(value, str): + for redacted_value in redacted_values: + value = value.replace(redacted_value, _REDACTION) + return value + if isinstance(value, dict): + return {k: _redact(v, redacted_values) for k, v in value.items()} + if isinstance(value, list): + return [_redact(item, redacted_values) for item in value] + return value + + +def _is_debug_trace_enabled() -> bool: + return settings.LOG_LEVEL == "DEBUG" or settings.ROOT_LOG_LEVEL == "DEBUG" + + +def _utcnow_iso() -> str: + return datetime.now(timezone.utc).isoformat() diff --git a/src/dstack/_internal/server/services/endpoints/agent/report.py b/src/dstack/_internal/server/services/endpoints/agent/report.py new file mode 100644 index 0000000000..a601869b79 --- /dev/null +++ b/src/dstack/_internal/server/services/endpoints/agent/report.py @@ -0,0 +1,50 @@ +import uuid +from typing import Optional + +from pydantic import root_validator + +from dstack._internal.core.models.common import CoreModel + +AGENT_FINAL_REPORT_JSON_SCHEMA = { + "type": "object", + "properties": { + "success": {"type": "boolean"}, + "run_id": {"type": "string"}, + "run_name": {"type": "string"}, + "service_yaml": {"type": "string"}, + "base": {"type": "string"}, + "model": {"type": "string"}, + "failure_summary": {"type": "string"}, + }, + "required": ["success"], + "additionalProperties": False, +} + + +class AgentFinalReport(CoreModel): + # TODO: Keep this contract intentionally minimal until real agent runs show + # which validation evidence is actually useful and reliable. + success: bool + run_id: Optional[uuid.UUID] = None + run_name: Optional[str] = None + service_yaml: Optional[str] = None + base: Optional[str] = None + model: Optional[str] = None + failure_summary: Optional[str] = None + + @root_validator + def _validate_report(cls, values: dict) -> dict: + success = values.get("success") + if success: + if values.get("run_id") is None: + raise ValueError("successful agent report must include run_id") + if not values.get("service_yaml"): + raise ValueError("successful agent report must include service_yaml") + if not values.get("base") or not values["base"].strip(): + raise ValueError("successful agent report must include base") + if not values.get("model") or not values["model"].strip(): + raise ValueError("successful agent report must include model") + else: + if not values.get("failure_summary"): + raise ValueError("failed agent report must include failure_summary") + return values diff --git a/src/dstack/_internal/server/services/endpoints/agent/resources/system_prompt.md b/src/dstack/_internal/server/services/endpoints/agent/resources/system_prompt.md new file mode 100644 index 0000000000..3c74417d31 --- /dev/null +++ b/src/dstack/_internal/server/services/endpoints/agent/resources/system_prompt.md @@ -0,0 +1,251 @@ +# Objective + +You are the endpoint deployment agent for dstack. Produce one final dstack +service. Report success only after that service answers a real request using +`service_model_name` from `Endpoint context:` through the dstack service URL. + +# Requested Model + +The `Endpoint context:` block contains either `model_repo` or `base_model`. + +- If it contains `model_repo`, deploy that repo/path exactly. +- If it contains `base_model`, choose a repo/path compatible with `base_model` + that best fits performance and hardware within the endpoint constraints, + allowed fleets, backends, and offers. A variant can be the base repo itself, a + different precision or quantization, or another trusted repo compatible with + `base_model`. + +Use the real `dstack` CLI and shell commands in this workspace. Load and follow +`/dstack` for dstack CLI/YAML syntax. Load and follow `/dstack-prototyping` for +how to test a model-serving recipe with tasks before verifying it as a service. +The skill files are installed under `.claude/skills` if you need to inspect +them directly. + +# Progress + +Write endpoint progress with: + +```bash +progress "message for dstack endpoint logs" +``` + +The helper appends to `progress.jsonl`; the server shows only the message text. + +Write progress before your first investigation action, whenever you submit a +run, when new evidence changes the plan, when model verification succeeds or +fails, and before the final report. Messages should explain what you tried, what +happened, and what you will do next. Do not put raw YAML, command output, long +tables, traces, or secrets in progress. + +For `dstack endpoint logs`, write a short progress message for every meaningful +choice or action. Each message should say what you did or chose, why, what +evidence you used, what happened, and what you will do next. Include the run +name when a run is involved. Include the fleet or backend when a fleet or +backend is involved. Do not use generic phase labels as the message. + +# Workspace Files + +`submissions.jsonl` is append-only. For every dstack task or service you submit, +append one JSON line when you submit it and more JSON lines when you learn its +run id, status, URL, or final outcome. + +Use these fields: + +- `event`: `submit`, `update`, or `final` +- `name`: submitted dstack run name +- `type`: `task` or `service` +- `status`: current known status +- `config_path`: YAML file path, when applicable +- `run_id`: run id, when known +- `reason`: why this run exists or why its status changed +- `service_url`: service URL, when known + +Example: + +```json +{"event":"submit","name":"qwen-endpoint-1","type":"task","status":"submitted","config_path":"qwen-endpoint-1.dstack.yml","reason":"test the serving image and local model request on an allowed fleet"} +{"event":"update","name":"qwen-endpoint-1","type":"task","status":"running","run_id":"..."} +``` + +Do not rewrite previous `submissions.jsonl` lines; the latest line for a run is +the current record. Stop runs you no longer need unless they are still needed for +attach/SSH debugging, logs, or backend diagnosis. + +After stopping a task or service, do not wait for it to disappear from +`dstack ps`. Confirm that the run reached a terminal status such as `stopped`, +`terminated`, or `failed`, then continue. + +# Run Names + +Do not use the endpoint name itself as a submitted run name. + +Use only `-` for dstack runs submitted for +this endpoint. The first submitted run is `-1`, the second is +`-2`, and so on. Do not add framework, hardware, role, or purpose +suffixes to run names; record those details in workspace files and progress. + +# Endpoint Constraints + +Use existing allowed fleets only. Do not create, delete, apply, or edit fleets, +including `nodes`, `target`, `idle_duration`, backends, resources, max nodes, or +ownership. + +If the endpoint config lists fleets, use only those fleets. Otherwise, use the +existing project/imported fleets supplied in the request. + +The request also lists fixed endpoint constraints such as max price, spot +policy, backends, regions, instance types, fleets, env keys, tags, or backend +options. Do not submit a task or service that conflicts with those values. + +Put each fixed constraint that has a dstack service YAML field into the final +service YAML. For example, if the endpoint is limited to a fleet, max price, or +spot policy, the final `service_yaml` should include the corresponding `fleets`, +`max_price`, or `spot_policy` fields. Use `/dstack` for exact field names. + +If you cannot submit a useful task or service within the allowed fleets and +constraints, write a failed `final_report.json`. The `failure_summary` should +say which fleet or constraint blocked the run and what the user/admin would need +to change. + +# Task Usage + +Use `/dstack-prototyping` to learn how to use tasks. Keep in mind that using a +task is a must. This means `sleep infinity` and directly attaching inside the +task via SSH to run commands as the only way to use tasks. If there is any +ambiguity, follow `/dstack-prototyping` and do nothing that contradicts it. + +# Backend/Fleet Selection, Idle Duration, Instance Volumes + +Use only the fleets allowed by the endpoint request. If the endpoint request +also has constraints such as `backends`, `regions`, `instance_types`, +`spot_policy`, or `max_price`, apply them when running `dstack offer` if the CLI +has matching flags, and always apply them when submitting runs. + +## Backend And Fleet Choice + +Use `/dstack-prototyping` to learn how to select backends and fleets. + +Follow `/dstack-prototyping` skill on using a VM-based backend, Kubernetes backend, or SSH fleet that supports idle instances/instance volumes if there is such an option for the required GPU class. + +If backend allows (see above), use instance volumes to mount cache and model weights between runs. + +## Validating Offers + +When selecting backends/fleets and evaluating hardware that can be used to run +the model, use `dstack offer --json` and pass `--fleet` explicitly. If the +endpoint request has `fleets`, pass those fleets. Otherwise, pass every existing +fleet. If `--fleet` is not passed, `dstack offer` can show offers that are not +applicable to the fleets allowed by the endpoint request. Use these offers when +selecting fleet, backend, and hardware. + +Choosing fleet/backend is gated by classifying each against `https://dstack.ai/docs/concepts/backends.md`. VM-based backends are listed under `## VM-based` (they support idle instances and instance volumes). Kubernetes backend is listed under `## Container-based`, but supports instance volumes and thus is preferred over other container-based backends. +SSH fleets can be treated as VM-based backends as they support both idle instances (its equivalent) and instance volumes. + +## Model And Compute Fit + +If `Endpoint context:` contains `model_repo`, choose fleet/backend/hardware that +can run `model_repo` within endpoint constraints. + +If `Endpoint context:` contains `base_model`, choose the repo/path and compute +together. You may pick the base repo or a compatible variant that fits the +allowed fleets, backends, and offers. + +If a task or service shows that the selected repo/path is a bad fit and +`Endpoint context:` contains `base_model`, pick another compatible variant if +available and test it in a task before submitting another service. + +## Submitting run + +When submitting a task or service, pass exact `fleets`, `backends`, and an +intentional `resources` range based on the choice made from offers, so the run +does not land outside the intended fleet/backend/hardware. + +## Endpoint Logs + +Endpoint logs should explain meaningful decisions and actions. + +When choosing repo/path for `base_model`, fleet, backend, or hardware, write a +log message that includes: + +- `service_model_name` and selected repo/path, when they differ; +- the selected fleet, backend, and resource range; +- the offer/docs evidence used for the choice; +- the viable alternatives not selected and the exact reason; +- the fleet, backend, and resources that will be used in the submitted YAML. +- how the selected and rejected fleets/backends were classified; + +Backend-choice log example (provide the same level of explanation for repo/path, +fleet, and hardware choice): + +"I chose backend ... on fleet ... because ...; I did not choose ... because ...; +I will submit YAML with fleets=..., backends=..., resources=...." + +# Final Service + +The final `service_yaml` is what dstack will save as the endpoint recipe. It +must contain the full service config: `type: service`, the final run name, the +service model name, image/commands/port, resources, env references, and the +fixed endpoint constraints that apply to service YAML. + +Set final `service_yaml.name` to the final service run name. +Set the final service model name to `service_model_name` from `Endpoint +context:`. If final `service_yaml.model` is a string, set it to +`service_model_name`; if it is an object, set `model.name` to +`service_model_name`. + +Choose service resources from the least restrictive requirements supported by +the evidence, not from the exact machine that happened to run. For example, if +the model worked on an A40 but the evidence only says it needs an NVIDIA GPU +with at least 16GB memory, use that broader requirement. Use an exact GPU, +region, backend, or instance type only when the endpoint constraints require it +or the tested recipe depends on that exact choice. + +Use run status to know whether the final service is still starting, running, or +failed. Use logs to understand failures. When dstack exposes the service URL, +send a real model request through that URL. + +Do not write a successful `final_report.json` until the final service answers a +request using `service_model_name` from `Endpoint context:` through the dstack +service URL. + +# Secrets + +Do not print, cat, copy, or summarize `~/.dstack/config.yml`, tokens, secrets, +or environment variables. The dstack CLI is already configured; use it directly. +For final service verification, build absolute URLs from +`DSTACK_ENDPOINT_SERVER_URL` when needed and use +`DSTACK_ENDPOINT_BEARER_TOKEN` only as the bearer token in the request header. + +# Resume + +On startup or resume, inspect `final_report.json` and `submissions.jsonl` +before submitting anything new. If `final_report.json` already reports success +for the current endpoint configuration, do not submit another run. Otherwise use +the next run number and record why the old run is not enough. + +# Final Report + +On success, write `final_report.json`, then return the structured final report. +The successful report must include: + +- `base`: the base model repo for the deployed model +- `model`: the exact repo/path loaded by the final service command + +Set `model` to the exact repo/path loaded by the service command. + +Set `base` as follows: + +- If `Endpoint context:` contains `base_model`, set `base` to `base_model`. +- If `Endpoint context:` contains `model_repo`, inspect the repo metadata, model + card, config, or another reliable source to identify the base model repo. +- If `model_repo` is itself the base model repo, set `base` to `model_repo`. +- Do not infer `base` only from the repo name. + +On failure, write `final_report.json` with a useful `failure_summary`, then +return the structured final report. + +`final_report.json` must contain only the schema fields: `success`, `run_id`, +`run_name`, `service_yaml`, `base`, `model`, and `failure_summary`. + +Stop after one correct service is verified. P/D disaggregation is not covered by +the current endpoint agent or `/dstack-prototyping` skill. diff --git a/src/dstack/_internal/server/services/endpoints/names.py b/src/dstack/_internal/server/services/endpoints/names.py new file mode 100644 index 0000000000..7552fc9022 --- /dev/null +++ b/src/dstack/_internal/server/services/endpoints/names.py @@ -0,0 +1,12 @@ +from typing import Optional + +_DSTACK_RESOURCE_NAME_MAX_LENGTH = 41 +_SERVING_RUN_SUFFIX = "-serving" + + +def get_endpoint_serving_run_name(endpoint_name: Optional[str]) -> Optional[str]: + if endpoint_name is None: + return None + if len(endpoint_name) + len(_SERVING_RUN_SUFFIX) <= _DSTACK_RESOURCE_NAME_MAX_LENGTH: + return f"{endpoint_name}{_SERVING_RUN_SUFFIX}" + return endpoint_name diff --git a/src/dstack/_internal/server/services/endpoints/planning.py b/src/dstack/_internal/server/services/endpoints/planning.py new file mode 100644 index 0000000000..2602ca92dc --- /dev/null +++ b/src/dstack/_internal/server/services/endpoints/planning.py @@ -0,0 +1,183 @@ +from dataclasses import dataclass +from typing import Optional, Sequence + +from sqlalchemy.ext.asyncio import AsyncSession + +from dstack._internal.core.errors import ServerClientError +from dstack._internal.core.models.configurations import ServiceConfiguration +from dstack._internal.core.models.endpoints import ( + EndpointConfiguration, +) +from dstack._internal.core.models.envs import EnvSentinel +from dstack._internal.core.models.profiles import ProfileParams +from dstack._internal.core.models.runs import JobPlan, RunPlan, RunSpec +from dstack._internal.server.models import ProjectModel, UserModel +from dstack._internal.server.services import runs as runs_services +from dstack._internal.server.services import users as users_services +from dstack._internal.server.services.endpoints.names import get_endpoint_serving_run_name +from dstack._internal.server.services.endpoints.presets import ( + EndpointPreset, + EndpointPresetRecipe, + EndpointPresetService, + get_endpoint_preset_service, +) +from dstack._internal.utils.logging import get_logger + +logger = get_logger(__name__) + + +@dataclass(frozen=True) +class EndpointPresetPlan: + """A preset selected by endpoint planning and the run plan computed from it.""" + + preset: EndpointPreset + recipe: EndpointPresetRecipe + run_plan: RunPlan + + +@dataclass(frozen=True) +class EndpointPresetPlanningResult: + """Preset planning result split into provisionable and no-offer matches.""" + + provisionable: Optional[EndpointPresetPlan] = None + unprovisionable: Optional[EndpointPresetPlan] = None + + +async def find_matching_preset_plan( + session: AsyncSession, + project: ProjectModel, + user: UserModel, + endpoint_name: Optional[str], + endpoint_configuration: EndpointConfiguration, + max_offers: int = 1, + preset_service: Optional[EndpointPresetService] = None, +) -> Optional[EndpointPresetPlan]: + result = await find_preset_planning_result( + session=session, + project=project, + user=user, + endpoint_name=endpoint_name, + endpoint_configuration=endpoint_configuration, + max_offers=max_offers, + preset_service=preset_service, + ) + return result.provisionable + + +async def find_preset_planning_result( + session: AsyncSession, + project: ProjectModel, + user: UserModel, + endpoint_name: Optional[str], + endpoint_configuration: EndpointConfiguration, + max_offers: int = 1, + preset_service: Optional[EndpointPresetService] = None, +) -> EndpointPresetPlanningResult: + if preset_service is None: + preset_service = get_endpoint_preset_service() + + endpoint_model = endpoint_configuration.model.api_model_name.lower() + presets = [ + preset + for preset in await preset_service.list_presets(project.name) + if preset.base.lower() == endpoint_model + ] + if not presets: + return EndpointPresetPlanningResult() + + exact_model_repo = endpoint_configuration.model.exact_repo + user = await _ensure_user_has_ssh_key(session=session, user=user) + first_unprovisionable_preset: Optional[EndpointPresetPlan] = None + for preset in presets: + for recipe in preset.recipes: + if exact_model_repo is not None and recipe.model != exact_model_repo: + continue + try: + run_spec = build_preset_run_spec( + endpoint_name=endpoint_name, + endpoint_configuration=endpoint_configuration, + recipe=recipe, + ) + _validate_run_spec_env_resolved(run_spec) + run_plan = await runs_services.get_plan( + session=session, + project=project, + user=user, + run_spec=run_spec, + max_offers=max_offers, + ) + except (ServerClientError, ValueError) as e: + logger.warning( + "Skipping endpoint preset %s recipe %s: %s", + preset.base, + recipe.id, + e, + ) + continue + preset_plan = EndpointPresetPlan(preset=preset, recipe=recipe, run_plan=run_plan) + if _run_plan_has_available_offers(run_plan.job_plans): + return EndpointPresetPlanningResult( + provisionable=preset_plan, + unprovisionable=first_unprovisionable_preset, + ) + if first_unprovisionable_preset is None: + first_unprovisionable_preset = preset_plan + return EndpointPresetPlanningResult(unprovisionable=first_unprovisionable_preset) + + +def build_preset_service_configuration( + endpoint_name: Optional[str], + endpoint_configuration: EndpointConfiguration, + recipe: EndpointPresetRecipe, +) -> ServiceConfiguration: + service_configuration = recipe.service.copy(deep=True) + service_configuration.name = get_endpoint_serving_run_name(endpoint_name) + service_configuration.env.update(endpoint_configuration.env) + for field in ProfileParams.__fields__: + value = getattr(endpoint_configuration, field) + if value is not None: + setattr(service_configuration, field, value) + return service_configuration + + +def build_preset_run_spec( + endpoint_name: Optional[str], + endpoint_configuration: EndpointConfiguration, + recipe: EndpointPresetRecipe, +) -> RunSpec: + service_configuration = build_preset_service_configuration( + endpoint_name=endpoint_name, + endpoint_configuration=endpoint_configuration, + recipe=recipe, + ) + return RunSpec( + run_name=service_configuration.name, + configuration=service_configuration, + ) + + +def _validate_run_spec_env_resolved(run_spec: RunSpec) -> None: + unresolved = [ + key for key, value in run_spec.configuration.env.items() if isinstance(value, EnvSentinel) + ] + if unresolved: + raise ValueError("preset env is unresolved: " + ", ".join(sorted(unresolved))) + + +async def _ensure_user_has_ssh_key( + session: AsyncSession, + user: UserModel, +) -> UserModel: + if user.ssh_public_key: + return user + refreshed_user = await users_services.refresh_ssh_key(session=session, actor=user) + if refreshed_user is None: + return user + return refreshed_user + + +def _run_plan_has_available_offers(job_plans: Sequence[JobPlan]) -> bool: + return bool(job_plans) and all( + any(offer.availability.is_available() for offer in job_plan.offers) + for job_plan in job_plans + ) diff --git a/src/dstack/_internal/server/services/endpoints/preset_building.py b/src/dstack/_internal/server/services/endpoints/preset_building.py new file mode 100644 index 0000000000..3d2d2054e5 --- /dev/null +++ b/src/dstack/_internal/server/services/endpoints/preset_building.py @@ -0,0 +1,139 @@ +from collections import defaultdict +from typing import Any, Optional + +from dstack._internal.core.models.configurations import ServiceConfiguration +from dstack._internal.core.models.instances import InstanceOfferWithAvailability, Resources +from dstack._internal.core.models.resources import ResourcesSpec +from dstack._internal.core.models.runs import JobStatus +from dstack._internal.server.models import JobModel, RunModel +from dstack._internal.server.services import jobs as jobs_services +from dstack._internal.server.services import runs as runs_services +from dstack._internal.server.services.endpoints.presets import ( + EndpointPreset, + EndpointPresetRecipe, + EndpointPresetValidation, + EndpointPresetValidationReplica, + make_endpoint_preset_recipe_id, + set_service_gpu_vendors_from_validations, +) +from dstack._internal.utils.common import format_mib_as_gb + + +def build_endpoint_preset_from_run( + run_model: RunModel, + base_model: Optional[str] = None, + recipe_model: Optional[str] = None, +) -> EndpointPreset: + run_spec = runs_services.get_run_spec(run_model) + configuration = run_spec.configuration + if not isinstance(configuration, ServiceConfiguration): + raise ValueError("endpoint preset can only be built from a service run") + if configuration.model is None: + raise ValueError("endpoint preset service must specify model") + if base_model is None: + base_model = configuration.model.name + if not base_model.strip(): + raise ValueError("endpoint preset base must be non-empty") + if recipe_model is None: + recipe_model = configuration.model.name + if not recipe_model.strip(): + raise ValueError("endpoint preset recipe must specify model") + + jobs_by_group = _get_current_registered_replica_jobs_by_group(run_model) + validation_replicas = [] + for replica_group in configuration.replica_groups: + group_name = replica_group.name + if group_name is None: + raise ValueError("endpoint preset cannot be built from unnamed service replicas") + group_jobs = jobs_by_group.get(group_name, []) + if not group_jobs: + raise ValueError( + f"endpoint preset cannot be built: replica group {group_name!r} " + "has no registered running replicas" + ) + validation_replicas.append( + EndpointPresetValidationReplica( + resources=[_get_job_resources_spec(job) for job in group_jobs], + ) + ) + + service = configuration.copy(deep=True) + service.name = None + validation = EndpointPresetValidation(replicas=validation_replicas) + set_service_gpu_vendors_from_validations( + service=service, + validations=[validation], + ) + return EndpointPreset( + base=base_model, + recipes=[ + EndpointPresetRecipe( + id=make_endpoint_preset_recipe_id(service), + model=recipe_model, + service=service, + validations=[validation], + ) + ], + ) + + +def _get_current_registered_replica_jobs_by_group( + run_model: RunModel, +) -> dict[str, list[JobModel]]: + jobs_by_group = defaultdict(list) + for job in sorted(run_model.jobs, key=lambda j: (j.replica_num, j.job_num)): + if job.deployment_num != run_model.deployment_num: + continue + if job.job_num != 0: + continue + if job.status != JobStatus.RUNNING or not job.registered: + continue + job_spec = jobs_services.get_job_spec(job) + jobs_by_group[job_spec.replica_group].append(job) + return jobs_by_group + + +def _get_job_resources_spec(job_model: JobModel) -> ResourcesSpec: + offer = _get_job_offer(job_model) + if offer is None: + raise ValueError("endpoint preset cannot be built without actual instance resources") + return _resources_spec_from_instance_resources(offer.instance.resources) + + +def _get_job_offer(job_model: JobModel) -> InstanceOfferWithAvailability | None: + job_runtime_data = jobs_services.get_job_runtime_data(job_model) + if job_runtime_data is not None and job_runtime_data.offer is not None: + return job_runtime_data.offer + instance = job_model.__dict__.get("instance") + if instance is not None and instance.offer is not None: + return InstanceOfferWithAvailability.__response__.parse_raw(instance.offer) + return None + + +def _resources_spec_from_instance_resources(resources: Resources) -> ResourcesSpec: + data: dict[str, Any] = { + "cpu": str(resources.cpus), + "memory": format_mib_as_gb(resources.memory_mib), + "disk": format_mib_as_gb(resources.disk.size_mib), + } + if resources.cpu_arch is not None: + data["cpu"] = f"{resources.cpu_arch.value}:{resources.cpus}" + if resources.gpus: + first_gpu = resources.gpus[0] + if any( + gpu.name != first_gpu.name + or gpu.memory_mib != first_gpu.memory_mib + or gpu.vendor != first_gpu.vendor + for gpu in resources.gpus + ): + raise ValueError("endpoint preset cannot be built from mixed-GPU instances") + data["gpu"] = { + "name": first_gpu.name, + "memory": format_mib_as_gb(first_gpu.memory_mib), + "count": len(resources.gpus), + } + if first_gpu.vendor is not None: + data["gpu"]["vendor"] = first_gpu.vendor.value + else: + data["gpu"] = 0 + return ResourcesSpec.parse_obj(data) diff --git a/src/dstack/_internal/server/services/endpoints/presets.py b/src/dstack/_internal/server/services/endpoints/presets.py new file mode 100644 index 0000000000..9361485d9b --- /dev/null +++ b/src/dstack/_internal/server/services/endpoints/presets.py @@ -0,0 +1,727 @@ +import hashlib +import json +import os +import re +import uuid +from abc import ABC, abstractmethod +from copy import deepcopy +from pathlib import Path +from typing import Any, NoReturn, Optional, Sequence + +import gpuhunt +import yaml +from pydantic import ValidationError, parse_obj_as + +from dstack._internal.core.models.configurations import ( + DEFAULT_REPLICA_GROUP_NAME, + ServiceConfiguration, +) +from dstack._internal.core.models.endpoint_presets import ( + EndpointPreset, + EndpointPresetRecipe, + EndpointPresetValidation, + EndpointPresetValidationReplica, +) +from dstack._internal.core.models.envs import EnvSentinel +from dstack._internal.core.models.profiles import ProfileParams +from dstack._internal.core.models.resources import CPUSpec, ResourcesSpec +from dstack._internal.server import settings +from dstack._internal.utils.common import run_async +from dstack._internal.utils.logging import get_logger + +_SECRET_ENV_PATTERN = re.compile(r"(token|key|secret|password)", re.IGNORECASE) +logger = get_logger(__name__) + + +class EndpointPresetService(ABC): + @abstractmethod + async def list_presets(self, project_name: str) -> list[EndpointPreset]: + pass + + @abstractmethod + async def get_preset(self, project_name: str, base: str) -> Optional[EndpointPreset]: + pass + + @abstractmethod + async def delete_preset(self, project_name: str, base: str) -> None: + pass + + @abstractmethod + async def save_preset( + self, + project_name: str, + preset: EndpointPreset, + comments: Optional[Sequence[str]] = None, + ) -> EndpointPreset: + """Merge and store a model-level preset.""" + pass + + +class LocalDirEndpointPresetService(EndpointPresetService): + def __init__(self, projects_dir: Path = settings.SERVER_PROJECTS_DIR_PATH) -> None: + self._projects_dir = projects_dir + + async def list_presets(self, project_name: str) -> list[EndpointPreset]: + return await run_async(self._list_presets, project_name) + + async def get_preset(self, project_name: str, base: str) -> Optional[EndpointPreset]: + return await run_async(self._get_preset, project_name, base) + + async def delete_preset(self, project_name: str, base: str) -> None: + return await run_async(self._delete_preset, project_name, base) + + async def save_preset( + self, + project_name: str, + preset: EndpointPreset, + comments: Optional[Sequence[str]] = None, + ) -> EndpointPreset: + return await run_async(self._save_preset, project_name, preset, comments or []) + + def _list_presets(self, project_name: str) -> list[EndpointPreset]: + return [preset for preset, _ in self._list_merged_presets(project_name)] + + def _get_preset(self, project_name: str, base: str) -> Optional[EndpointPreset]: + preset, _ = self._find_preset_by_base(project_name, base) + return preset + + def _delete_preset(self, project_name: str, base: str) -> None: + _, paths = self._find_preset_by_base(project_name, base) + if not paths: + raise FileNotFoundError(base) + for path in paths: + path.unlink() + + def _save_preset( + self, + project_name: str, + preset: EndpointPreset, + comments: Sequence[str], + ) -> EndpointPreset: + _validate_preset(preset) + presets_dir = self._get_project_presets_dir(project_name) + presets_dir.mkdir(parents=True, exist_ok=True) + + existing_preset, existing_paths = self._find_preset_by_base(project_name, preset.base) + if existing_preset is not None: + saved_preset = _merge_presets(existing_preset, preset, prefer_incoming=True) + path = existing_paths[0] + else: + saved_preset = preset + path = self._get_available_preset_path(presets_dir, preset.base) + + data = _preset_to_data(saved_preset) + content = _format_preset_comments(comments) + yaml.safe_dump(data, sort_keys=False) + self._replace_preset_file(presets_dir=presets_dir, path=path, content=content) + for duplicate_path in existing_paths[1:]: + duplicate_path.unlink(missing_ok=True) + loaded_preset = self._load_preset(path) + if loaded_preset is None: + raise ValueError(f"saved endpoint preset {path} could not be loaded") + return loaded_preset + + def _find_preset_by_base( + self, + project_name: str, + base: str, + ) -> tuple[Optional[EndpointPreset], list[Path]]: + base_key = base.lower() + for preset, paths in self._list_merged_presets(project_name): + if preset.base.lower() == base_key: + return preset, paths + return None, [] + + def _list_merged_presets(self, project_name: str) -> list[tuple[EndpointPreset, list[Path]]]: + presets_dir = self._get_project_presets_dir(project_name) + if not presets_dir.exists(): + return [] + + presets_by_base: dict[str, EndpointPreset] = {} + paths_by_base: dict[str, list[Path]] = {} + base_order: list[str] = [] + for path in self._iter_preset_paths(presets_dir): + preset = self._load_preset(path) + if preset is None: + continue + base_key = preset.base.lower() + existing = presets_by_base.get(base_key) + if existing is None: + presets_by_base[base_key] = preset + paths_by_base[base_key] = [path] + base_order.append(base_key) + continue + try: + presets_by_base[base_key] = _merge_presets(existing, preset) + except ValueError as e: + logger.warning("Skipping endpoint preset %s: %s", path, e) + continue + paths_by_base[base_key].append(path) + return [(presets_by_base[base], paths_by_base[base]) for base in base_order] + + def _iter_preset_paths(self, presets_dir: Path) -> list[Path]: + return [path for path in sorted(presets_dir.iterdir()) if path.suffix in [".yml", ".yaml"]] + + def _load_preset(self, path: Path) -> Optional[EndpointPreset]: + try: + with path.open("r") as f: + data = yaml.safe_load(f) + preset = _preset_from_data(data) + _validate_preset(preset) + return preset + except (OSError, ValidationError, ValueError, yaml.YAMLError) as e: + logger.warning("Skipping endpoint preset %s: %s", path, e) + return None + + def _get_project_presets_dir(self, project_name: str) -> Path: + return self._projects_dir / project_name / "presets" + + def _get_available_preset_path(self, presets_dir: Path, model: str) -> Path: + base_name = _slugify_model(model) + suffix = 0 + while True: + name = base_name if suffix == 0 else f"{base_name}-{suffix + 1}" + path = presets_dir / f"{name}.dstack.yml" + if not path.exists(): + return path + suffix += 1 + + def _replace_preset_file(self, presets_dir: Path, path: Path, content: str) -> None: + tmp_path = presets_dir / f".{path.name}.{uuid.uuid4().hex}.tmp" + tmp_path.write_text(content, encoding="utf-8") + try: + os.replace(tmp_path, path) + finally: + tmp_path.unlink(missing_ok=True) + + +_endpoint_preset_service: EndpointPresetService = LocalDirEndpointPresetService() + + +def get_endpoint_preset_service() -> EndpointPresetService: + return _endpoint_preset_service + + +def endpoint_preset_to_api_model(preset: EndpointPreset) -> EndpointPreset: + return EndpointPreset.parse_obj(preset.dict()) + + +def endpoint_preset_to_api_details(preset: EndpointPreset) -> EndpointPreset: + return EndpointPreset.parse_obj(preset.dict()) + + +def make_endpoint_preset_recipe_id(service: ServiceConfiguration) -> str: + service_data = _service_configuration_to_preset_data(service) + payload = json.dumps(service_data, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode()).hexdigest()[:8] + + +def _preset_from_data(data: Any) -> EndpointPreset: + if not isinstance(data, dict): + raise ValueError("preset must be a YAML object") + if data.get("type") != "endpoint-preset": + raise ValueError("preset must be an endpoint preset") + base = data.get("base", data.get("model")) + if not isinstance(base, str) or base == "": + raise ValueError("preset must specify a base") + legacy_model = data.get("model") + if isinstance(legacy_model, str) and "base" in data and legacy_model != base: + raise ValueError("preset base must match legacy model") + if "recipes" in data: + return EndpointPreset( + base=base, + recipes=_get_preset_recipes(data=data, base=base), + ) + return EndpointPreset( + base=base, + recipes=[_get_legacy_preset_recipe(data=data, base=base)], + ) + + +def _get_preset_recipes(data: dict[str, Any], base: str) -> list[EndpointPresetRecipe]: + raw_recipes = data.get("recipes") + if not isinstance(raw_recipes, list) or not raw_recipes: + raise ValueError("preset must specify non-empty recipes") + recipes = [_get_preset_recipe(raw_recipe, base=base) for raw_recipe in raw_recipes] + recipe_ids = [recipe.id for recipe in recipes] + if len(recipe_ids) != len(set(recipe_ids)): + raise ValueError("preset recipes must have unique ids") + return recipes + + +def _get_preset_recipe(raw_recipe: Any, base: str) -> EndpointPresetRecipe: + if not isinstance(raw_recipe, dict): + raise ValueError("preset recipes must be objects") + recipe_id = raw_recipe.get("id") + if not isinstance(recipe_id, str) or recipe_id == "": + raise ValueError("preset recipe must specify id") + recipe_model = raw_recipe.get("model", base) + if not isinstance(recipe_model, str) or not recipe_model.strip(): + raise ValueError("preset recipe must specify model") + service_data = _get_service_data(raw_recipe, base=base) + service = ServiceConfiguration.parse_obj(service_data) + validations = _get_validations(raw_recipe, service=service) + return EndpointPresetRecipe( + id=recipe_id, + model=recipe_model, + service=service, + validations=validations, + ) + + +def _get_legacy_preset_recipe(data: dict[str, Any], base: str) -> EndpointPresetRecipe: + service_data = _get_service_data(data, base=base, allow_missing_resources=True) + raw_groups = _get_legacy_replica_spec_groups(data=data, service_data=service_data) + _apply_legacy_replica_group_resources(service_data=service_data, groups=raw_groups) + service = ServiceConfiguration.parse_obj(service_data) + validation = EndpointPresetValidation( + replicas=[ + EndpointPresetValidationReplica( + resources=[ + ResourcesSpec.parse_obj(resources) + for resources in _get_legacy_group_tested_resources(group) + ] + ) + for group in raw_groups + ] + ) + return EndpointPresetRecipe( + id=make_endpoint_preset_recipe_id(service), + model=base, + service=service, + validations=[validation], + ) + + +def _get_service_data( + data: dict[str, Any], + base: str, + allow_missing_resources: bool = False, +) -> dict[str, Any]: + service_data = data.get("service") + if not isinstance(service_data, dict): + raise ValueError("preset recipe must specify a service object") + service_data = deepcopy(service_data) + if service_data.get("type") not in (None, "service"): + raise ValueError("preset service object must be a service configuration") + if "name" in service_data: + raise ValueError("preset service object must not specify name") + profile_fields = sorted(set(ProfileParams.__fields__) & set(service_data)) + if profile_fields: + raise ValueError( + "preset service object must not specify profile fields: " + ", ".join(profile_fields) + ) + if not allow_missing_resources: + _validate_service_has_resources(service_data) + service_data.setdefault("model", base) + service_data["type"] = "service" + configuration = ServiceConfiguration.parse_obj(service_data) + if configuration.model is None: + raise ValueError("preset service configuration must specify model") + if configuration.model.name.lower() != base.lower(): + raise ValueError("preset base must match the service model") + return service_data + + +def _validate_service_has_resources(service_data: dict[str, Any]) -> None: + replicas = service_data.get("replicas") + if isinstance(replicas, list): + for group in replicas: + if not isinstance(group, dict): + raise ValueError("preset service replica groups must be objects") + if "resources" not in group: + raise ValueError("preset service replica groups must specify resources") + return + if "resources" not in service_data: + raise ValueError("preset service object must specify resources") + + +def _get_validations( + raw_recipe: dict[str, Any], + service: ServiceConfiguration, +) -> list[EndpointPresetValidation]: + raw_validations = raw_recipe.get("validations") + if not isinstance(raw_validations, list) or not raw_validations: + raise ValueError("preset recipe must specify non-empty validations") + validations = [ + EndpointPresetValidation.parse_obj(raw_validation) for raw_validation in raw_validations + ] + _validate_validations(validations=validations, service=service) + return validations + + +def _validate_preset(preset: EndpointPreset) -> None: + if not preset.recipes: + raise ValueError("preset must specify non-empty recipes") + recipe_ids = [recipe.id for recipe in preset.recipes] + if len(recipe_ids) != len(set(recipe_ids)): + raise ValueError("preset recipes must have unique ids") + for recipe in preset.recipes: + if not recipe.model.strip(): + raise ValueError("preset recipe must specify model") + if recipe.service.model is None: + raise ValueError("preset recipe service must specify model") + if recipe.service.model.name.lower() != preset.base.lower(): + raise ValueError("preset base must match the recipe service model") + _validate_validations(validations=recipe.validations, service=recipe.service) + _validate_service_gpu_vendor_matches_validations( + service=recipe.service, + validations=recipe.validations, + ) + + +def set_service_gpu_vendors_from_validations( + service: ServiceConfiguration, + validations: list[EndpointPresetValidation], +) -> None: + for group_num, group in enumerate(service.replica_groups): + resources = group.resources + if resources is None or not _requires_gpu(resources): + continue + validation_vendor = _get_validation_group_gpu_vendor( + validations=validations, + group_num=group_num, + ) + if validation_vendor is None: + continue + if resources.gpu is None: + continue + if resources.gpu.vendor is not None and resources.gpu.vendor != validation_vendor: + raise ValueError("preset service GPU vendor does not match validation") + group_resources = _get_service_group_resources(service, group_num) + if group_resources.gpu is None: + continue + group_resources.gpu.vendor = validation_vendor + + +def _validate_service_gpu_vendor_matches_validations( + service: ServiceConfiguration, + validations: list[EndpointPresetValidation], +) -> None: + for group_num, group in enumerate(service.replica_groups): + resources = group.resources + if resources is None or not _requires_gpu(resources): + continue + if resources.gpu is None or resources.gpu.vendor is None: + continue + validation_vendor = _get_validation_group_gpu_vendor( + validations=validations, + group_num=group_num, + ) + if validation_vendor is not None and resources.gpu.vendor != validation_vendor: + raise ValueError("preset service GPU vendor does not match validation") + + +def _get_validation_group_gpu_vendor( + validations: list[EndpointPresetValidation], + group_num: int, +): + vendors = set() + for validation in validations: + for resources in validation.replicas[group_num].resources: + if resources.gpu is None or resources.gpu.count.min == 0: + continue + vendor = _get_validation_resources_gpu_vendor(resources) + if vendor is not None: + vendors.add(vendor) + if len(vendors) > 1: + raise ValueError("preset validations must not mix GPU vendors in a replica group") + return next(iter(vendors), None) + + +def _get_validation_resources_gpu_vendor(resources: ResourcesSpec): + gpu = resources.gpu + if gpu is None: + return None + if gpu.vendor is not None: + return gpu.vendor + if not gpu.name: + return None + vendors = {_get_gpu_name_vendor(name) for name in gpu.name} + vendors.discard(None) + if len(vendors) > 1: + raise ValueError("preset validations must not mix GPU vendors in a replica group") + return next(iter(vendors), None) + + +def _get_gpu_name_vendor(name: str): + if _is_known_gpu_name(name, gpuhunt.KNOWN_NVIDIA_GPUS): + return gpuhunt.AcceleratorVendor.NVIDIA + if _is_known_gpu_name(name, gpuhunt.KNOWN_AMD_GPUS): + return gpuhunt.AcceleratorVendor.AMD + if _is_known_gpu_name(name, gpuhunt.KNOWN_INTEL_ACCELERATORS): + return gpuhunt.AcceleratorVendor.INTEL + if _is_known_gpu_name(name, gpuhunt.KNOWN_TENSTORRENT_ACCELERATORS): + return gpuhunt.AcceleratorVendor.TENSTORRENT + if name.startswith("tpu-"): + return gpuhunt.AcceleratorVendor.GOOGLE + return None + + +def _is_known_gpu_name(name: str, known_gpus: list) -> bool: + return any(gpu.name.lower() == name.lower() for gpu in known_gpus) + + +def _get_service_group_resources( + service: ServiceConfiguration, + group_num: int, +) -> ResourcesSpec: + if isinstance(service.replicas, list): + resources = service.replicas[group_num].resources + else: + resources = service.resources + if resources is None: + raise ValueError("preset service object must specify resources") + return resources + + +def _requires_gpu(resources: ResourcesSpec) -> bool: + gpu = resources.gpu + if gpu is None: + return False + if gpu.count.max == 0: + return False + return gpu.count.min != 0 or gpu.count.max is not None + + +def _validate_validations( + validations: list[EndpointPresetValidation], + service: ServiceConfiguration, +) -> None: + expected_group_count = len(service.replica_groups) + for validation in validations: + if len(validation.replicas) != expected_group_count: + raise ValueError("preset validation replicas must match service replica group order") + for group in validation.replicas: + if not group.resources: + raise ValueError("preset validation replicas must specify resources") + for resources in group.resources: + _validate_replica_resources_are_exact(resources) + + +def _merge_presets( + existing: EndpointPreset, + incoming: EndpointPreset, + prefer_incoming: bool = False, +) -> EndpointPreset: + if existing.base.lower() != incoming.base.lower(): + raise ValueError("cannot merge endpoint presets for different bases") + recipes = [EndpointPresetRecipe.parse_obj(recipe.dict()) for recipe in existing.recipes] + recipes_by_id = {recipe.id: recipe for recipe in recipes} + incoming_recipe_ids = [] + for incoming_recipe in incoming.recipes: + existing_recipe = recipes_by_id.get(incoming_recipe.id) + if existing_recipe is None: + recipe = EndpointPresetRecipe.parse_obj(incoming_recipe.dict()) + recipes.append(recipe) + recipes_by_id[recipe.id] = recipe + incoming_recipe_ids.append(recipe.id) + continue + if ( + _service_configuration_to_preset_data(existing_recipe.service) + != (_service_configuration_to_preset_data(incoming_recipe.service)) + or existing_recipe.model != incoming_recipe.model + ): + raise ValueError(f"endpoint preset recipe id conflict: {incoming_recipe.id}") + _merge_recipe_validations(existing_recipe, incoming_recipe.validations) + incoming_recipe_ids.append(existing_recipe.id) + if prefer_incoming: + preferred_ids = set(incoming_recipe_ids) + recipes = [recipes_by_id[recipe_id] for recipe_id in incoming_recipe_ids] + [ + recipe for recipe in recipes if recipe.id not in preferred_ids + ] + return EndpointPreset(base=existing.base, recipes=recipes) + + +def _merge_recipe_validations( + recipe: EndpointPresetRecipe, + incoming_validations: list[EndpointPresetValidation], +) -> None: + existing_keys = {_validation_key(validation) for validation in recipe.validations} + for validation in incoming_validations: + key = _validation_key(validation) + if key in existing_keys: + continue + recipe.validations.append(EndpointPresetValidation.parse_obj(validation.dict())) + existing_keys.add(key) + + +def _validation_key(validation: EndpointPresetValidation) -> str: + data = json.loads(validation.json(exclude_none=True)) + return json.dumps(data, sort_keys=True, separators=(",", ":")) + + +def _preset_to_data(preset: EndpointPreset) -> dict[str, Any]: + return { + "type": "endpoint-preset", + "base": preset.base, + "recipes": [_recipe_to_data(recipe) for recipe in preset.recipes], + } + + +def _recipe_to_data(recipe: EndpointPresetRecipe) -> dict[str, Any]: + return { + "id": recipe.id, + "model": recipe.model, + "service": _service_configuration_to_preset_data(recipe.service), + "validations": [ + json.loads(validation.json(exclude_none=True)) for validation in recipe.validations + ], + } + + +def _service_configuration_to_preset_data( + configuration: ServiceConfiguration, +) -> dict[str, Any]: + service_data = json.loads(configuration.json(exclude_none=True)) + service_data.pop("type", None) + service_data.pop("name", None) + for field in ProfileParams.__fields__: + service_data.pop(field, None) + if configuration.env: + service_data["env"] = [ + _env_item_to_preset_data(key, value) + for key, value in sorted(configuration.env.items()) + ] + else: + service_data.pop("env", None) + for field, value in list(service_data.items()): + if value in ({}, []): + service_data.pop(field) + return service_data + + +def _env_item_to_preset_data(key: str, value: str | EnvSentinel) -> str: + if isinstance(value, EnvSentinel) or _is_secret_like_env(key=key, value=value): + return key + return f"{key}={value}" + + +def _is_secret_like_env(key: str, value: str | EnvSentinel) -> bool: + if _SECRET_ENV_PATTERN.search(key): + return True + return isinstance(value, str) and _SECRET_ENV_PATTERN.search(value) is not None + + +def _format_preset_comments(comments: Sequence[str]) -> str: + if not comments: + return "" + lines = [] + for comment in comments: + for line in comment.splitlines() or [""]: + lines.append(f"# {line}".rstrip()) + return "\n".join(lines) + "\n" + + +def _slugify_model(model: str) -> str: + slug = re.sub(r"[^a-z0-9]+", "-", model.lower()).strip("-") + return slug or "endpoint-preset" + + +def _get_legacy_replica_spec_groups( + *, + data: dict[str, Any], + service_data: dict[str, Any], +) -> list[dict[str, Any]]: + raw_groups = data.get("replica_spec_groups") + if not isinstance(raw_groups, list) or not raw_groups: + raise ValueError("legacy preset must specify non-empty replica_spec_groups") + if not all(isinstance(group, dict) for group in raw_groups): + raise ValueError("legacy preset replica_spec_groups must be objects") + groups = [_normalize_legacy_replica_spec_group(dict(group)) for group in raw_groups] + expected_names = _get_expected_replica_group_names(service_data) + if expected_names == [DEFAULT_REPLICA_GROUP_NAME] and "name" not in groups[0]: + groups[0]["name"] = DEFAULT_REPLICA_GROUP_NAME + group_names = [group.get("name") for group in groups] + if group_names != expected_names: + raise ValueError( + "legacy preset replica_spec_groups must match replica group order: " + + ", ".join(expected_names) + ) + return groups + + +def _normalize_legacy_replica_spec_group(group: dict[str, Any]) -> dict[str, Any]: + if "resources" in group or "tested_resources" in group: + if "resources" not in group or "tested_resources" not in group: + raise ValueError( + "legacy preset replica_spec_groups must specify resources and tested_resources" + ) + return group + replica_specs = group.get("replica_specs") + if not isinstance(replica_specs, list) or not replica_specs: + raise ValueError("legacy preset replica_spec_groups must specify resources") + group["resources"] = replica_specs[0] + group["tested_resources"] = replica_specs + group.pop("replica_specs", None) + return group + + +def _get_expected_replica_group_names(service_data: dict[str, Any]) -> list[str]: + replicas = service_data.get("replicas") + if not isinstance(replicas, list): + return [DEFAULT_REPLICA_GROUP_NAME] + names = [] + for group in replicas: + if not isinstance(group, dict): + raise ValueError("preset service replica groups must be objects") + name = group.get("name") + if not isinstance(name, str) or name == "": + raise ValueError("preset service replica groups must specify names") + names.append(name) + return names + + +def _apply_legacy_replica_group_resources( + *, + service_data: dict[str, Any], + groups: list[dict[str, Any]], +) -> None: + replicas = service_data.get("replicas") + if not isinstance(replicas, list): + service_data["resources"] = groups[0]["resources"] + return + resources_by_group = {group["name"]: group["resources"] for group in groups} + for group in replicas: + group["resources"] = resources_by_group[group["name"]] + + +def _get_legacy_group_tested_resources(group: dict[str, Any]) -> list[Any]: + tested_resources = group.get("tested_resources") + if not isinstance(tested_resources, list) or not tested_resources: + raise ValueError("legacy preset group must specify non-empty tested_resources") + return tested_resources + + +def _validate_replica_resources_are_exact(resources: ResourcesSpec) -> None: + cpu = parse_obj_as(CPUSpec, resources.cpu) + if not _is_exact_range(cpu.count): + _raise_loose_replica_resources() + if not _is_exact_range(resources.memory): + _raise_loose_replica_resources() + if resources.disk is None or not _is_exact_range(resources.disk.size): + _raise_loose_replica_resources() + gpu = resources.gpu + if gpu is None or not _is_exact_range(gpu.count): + _raise_loose_replica_resources() + gpu_count = gpu.count.min + if gpu_count == 0: + return + if gpu.name is None or len(gpu.name) != 1: + _raise_loose_replica_resources() + if gpu.memory is None or not _is_exact_range(gpu.memory): + _raise_loose_replica_resources() + if gpu.compute_capability is not None: + _raise_loose_replica_resources() + + +def _raise_loose_replica_resources() -> NoReturn: + raise ValueError("preset validations must use exact replica resources") + + +def _is_exact_range(value) -> bool: + return ( + value is not None + and value.min is not None + and value.max is not None + and value.min == value.max + ) diff --git a/src/dstack/_internal/server/services/events.py b/src/dstack/_internal/server/services/events.py index dd7b33dc7f..960fa4379f 100644 --- a/src/dstack/_internal/server/services/events.py +++ b/src/dstack/_internal/server/services/events.py @@ -11,6 +11,7 @@ from dstack._internal.core.models.users import GlobalRole from dstack._internal.server import settings from dstack._internal.server.models import ( + EndpointModel, EventModel, EventTargetModel, FleetModel, @@ -97,6 +98,7 @@ def from_model( SecretModel, UserModel, VolumeModel, + EndpointModel, ], ) -> "Target": if isinstance(model, FleetModel): @@ -162,6 +164,13 @@ def from_model( id=model.id, name=model.name, ) + if isinstance(model, EndpointModel): + return Target( + type=EventTargetType.ENDPOINT, + project_id=model.project_id or model.project.id, + id=model.id, + name=model.name, + ) raise ValueError(f"Unsupported model type: {type(model)}") def fmt(self) -> str: diff --git a/src/dstack/_internal/server/settings.py b/src/dstack/_internal/server/settings.py index 27a97a6db5..e0892fca0c 100644 --- a/src/dstack/_internal/server/settings.py +++ b/src/dstack/_internal/server/settings.py @@ -12,6 +12,14 @@ logger = get_logger(__name__) + +def _getenv_non_empty(name: str) -> str | None: + value = os.getenv(name) + if value is None or value == "": + return None + return value + + DSTACK_DIR_PATH = Path("~/.dstack/").expanduser() SERVER_DIR_PATH = Path(os.getenv("DSTACK_SERVER_DIR", DSTACK_DIR_PATH / "server")).resolve() @@ -21,6 +29,21 @@ SERVER_DATA_DIR_PATH = SERVER_DIR_PATH / "data" SERVER_DATA_DIR_PATH.mkdir(parents=True, exist_ok=True) +SERVER_PROJECTS_DIR_PATH = SERVER_DIR_PATH / "projects" + +AGENT_ANTHROPIC_API_KEY = _getenv_non_empty("DSTACK_AGENT_ANTHROPIC_API_KEY") +AGENT_CLAUDE_PATH = _getenv_non_empty("DSTACK_AGENT_CLAUDE_PATH") +# Development-only opt-in. Do not set together with +# DSTACK_AGENT_ANTHROPIC_API_KEY. Production servers should set +# DSTACK_AGENT_ANTHROPIC_API_KEY instead. Existing Claude CLI auth requires +# running without `claude --bare`, so the agent may read the server user's +# Claude CLI auth/settings. +AGENT_CLAUDE_USE_EXISTING_AUTH = environ.get_bool( + "DSTACK_AGENT_CLAUDE_USE_EXISTING_AUTH", default=False +) +AGENT_CLAUDE_EFFORT = _getenv_non_empty("DSTACK_AGENT_CLAUDE_EFFORT") +AGENT_ANTHROPIC_MODEL = os.getenv("DSTACK_AGENT_ANTHROPIC_MODEL", "claude-opus-4-8") + DATABASE_URL = os.getenv( "DSTACK_DATABASE_URL", f"sqlite+aiosqlite:///{str(SERVER_DATA_DIR_PATH.absolute())}/sqlite.db" ) diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py index 3b24ef96e5..142faac84c 100644 --- a/src/dstack/_internal/server/testing/common.py +++ b/src/dstack/_internal/server/testing/common.py @@ -438,12 +438,18 @@ async def create_job( disconnected_at: Optional[datetime] = None, registered: bool = False, waiting_master_job: Optional[bool] = None, + replica_group_name: Optional[str] = None, ) -> JobModel: if deployment_num is None: deployment_num = run.deployment_num run_spec = RunSpec.parse_raw(run.run_spec) job_spec = ( - await get_job_specs_from_run_spec(run_spec=run_spec, secrets={}, replica_num=replica_num) + await get_job_specs_from_run_spec( + run_spec=run_spec, + secrets={}, + replica_num=replica_num, + replica_group_name=replica_group_name, + ) )[0] job_spec.job_num = job_num job = JobModel( diff --git a/src/dstack/api/server/__init__.py b/src/dstack/api/server/__init__.py index 82b009863a..80e779a973 100644 --- a/src/dstack/api/server/__init__.py +++ b/src/dstack/api/server/__init__.py @@ -16,6 +16,8 @@ from dstack._internal.utils.logging import get_logger from dstack.api.server._auth import AuthAPIClient from dstack.api.server._backends import BackendsAPIClient +from dstack.api.server._endpoint_presets import EndpointPresetsAPIClient +from dstack.api.server._endpoints import EndpointsAPIClient from dstack.api.server._events import EventsAPIClient from dstack.api.server._exports import ExportsAPIClient from dstack.api.server._files import FilesAPIClient @@ -52,6 +54,7 @@ class APIClient: logs: operations with logs gateways: operations with gateways volumes: operations with volumes + endpoints: operations with endpoints exports: operations with exports files: operations with files """ @@ -129,6 +132,14 @@ def gateways(self) -> GatewaysAPIClient: def volumes(self) -> VolumesAPIClient: return VolumesAPIClient(self._request, self._logger) + @property + def endpoints(self) -> EndpointsAPIClient: + return EndpointsAPIClient(self._request, self._logger) + + @property + def endpoint_presets(self) -> EndpointPresetsAPIClient: + return EndpointPresetsAPIClient(self._request, self._logger) + @property def exports(self) -> ExportsAPIClient: return ExportsAPIClient(self._request, self._logger) diff --git a/src/dstack/api/server/_endpoint_presets.py b/src/dstack/api/server/_endpoint_presets.py new file mode 100644 index 0000000000..9542cdce64 --- /dev/null +++ b/src/dstack/api/server/_endpoint_presets.py @@ -0,0 +1,31 @@ +from typing import List + +from pydantic import parse_obj_as + +from dstack._internal.core.models.endpoint_presets import EndpointPreset +from dstack._internal.server.schemas.endpoint_presets import ( + DeleteEndpointPresetsRequest, + GetEndpointPresetRequest, +) +from dstack.api.server._group import APIClientGroup + + +class EndpointPresetsAPIClient(APIClientGroup): + def list(self, project_name: str) -> List[EndpointPreset]: + resp = self._request(f"/api/project/{project_name}/endpoints/presets/list") + return parse_obj_as(List[EndpointPreset.__response__], resp.json()) + + def get(self, project_name: str, model: str) -> EndpointPreset: + body = GetEndpointPresetRequest(model=model) + resp = self._request( + f"/api/project/{project_name}/endpoints/presets/get", + body=body.json(), + ) + return parse_obj_as(EndpointPreset.__response__, resp.json()) + + def delete(self, project_name: str, models: List[str]) -> None: + body = DeleteEndpointPresetsRequest(models=models) + self._request( + f"/api/project/{project_name}/endpoints/presets/delete", + body=body.json(), + ) diff --git a/src/dstack/api/server/_endpoints.py b/src/dstack/api/server/_endpoints.py new file mode 100644 index 0000000000..aaac8f439d --- /dev/null +++ b/src/dstack/api/server/_endpoints.py @@ -0,0 +1,49 @@ +from typing import List + +from pydantic import parse_obj_as + +from dstack._internal.core.models.endpoints import Endpoint, EndpointConfiguration, EndpointPlan +from dstack._internal.server.schemas.endpoints import ( + CreateEndpointRequest, + GetEndpointPlanRequest, + GetEndpointRequest, + StopEndpointsRequest, +) +from dstack.api.server._group import APIClientGroup + + +class EndpointsAPIClient(APIClientGroup): + def list(self, project_name: str) -> List[Endpoint]: + resp = self._request(f"/api/project/{project_name}/endpoints/list") + return parse_obj_as(List[Endpoint.__response__], resp.json()) + + def get(self, project_name: str, name: str) -> Endpoint: + body = GetEndpointRequest(name=name) + resp = self._request(f"/api/project/{project_name}/endpoints/get", body=body.json()) + return parse_obj_as(Endpoint.__response__, resp.json()) + + def get_plan( + self, + project_name: str, + configuration: EndpointConfiguration, + configuration_path: str, + ) -> EndpointPlan: + body = GetEndpointPlanRequest( + configuration=configuration, + configuration_path=configuration_path, + ) + resp = self._request(f"/api/project/{project_name}/endpoints/get_plan", body=body.json()) + return parse_obj_as(EndpointPlan.__response__, resp.json()) + + def create( + self, + project_name: str, + configuration: EndpointConfiguration, + ) -> Endpoint: + body = CreateEndpointRequest(configuration=configuration) + resp = self._request(f"/api/project/{project_name}/endpoints/create", body=body.json()) + return parse_obj_as(Endpoint.__response__, resp.json()) + + def stop(self, project_name: str, names: List[str]) -> None: + body = StopEndpointsRequest(names=names) + self._request(f"/api/project/{project_name}/endpoints/stop", body=body.json()) diff --git a/src/tests/_internal/cli/commands/test_endpoint.py b/src/tests/_internal/cli/commands/test_endpoint.py new file mode 100644 index 0000000000..74cea4dbad --- /dev/null +++ b/src/tests/_internal/cli/commands/test_endpoint.py @@ -0,0 +1,158 @@ +import argparse +import base64 +from datetime import datetime, timezone +from types import SimpleNamespace +from uuid import uuid4 + +from dstack._internal.cli.commands.endpoint import EndpointCommand +from dstack._internal.core.models.endpoints import Endpoint, EndpointConfiguration, EndpointStatus +from dstack._internal.core.models.logs import JobSubmissionLogs, LogEvent, LogEventSource + + +def _get_endpoint(status: EndpointStatus = EndpointStatus.RUNNING) -> Endpoint: + return Endpoint( + id=uuid4(), + name="qwen-endpoint", + project_name="main", + user="test-user", + configuration=EndpointConfiguration(name="qwen-endpoint", model="Qwen/Qwen3-0.6B"), + created_at=datetime.now(timezone.utc), + last_processed_at=datetime.now(timezone.utc), + status=status, + run_name="qwen-endpoint-1", + ) + + +class _FakeEndpoints: + def __init__(self, endpoint: Endpoint): + self._endpoint = endpoint + self.get_requests = [] + self.stopped_names = [] + + def get(self, project_name, name): + self.get_requests.append( + { + "project_name": project_name, + "name": name, + } + ) + return self._endpoint + + def stop(self, project_name, names): + self.stopped_names.extend(names) + + +class _FakeLogs: + def __init__(self, responses=None): + self._responses = responses + self.requests = [] + + def poll(self, project_name, body): + self.requests.append(body) + if self._responses is not None: + return self._responses.pop(0) + return JobSubmissionLogs( + logs=[ + LogEvent( + timestamp=datetime.now(timezone.utc), + log_source=LogEventSource.STDOUT, + message=base64.b64encode(b"agent log\n").decode(), + ) + ], + ) + + +def _get_command(endpoint: Endpoint, logs=None) -> EndpointCommand: + command = EndpointCommand.__new__(EndpointCommand) + command.api = SimpleNamespace( + project="main", + client=SimpleNamespace( + endpoints=_FakeEndpoints(endpoint), + logs=logs or _FakeLogs(), + ), + ) + return command + + +class TestEndpointCommand: + def test_preset_list_verbose_can_be_before_or_after_list_action(self): + parser = argparse.ArgumentParser() + EndpointCommand(parser)._register() + + assert parser.parse_args(["preset", "-v", "list"]).verbose is True + assert parser.parse_args(["preset", "list", "-v"]).verbose is True + assert parser.parse_args(["preset", "list"]).verbose is False + + def test_reads_endpoint_logs(self): + endpoint = _get_endpoint() + command = _get_command(endpoint) + + logs = list(command._get_endpoint_logs(endpoint=endpoint, start_time=None)) + + assert len(logs) == 1 + assert logs[0].message == "agent log\n" + request = command.api.client.logs.requests[0] + assert request.run_name == endpoint.name + assert request.job_submission_id == endpoint.id + + def test_watch_endpoint_logs_does_not_swallow_same_timestamp_logs(self, monkeypatch): + endpoint = _get_endpoint() + timestamp = datetime.now(timezone.utc) + logs_api = _FakeLogs( + responses=[ + JobSubmissionLogs( + logs=[ + LogEvent( + timestamp=timestamp, + log_source=LogEventSource.STDOUT, + message=base64.b64encode(b"first\n").decode(), + ), + ], + ), + JobSubmissionLogs( + logs=[ + LogEvent( + timestamp=timestamp, + log_source=LogEventSource.STDOUT, + message=base64.b64encode(b"first\n").decode(), + ), + LogEvent( + timestamp=timestamp, + log_source=LogEventSource.STDOUT, + message=base64.b64encode(b"second\n").decode(), + ), + ], + ), + ], + ) + command = _get_command(endpoint, logs=logs_api) + monkeypatch.setattr("dstack._internal.cli.commands.endpoint.time.sleep", lambda _: None) + + logs = command._get_endpoint_logs(endpoint=endpoint, start_time=None, watch=True) + + assert next(logs).message == "first\n" + assert next(logs).message == "second\n" + + def test_stop_endpoint(self): + endpoint = _get_endpoint() + command = _get_command(endpoint) + + command._stop(argparse.Namespace(name=endpoint.name, yes=True)) + + assert command.api.client.endpoints.stopped_names == [endpoint.name] + + def test_stop_already_stopped_endpoint_is_noop(self): + endpoint = _get_endpoint(status=EndpointStatus.STOPPED) + command = _get_command(endpoint) + + command._stop(argparse.Namespace(name=endpoint.name, yes=True)) + + assert command.api.client.endpoints.stopped_names == [] + + def test_stop_failed_endpoint_is_noop(self): + endpoint = _get_endpoint(status=EndpointStatus.FAILED) + command = _get_command(endpoint) + + command._stop(argparse.Namespace(name=endpoint.name, yes=True)) + + assert command.api.client.endpoints.stopped_names == [] diff --git a/src/tests/_internal/cli/commands/test_logs.py b/src/tests/_internal/cli/commands/test_logs.py new file mode 100644 index 0000000000..4600c91311 --- /dev/null +++ b/src/tests/_internal/cli/commands/test_logs.py @@ -0,0 +1,55 @@ +import argparse +from types import SimpleNamespace + +import pytest + +from dstack._internal.cli.commands.logs import LogsCommand +from dstack._internal.core.errors import CLIError + + +class _FakeRun: + def logs(self, **kwargs): + yield b"run log\n" + + +class _FakeRuns: + def __init__(self, run=None): + self._run = run + self.requested_names = [] + + def get(self, name): + self.requested_names.append(name) + return self._run + + +def _get_command(run=None) -> LogsCommand: + command = LogsCommand.__new__(LogsCommand) + command.api = SimpleNamespace( + runs=_FakeRuns(run=run), + ) + return command + + +def _get_args(name="qwen-run"): + return argparse.Namespace( + run_name=name, + diagnose=False, + replica=0, + job=0, + ) + + +class TestLogsCommand: + def test_reads_run_logs(self): + command = _get_command(run=_FakeRun()) + + logs = list(command._get_logs(args=_get_args(), start_time=None)) + + assert logs == [b"run log\n"] + assert command.api.runs.requested_names == ["qwen-run"] + + def test_errors_when_run_is_not_found(self): + command = _get_command(run=None) + + with pytest.raises(CLIError, match="Run qwen-run not found"): + list(command._get_logs(args=_get_args(), start_time=None)) diff --git a/src/tests/_internal/cli/services/configurators/test_endpoint.py b/src/tests/_internal/cli/services/configurators/test_endpoint.py new file mode 100644 index 0000000000..edb057741c --- /dev/null +++ b/src/tests/_internal/cli/services/configurators/test_endpoint.py @@ -0,0 +1,242 @@ +from datetime import datetime, timezone +from uuid import uuid4 + +import pytest +from rich.console import Console + +import dstack._internal.cli.services.configurators.endpoint as endpoint_configurator_module +from dstack._internal.cli.services.configurators.endpoint import ( + EndpointConfigurator, + _get_apply_confirm_message, + _make_endpoint_url_absolute, + _print_endpoint_plan, + _print_finished_endpoint_message, + _print_submitted_endpoint_message, +) +from dstack._internal.core.errors import ConfigurationError +from dstack._internal.core.models.common import ApplyAction +from dstack._internal.core.models.endpoints import ( + Endpoint, + EndpointConfiguration, + EndpointPlan, + EndpointPlanJobOffers, + EndpointPresetPolicy, + EndpointProvisioningPlanAgent, + EndpointProvisioningPlanNone, + EndpointProvisioningPlanPreset, + EndpointStatus, +) +from dstack._internal.core.models.resources import ResourcesSpec + + +def _get_endpoint_plan( + provisioning_plan: ( + EndpointProvisioningPlanNone + | EndpointProvisioningPlanPreset + | EndpointProvisioningPlanAgent + ), + current_resource: Endpoint | None = None, + preset_policy: EndpointPresetPolicy = EndpointPresetPolicy.REUSE_OR_CREATE, +) -> EndpointPlan: + configuration = EndpointConfiguration(name="qwen-endpoint", model="Qwen/Qwen3-0.6B") + return EndpointPlan( + project_name="main", + user="test-user", + configuration=configuration, + configuration_path="endpoint.dstack.yml", + current_resource=current_resource, + action=ApplyAction.CREATE if current_resource is None else ApplyAction.UPDATE, + preset_policy=preset_policy, + provisioning_plan=provisioning_plan, + ) + + +def _get_no_provisioning_plan() -> EndpointProvisioningPlanNone: + return EndpointProvisioningPlanNone( + reason=( + "No matching endpoint presets found. " + "Creating a preset requires the server agent, but " + "DSTACK_AGENT_ANTHROPIC_API_KEY is not set." + ) + ) + + +def _get_current_endpoint() -> Endpoint: + configuration = EndpointConfiguration(name="qwen-endpoint", model="Qwen/Qwen3-0.6B") + return Endpoint( + id=uuid4(), + name="qwen-endpoint", + project_name="main", + user="test-user", + configuration=configuration, + created_at=datetime.now(timezone.utc), + last_processed_at=datetime.now(timezone.utc), + status=EndpointStatus.RUNNING, + ) + + +def _get_failed_endpoint() -> Endpoint: + endpoint = _get_current_endpoint() + endpoint.status = EndpointStatus.FAILED + endpoint.status_message = ( + "Preset policy create requires the server agent, but " + "DSTACK_AGENT_ANTHROPIC_API_KEY is not set." + ) + return endpoint + + +def _patch_console(monkeypatch: pytest.MonkeyPatch) -> Console: + console = Console(record=True, force_terminal=False, color_system=None, width=120) + monkeypatch.setattr(endpoint_configurator_module, "console", console) + return console + + +class TestPrintEndpointPlan: + def test_prints_stable_properties_for_missing_provisioning_path( + self, monkeypatch: pytest.MonkeyPatch + ): + console = _patch_console(monkeypatch) + plan = _get_endpoint_plan(_get_no_provisioning_plan()) + + _print_endpoint_plan(plan) + + output = console.export_text() + assert "Project main" in output + assert "Endpoint" not in output + assert "Resources" not in output + assert "Spot policy auto" in output + assert "Max price off" in output + assert "Preset policy reuse-or-create" in output + assert "No matching endpoint presets found" in output + assert "Creating a preset requires the server agent" in output + assert "Model" not in output + assert "Action" not in output + assert "Provisioning" not in output + + def test_prints_preset_without_resources_property(self, monkeypatch: pytest.MonkeyPatch): + console = _patch_console(monkeypatch) + resources = ResourcesSpec.parse_obj({"gpu": "16GB", "disk": "60GB"}) + plan = _get_endpoint_plan( + EndpointProvisioningPlanPreset( + preset_base="Qwen/Qwen3-0.6B", + recipe_id="vllm-a40", + service_name="qwen-endpoint-serving", + job_offers=[ + EndpointPlanJobOffers( + replica_group="0", + resources=resources, + spot=False, + max_price=0.5, + offers=[], + total_offers=0, + max_offer_price=None, + ) + ], + ) + ) + + _print_endpoint_plan(plan) + + output = console.export_text() + assert "Resources" not in output + assert "Preset Qwen/Qwen3-0.6B" in output + assert "Recipe vllm-a40" in output + + def test_prints_agent_reason_when_preset_has_no_offers(self, monkeypatch: pytest.MonkeyPatch): + console = _patch_console(monkeypatch) + plan = _get_endpoint_plan( + EndpointProvisioningPlanAgent( + agent_model="test-agent", + reason="Endpoint preset qwen matched but has no available offers.", + ) + ) + + _print_endpoint_plan(plan) + + output = console.export_text() + assert "Preset policy reuse-or-create" in output + assert "Endpoint preset qwen matched but has no available offers." in output + assert "Agent" not in output + + +class TestPrintSubmittedEndpointMessage: + def test_prints_generic_detach_message(self, monkeypatch: pytest.MonkeyPatch): + console = _patch_console(monkeypatch) + + _print_submitted_endpoint_message(_get_current_endpoint()) + + output = console.export_text() + assert "Endpoint qwen-endpoint submitted, detaching..." in output + assert "without a provisioning path" not in output + + +class TestPrintFinishedEndpointMessage: + def test_prints_running_url(self, monkeypatch: pytest.MonkeyPatch): + console = _patch_console(monkeypatch) + endpoint = _get_current_endpoint() + endpoint.url = "/proxy/services/main/qwen-endpoint/v1" + + _print_finished_endpoint_message(endpoint) + + output = console.export_text() + assert "/proxy/services/main/qwen-endpoint/v1" in output + + def test_makes_relative_url_absolute(self): + endpoint = _get_current_endpoint() + endpoint.url = "/proxy/services/main/qwen-endpoint/v1" + + _make_endpoint_url_absolute(endpoint, "http://127.0.0.1:8000") + + assert endpoint.url == "http://127.0.0.1:8000/proxy/services/main/qwen-endpoint/v1" + + def test_prints_failed_message_without_logs_hint(self, monkeypatch: pytest.MonkeyPatch): + console = _patch_console(monkeypatch) + + _print_finished_endpoint_message(_get_failed_endpoint()) + + output = console.export_text() + assert "DSTACK_AGENT_ANTHROPIC_API_KEY" in output + assert "dstack logs" not in output + + +class TestGetApplyConfirmMessage: + def test_asks_to_create_without_provisioning_path(self): + plan = _get_endpoint_plan(_get_no_provisioning_plan()) + + assert _get_apply_confirm_message(plan) == "Create the endpoint?" + + def test_asks_to_override_without_provisioning_path(self): + plan = _get_endpoint_plan( + _get_no_provisioning_plan(), + current_resource=_get_current_endpoint(), + ) + + assert ( + _get_apply_confirm_message(plan) + == "Stop and override the endpoint [code]qwen-endpoint[/]?" + ) + + def test_asks_to_create_when_existing_endpoint_is_terminal(self): + plan = _get_endpoint_plan( + _get_no_provisioning_plan(), + current_resource=_get_failed_endpoint(), + ) + + assert _get_apply_confirm_message(plan) == "Create the endpoint?" + + def test_asks_to_override_same_name_run(self): + plan = _get_endpoint_plan(EndpointProvisioningPlanAgent(agent_model="test-agent")) + + assert ( + _get_apply_confirm_message(plan, stop_run_name="qwen-endpoint") + == "Stop and override the run [code]qwen-endpoint[/]?" + ) + + +class TestEndpointConfigurator: + def test_delete_configuration_is_not_supported(self): + configurator = EndpointConfigurator.__new__(EndpointConfigurator) + conf = EndpointConfiguration(name="qwen-endpoint", model="Qwen/Qwen3-0.6B") + + with pytest.raises(ConfigurationError, match="dstack endpoint stop"): + configurator.delete_configuration(conf, "endpoint.dstack.yml", command_args=None) diff --git a/src/tests/_internal/cli/utils/test_endpoint.py b/src/tests/_internal/cli/utils/test_endpoint.py new file mode 100644 index 0000000000..850854d64e --- /dev/null +++ b/src/tests/_internal/cli/utils/test_endpoint.py @@ -0,0 +1,286 @@ +from datetime import datetime, timezone +from uuid import uuid4 + +from dstack._internal.cli.utils.endpoint import ( + filter_endpoints_for_listing, + get_endpoint_table, + get_endpoints_table, +) +from dstack._internal.core.models.endpoints import ( + Endpoint, + EndpointConfiguration, + EndpointStatus, +) + + +def _get_endpoint( + name: str = "qwen-endpoint", + status: EndpointStatus = EndpointStatus.FAILED, + status_message: str | None = "No matching endpoint presets found.", + created_at: datetime | None = None, + model_repo: str | None = None, +) -> Endpoint: + if created_at is None: + created_at = datetime.now(timezone.utc) + return Endpoint( + id=uuid4(), + name=name, + project_name="main", + user="test-user", + configuration=EndpointConfiguration(name=name, model="Qwen/Qwen3-0.6B"), + created_at=created_at, + last_processed_at=created_at, + status=status, + status_message=status_message, + model_repo=model_repo, + ) + + +class TestGetEndpointsTable: + def test_default_table_does_not_show_status_message(self): + table = get_endpoints_table([_get_endpoint()]) + + assert "ERROR" not in [column.header for column in table.columns] + + def test_default_table_shows_preset_policy_not_service_run(self): + table = get_endpoints_table([_get_endpoint()]) + + headers = [column.header for column in table.columns] + assert "POLICY" in headers + assert "SERVICE RUN" not in headers + + def test_verbose_table_shows_status_message(self): + table = get_endpoints_table([_get_endpoint()], verbose=True) + + assert "ERROR" in [column.header for column in table.columns] + assert "SERVICE RUN" in [column.header for column in table.columns] + + def test_failed_status_without_reason_is_colored(self): + table = get_endpoints_table([_get_endpoint(status_message=None)]) + + status_column = next(column for column in table.columns if column.header == "STATUS") + assert status_column._cells == ["[indian_red1]failed[/]"] + + def test_no_preset_status_is_colored(self): + table = get_endpoints_table([_get_endpoint()]) + + status_column = next(column for column in table.columns if column.header == "STATUS") + assert status_column._cells == ["[indian_red1]no preset[/]"] + + def test_no_agent_status_is_colored(self): + table = get_endpoints_table( + [ + _get_endpoint( + status_message=( + "No matching endpoint presets found. Creating a preset requires " + "the server agent, but DSTACK_AGENT_ANTHROPIC_API_KEY is not set." + ) + ) + ] + ) + + status_column = next(column for column in table.columns if column.header == "STATUS") + assert status_column._cells == ["[indian_red1]no agent[/]"] + + def test_no_offers_status_is_colored(self): + table = get_endpoints_table( + [ + _get_endpoint( + status_message=( + "No dstack service could be deployed because max_price matches " + "ZERO offers." + ) + ) + ] + ) + + status_column = next(column for column in table.columns if column.header == "STATUS") + assert status_column._cells == ["[gold1]no offers[/]"] + + def test_agent_failed_status_is_colored(self): + table = get_endpoints_table( + [_get_endpoint(status_message="Server agent process exited without a final report")] + ) + + status_column = next(column for column in table.columns if column.header == "STATUS") + assert status_column._cells == ["[indian_red1]agent failed[/]"] + + def test_running_status_is_colored(self): + table = get_endpoints_table([_get_endpoint(status=EndpointStatus.RUNNING)]) + + status_column = next(column for column in table.columns if column.header == "STATUS") + assert status_column._cells == ["[bold sea_green3]running[/]"] + + def test_prototyping_status_is_colored(self): + table = get_endpoints_table([_get_endpoint(status=EndpointStatus.PROTOTYPING)]) + + status_column = next(column for column in table.columns if column.header == "STATUS") + assert status_column._cells == ["[bold medium_purple1]prototyping[/]"] + + def test_stopping_status_is_colored(self): + table = get_endpoints_table([_get_endpoint(status=EndpointStatus.STOPPING)]) + + status_column = next(column for column in table.columns if column.header == "STATUS") + assert status_column._cells == ["[bold deep_sky_blue1]stopping[/]"] + + def test_stopped_status_is_colored(self): + table = get_endpoints_table([_get_endpoint(status=EndpointStatus.STOPPED)]) + + status_column = next(column for column in table.columns if column.header == "STATUS") + assert status_column._cells == ["[grey62]stopped[/]"] + + def test_shows_repo_row_when_repo_differs_from_model(self): + table = get_endpoints_table([_get_endpoint(model_repo="groxaxo/Qwen3-0.6B-GPTQ-4Bit")]) + + model_column = next(column for column in table.columns if column.header == "MODEL") + assert model_column._cells == [ + "Qwen/Qwen3-0.6B", + " repo=groxaxo/Qwen3-0.6B-GPTQ-4Bit", + ] + + +class TestGetEndpointTable: + def test_shows_endpoint_details(self): + endpoint = _get_endpoint( + status_message="No matching endpoint presets found.", + model_repo="groxaxo/Qwen3-0.6B-GPTQ-4Bit", + ) + + table = get_endpoint_table(endpoint, format_date=lambda _: "now") + + assert table.columns[0]._cells == [ + "[bold]Project[/bold]", + "[bold]User[/bold]", + "[bold]Endpoint[/bold]", + "[bold]Model[/bold]", + "[bold]Repo[/bold]", + "[bold]Status[/bold]", + "[bold]Preset policy[/bold]", + "[bold]Service run[/bold]", + "[bold]URL[/bold]", + "[bold]Created[/bold]", + "[bold]Error[/bold]", + ] + assert table.columns[1]._cells == [ + "main", + "test-user", + "qwen-endpoint", + "Qwen/Qwen3-0.6B", + "groxaxo/Qwen3-0.6B-GPTQ-4Bit", + "[indian_red1]no preset[/]", + "reuse-or-create", + "-", + "-", + "now", + "No matching endpoint presets found.", + ] + + +class TestFilterEndpointsForListing: + def test_default_shows_unfinished_only_when_present(self): + endpoints = [ + _get_endpoint( + name="failed-old", + status=EndpointStatus.FAILED, + created_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + ), + _get_endpoint( + name="running", + status=EndpointStatus.RUNNING, + created_at=datetime(2026, 1, 2, tzinfo=timezone.utc), + ), + _get_endpoint( + name="failed-new", + status=EndpointStatus.FAILED, + created_at=datetime(2026, 1, 3, tzinfo=timezone.utc), + ), + _get_endpoint( + name="prototyping", + status=EndpointStatus.PROTOTYPING, + created_at=datetime(2026, 1, 4, tzinfo=timezone.utc), + ), + ] + + filtered = filter_endpoints_for_listing(endpoints) + + assert [endpoint.name for endpoint in filtered] == [ + "prototyping", + "running", + ] + + def test_default_shows_latest_finished_when_none_are_unfinished(self): + endpoints = [ + _get_endpoint( + name="failed-old", + status=EndpointStatus.FAILED, + created_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + ), + _get_endpoint( + name="failed-new", + status=EndpointStatus.FAILED, + created_at=datetime(2026, 1, 2, tzinfo=timezone.utc), + ), + ] + + filtered = filter_endpoints_for_listing(endpoints) + + assert [endpoint.name for endpoint in filtered] == ["failed-new"] + + def test_default_hides_latest_finished_when_unfinished_exist(self): + endpoints = [ + _get_endpoint( + name="failed-new", + status=EndpointStatus.FAILED, + created_at=datetime(2026, 1, 3, tzinfo=timezone.utc), + ), + _get_endpoint( + name="prototyping", + status=EndpointStatus.PROTOTYPING, + created_at=datetime(2026, 1, 4, tzinfo=timezone.utc), + ), + ] + + filtered = filter_endpoints_for_listing(endpoints) + + assert [endpoint.name for endpoint in filtered] == ["prototyping"] + + def test_all_shows_all_sorted_newest_first(self): + endpoints = [ + _get_endpoint( + name="older", + status=EndpointStatus.FAILED, + created_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + ), + _get_endpoint( + name="newer", + status=EndpointStatus.FAILED, + created_at=datetime(2026, 1, 2, tzinfo=timezone.utc), + ), + ] + + filtered = filter_endpoints_for_listing(endpoints, show_all=True) + + assert [endpoint.name for endpoint in filtered] == ["newer", "older"] + + def test_last_implies_all(self): + endpoints = [ + _get_endpoint( + name="oldest", + status=EndpointStatus.FAILED, + created_at=datetime(2026, 1, 1, tzinfo=timezone.utc), + ), + _get_endpoint( + name="middle", + status=EndpointStatus.FAILED, + created_at=datetime(2026, 1, 2, tzinfo=timezone.utc), + ), + _get_endpoint( + name="newest", + status=EndpointStatus.FAILED, + created_at=datetime(2026, 1, 3, tzinfo=timezone.utc), + ), + ] + + filtered = filter_endpoints_for_listing(endpoints, limit=2) + + assert [endpoint.name for endpoint in filtered] == ["newest", "middle"] diff --git a/src/tests/_internal/cli/utils/test_preset.py b/src/tests/_internal/cli/utils/test_preset.py new file mode 100644 index 0000000000..dde551f0e1 --- /dev/null +++ b/src/tests/_internal/cli/utils/test_preset.py @@ -0,0 +1,193 @@ +from dstack._internal.cli.utils.preset import get_endpoint_presets_table +from dstack._internal.core.models.configurations import ServiceConfiguration +from dstack._internal.core.models.endpoint_presets import ( + EndpointPreset, + EndpointPresetRecipe, + EndpointPresetValidation, + EndpointPresetValidationReplica, +) + + +class TestGetEndpointPresetsTable: + def test_shows_single_group_preset(self): + preset = _endpoint_preset( + model="Qwen/Qwen3-0.6B", + recipe_id="vllm-a40", + service={"resources": {"gpu": "nvidia:16GB"}}, + ) + + table = get_endpoint_presets_table([preset]) + + assert [column.header for column in table.columns] == ["MODEL", "GPU"] + assert table.columns[0]._cells == ["[bold]Qwen/Qwen3-0.6B[/]"] + assert table.columns[1]._cells == [ + "nvidia:16GB:1..", + ] + + def test_verbose_shows_full_resources(self): + preset = _endpoint_preset( + model="Qwen/Qwen3-0.6B", + recipe_id="vllm-a40", + service={"resources": {"gpu": "nvidia:16GB"}}, + ) + + table = get_endpoint_presets_table([preset], verbose=True) + + assert [column.header for column in table.columns] == ["MODEL", "RESOURCES"] + assert table.columns[0]._cells == ["[bold]Qwen/Qwen3-0.6B[/]"] + assert table.columns[1]._cells == [ + "cpu=2.. mem=8GB.. disk=100GB.. gpu=nvidia:16GB:1..", + ] + + def test_shows_recipe_model_when_it_differs_from_preset_base(self): + preset = _endpoint_preset( + model="Qwen/Qwen3-0.6B", + recipe_id="vllm-a40", + recipe_model="groxaxo/Qwen3-0.6B-GPTQ-4Bit", + service={"resources": {"gpu": "nvidia:16GB"}}, + ) + + table = get_endpoint_presets_table([preset]) + + assert table.columns[0]._cells == [ + "[bold]Qwen/Qwen3-0.6B[/]", + " repo=groxaxo/Qwen3-0.6B-GPTQ-4Bit", + ] + assert table.columns[1]._cells == [ + "nvidia:16GB:1..", + "", + ] + + def test_shows_single_implicit_group_once_even_with_multiple_tested_replicas(self): + preset = _endpoint_preset( + model="Qwen/Qwen3-0.6B", + recipe_id="vllm-a40", + service={"resources": {"gpu": "nvidia:16GB"}}, + validation_replicas=[ + { + "resources": [ + _exact_resources(), + _exact_resources(), + ] + } + ], + ) + + table = get_endpoint_presets_table([preset]) + + assert table.columns[0]._cells == ["[bold]Qwen/Qwen3-0.6B[/]"] + assert table.columns[1]._cells == [ + "nvidia:16GB:1..", + ] + + def test_shows_grouped_replica_specs(self): + preset = _endpoint_preset( + model="Qwen/Qwen3-30B-A3B", + recipe_id="pd-l4", + service={ + "replicas": [ + { + "name": "router", + "count": 1, + "commands": ["python router.py"], + "resources": {"cpu": 4}, + }, + { + "name": "worker", + "count": 2, + "commands": ["vllm serve Qwen/Qwen3-30B-A3B"], + "resources": {"gpu": "nvidia:24GB"}, + }, + ] + }, + ) + + table = get_endpoint_presets_table([preset]) + + assert table.columns[0]._cells == [ + "[bold]Qwen/Qwen3-30B-A3B[/]", + "[secondary] group=router[/]", + "[secondary] group=worker[/]", + ] + assert table.columns[1]._cells == [ + "", + "-", + "nvidia:24GB:1..", + ] + + def test_groups_multiple_recipes_by_model(self): + preset = EndpointPreset( + base="Qwen/Qwen3-0.6B", + recipes=[ + _endpoint_preset( + model="Qwen/Qwen3-0.6B", + recipe_id="vllm-a40", + service={"resources": {"gpu": "nvidia:16GB"}}, + ).recipes[0], + _endpoint_preset( + model="Qwen/Qwen3-0.6B", + recipe_id="vllm-l4", + service={"resources": {"gpu": "nvidia:24GB"}}, + ).recipes[0], + ], + ) + + table = get_endpoint_presets_table([preset]) + + assert table.columns[0]._cells == [ + "[bold]Qwen/Qwen3-0.6B[/]", + "[secondary] recipe=0[/]", + "[secondary] recipe=1[/]", + ] + assert table.columns[1]._cells == [ + "", + "nvidia:16GB:1..", + "nvidia:24GB:1..", + ] + + +def _endpoint_preset( + *, + model: str, + recipe_id: str, + service: dict, + validation_replicas: list[dict] | None = None, + recipe_model: str | None = None, +) -> EndpointPreset: + service_data = { + "type": "service", + "port": 8000, + "model": model, + **service, + } + if "replicas" not in service_data: + service_data["commands"] = [f"vllm serve {model}"] + return EndpointPreset( + base=model, + recipes=[ + EndpointPresetRecipe( + id=recipe_id, + model=recipe_model or model, + service=ServiceConfiguration.parse_obj(service_data), + validations=[ + EndpointPresetValidation( + replicas=[ + EndpointPresetValidationReplica.parse_obj(replica) + for replica in ( + validation_replicas or [{"resources": [_exact_resources()]}] + ) + ] + ) + ], + ) + ], + ) + + +def _exact_resources() -> dict: + return { + "cpu": 9, + "memory": "50GB", + "disk": "60GB", + "gpu": {"name": "A40", "memory": "48GB", "count": 1}, + } diff --git a/src/tests/_internal/core/models/test_endpoints.py b/src/tests/_internal/core/models/test_endpoints.py new file mode 100644 index 0000000000..19f033d68f --- /dev/null +++ b/src/tests/_internal/core/models/test_endpoints.py @@ -0,0 +1,119 @@ +import pytest + +from dstack._internal.core.errors import ConfigurationError +from dstack._internal.core.models.configurations import parse_apply_configuration +from dstack._internal.core.models.endpoints import ( + EndpointConfiguration, + EndpointModelBase, + EndpointModelRepo, + EndpointPresetPolicy, +) +from dstack._internal.core.models.profiles import CreationPolicy + + +class TestEndpointConfiguration: + def test_parses_endpoint_configuration(self): + conf = parse_apply_configuration( + { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + "env": {"HF_TOKEN": "secret"}, + "fleets": ["gpu-fleet"], + "creation_policy": "reuse", + "preset_policy": "reuse", + } + ) + + assert isinstance(conf, EndpointConfiguration) + assert conf.name == "qwen-endpoint" + assert isinstance(conf.model, EndpointModelRepo) + assert conf.model.repo == "Qwen/Qwen3-0.6B" + assert conf.env.as_dict() == {"HF_TOKEN": "secret"} + assert conf.fleets is not None + fleet = conf.fleets[0] + assert not isinstance(fleet, str) + assert fleet.name == "gpu-fleet" + assert conf.creation_policy == CreationPolicy.REUSE + assert conf.preset_policy == EndpointPresetPolicy.REUSE + + def test_defaults_to_reuse_or_create_preset_policy(self): + conf = EndpointConfiguration(name="qwen-endpoint", model="Qwen/Qwen3-0.6B") + + assert conf.preset_policy == EndpointPresetPolicy.REUSE_OR_CREATE + + def test_parses_exact_repo_model(self): + conf = parse_apply_configuration( + { + "type": "endpoint", + "name": "qwen-endpoint", + "model": {"repo": "groxaxo/Qwen3.6-27B-GPTQ-Pro-4Bit"}, + } + ) + + assert isinstance(conf.model, EndpointModelRepo) + assert conf.model.api_model_name == "groxaxo/Qwen3.6-27B-GPTQ-Pro-4Bit" + assert conf.model.exact_repo == "groxaxo/Qwen3.6-27B-GPTQ-Pro-4Bit" + assert not conf.model.allows_variant_selection + + def test_parses_exact_repo_with_api_model_name(self): + conf = parse_apply_configuration( + { + "type": "endpoint", + "name": "qwen-endpoint", + "model": { + "repo": "groxaxo/Qwen3.6-27B-GPTQ-Pro-4Bit", + "name": "Qwen/Qwen3.6-27B", + }, + } + ) + + assert isinstance(conf.model, EndpointModelRepo) + assert conf.model.api_model_name == "Qwen/Qwen3.6-27B" + assert conf.model.exact_repo == "groxaxo/Qwen3.6-27B-GPTQ-Pro-4Bit" + assert not conf.model.allows_variant_selection + + def test_parses_base_model(self): + conf = parse_apply_configuration( + { + "type": "endpoint", + "name": "qwen-endpoint", + "model": {"base": "Qwen/Qwen3.6-27B"}, + } + ) + + assert isinstance(conf.model, EndpointModelBase) + assert conf.model.api_model_name == "Qwen/Qwen3.6-27B" + assert conf.model.exact_repo is None + assert conf.model.allows_variant_selection + + def test_rejects_model_with_repo_and_base(self): + with pytest.raises(ConfigurationError): + parse_apply_configuration( + { + "type": "endpoint", + "name": "qwen-endpoint", + "model": {"repo": "Qwen/Qwen3.6-27B", "base": "Qwen/Qwen3.6-27B"}, + } + ) + + def test_rejects_empty_model_repo(self): + with pytest.raises(ConfigurationError): + parse_apply_configuration( + { + "type": "endpoint", + "name": "qwen-endpoint", + "model": {"repo": ""}, + } + ) + + def test_rejects_unknown_fields(self): + with pytest.raises(ConfigurationError): + parse_apply_configuration( + { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + "commands": ["echo not supported"], + } + ) diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_endpoints.py b/src/tests/_internal/server/background/pipeline_tasks/test_endpoints.py new file mode 100644 index 0000000000..23a0401f24 --- /dev/null +++ b/src/tests/_internal/server/background/pipeline_tasks/test_endpoints.py @@ -0,0 +1,2170 @@ +import uuid +from datetime import datetime, timezone +from unittest.mock import AsyncMock, Mock, patch + +import pytest +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.orm import joinedload + +from dstack._internal.core.errors import ServerClientError +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.configurations import ServiceConfiguration +from dstack._internal.core.models.endpoints import ( + EndpointConfiguration, + EndpointModelBase, + EndpointModelRepo, + EndpointModelSpec, + EndpointPresetPolicy, + EndpointStatus, +) +from dstack._internal.core.models.envs import Env +from dstack._internal.core.models.instances import ( + Disk, + Gpu, + InstanceAvailability, + InstanceOfferWithAvailability, + InstanceRuntime, + InstanceType, + Resources, +) +from dstack._internal.core.models.runs import JobStatus, RunSpec, RunStatus, ServiceSpec +from dstack._internal.core.models.users import GlobalRole, ProjectRole +from dstack._internal.server.background.pipeline_tasks.endpoints import ( + EndpointPipelineItem, + EndpointWorker, +) +from dstack._internal.server.db import get_session_ctx +from dstack._internal.server.models import EndpointModel, EndpointRunSubmissionModel, RunModel +from dstack._internal.server.services.endpoints import record_endpoint_run_submission +from dstack._internal.server.services.endpoints.agent import AgentPlan, AgentProvisioningResult +from dstack._internal.server.services.endpoints.agent.report import AgentFinalReport +from dstack._internal.server.services.endpoints.names import get_endpoint_serving_run_name +from dstack._internal.server.services.endpoints.planning import EndpointPresetPlanningResult +from dstack._internal.server.services.projects import add_project_member +from dstack._internal.server.testing.common import ( + create_fleet, + create_job, + create_project, + create_repo, + create_run, + create_user, + get_job_runtime_data, + list_events, +) + + +@pytest.fixture +def worker() -> EndpointWorker: + return EndpointWorker(queue=Mock(), heartbeater=Mock(), pipeline_hinter=Mock()) + + +class _FakeAgentService: + def __init__(self, result: AgentProvisioningResult | None = None, side_effect=None): + self.provision_endpoint = AsyncMock( + return_value=result or AgentProvisioningResult(), + side_effect=side_effect, + ) + self.abort_endpoint = AsyncMock(return_value=True) + + def is_enabled(self) -> bool: + return True + + def get_plan(self) -> AgentPlan: + return AgentPlan(model="test-agent") + + +def _endpoint_to_pipeline_item(endpoint_model: EndpointModel) -> EndpointPipelineItem: + assert endpoint_model.lock_token is not None + assert endpoint_model.lock_expires_at is not None + return EndpointPipelineItem( + __tablename__=endpoint_model.__tablename__, + id=endpoint_model.id, + lock_token=endpoint_model.lock_token, + lock_expires_at=endpoint_model.lock_expires_at, + prev_lock_expired=False, + status=endpoint_model.status, + ) + + +async def _lock_endpoint_model(session: AsyncSession, endpoint_model: EndpointModel) -> None: + endpoint_model.lock_token = uuid.uuid4() + endpoint_model.lock_expires_at = datetime(2025, 1, 2, 3, 4, tzinfo=timezone.utc) + await session.commit() + + +async def _create_endpoint_model( + session: AsyncSession, + status: EndpointStatus = EndpointStatus.SUBMITTED, + model: str | EndpointModelSpec = "Qwen/Qwen3-0.6B", + user_ssh_public_key: str | None = None, + user_global_role: GlobalRole = GlobalRole.ADMIN, + project_role: ProjectRole | None = None, +) -> EndpointModel: + project = await create_project(session=session) + user = await create_user( + session=session, + global_role=user_global_role, + ssh_public_key=user_ssh_public_key, + ) + if project_role is not None: + await add_project_member( + session=session, + project=project, + user=user, + project_role=project_role, + ) + configuration = EndpointConfiguration( + name="qwen-endpoint", + model=EndpointModelRepo(repo=model) if isinstance(model, str) else model, + env=Env.parse_obj({"HF_TOKEN": "secret"}), + ) + endpoint_model = EndpointModel( + project=project, + user=user, + name=configuration.name, + status=status, + configuration=configuration.json(), + created_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), + last_processed_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), + ) + endpoint_model.lock_token = uuid.uuid4() + endpoint_model.lock_expires_at = datetime(2025, 1, 2, 3, 4, tzinfo=timezone.utc) + session.add(endpoint_model) + await session.commit() + return endpoint_model + + +async def _create_backing_service_run( + session: AsyncSession, + endpoint_model: EndpointModel, + status: RunStatus = RunStatus.RUNNING, + job_status: JobStatus = JobStatus.RUNNING, + registered: bool = True, + deleted: bool = False, + service_model_base_url: str | None = "/proxy/services/main/qwen-endpoint-serving/v1", + link_endpoint: bool = True, + run_name: str | None = None, +): + if run_name is None: + run_name = get_endpoint_serving_run_name(endpoint_model.name) + assert run_name is not None + repo = await create_repo( + session=session, + project_id=endpoint_model.project_id, + repo_name=f"{run_name}-repo", + ) + run_spec = RunSpec( + run_name=run_name, + configuration=ServiceConfiguration.parse_obj( + { + "type": "service", + "name": run_name, + "commands": [ + "vllm serve Qwen/Qwen3-0.6B --host 0.0.0.0 --port 8000", + ], + "port": 8000, + "model": "Qwen/Qwen3-0.6B", + "resources": {"gpu": "16GB"}, + } + ), + ) + run = await create_run( + session=session, + project=endpoint_model.project, + repo=repo, + user=endpoint_model.user, + run_name=run_name, + status=status, + run_spec=run_spec, + deleted=deleted, + ) + if service_model_base_url is not None: + run.service_spec = ServiceSpec.parse_obj( + { + "url": "/proxy/services/main/qwen-endpoint-serving/", + "model": { + "name": "Qwen/Qwen3-0.6B", + "base_url": service_model_base_url, + "type": "chat", + }, + } + ).json() + if link_endpoint: + endpoint_model.service_run_id = run.id + await create_job( + session=session, + run=run, + status=job_status, + registered=registered, + job_runtime_data=get_job_runtime_data(offer=_instance_offer()), + ) + await session.commit() + return run + + +def _instance_offer() -> InstanceOfferWithAvailability: + return InstanceOfferWithAvailability( + backend=BackendType.AWS, + instance=InstanceType( + name="g5.xlarge", + resources=Resources( + cpus=4, + memory_mib=16 * 1024, + gpus=[Gpu(name="T4", memory_mib=16 * 1024)], + spot=False, + disk=Disk(size_mib=100 * 1024), + ), + ), + region="us-east-1", + price=1.0, + availability=InstanceAvailability.AVAILABLE, + instance_runtime=InstanceRuntime.SHIM, + ) + + +async def _create_ready_backing_service_run_for_agent( + endpoint_id: uuid.UUID, +) -> AgentProvisioningResult: + async with get_session_ctx() as session: + res = await session.execute( + select(EndpointModel) + .where(EndpointModel.id == endpoint_id) + .options(joinedload(EndpointModel.project)) + .options(joinedload(EndpointModel.user)) + ) + endpoint_model = res.unique().scalar_one() + run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + link_endpoint=False, + run_name=_get_agent_run_name(), + ) + return _get_verified_agent_result(run) + + +def _get_verified_agent_result( + run: RunModel, + *, + run_id: uuid.UUID | None = None, + run_name: str | None = None, + base: str = "Qwen/Qwen3-0.6B", + model: str = "Qwen/Qwen3-0.6B", +) -> AgentProvisioningResult: + return AgentProvisioningResult( + run_id=run_id if run_id is not None else run.id, + run_name=run_name if run_name is not None else run.run_name, + final_report=AgentFinalReport( + success=True, + run_id=run.id, + run_name=run.run_name, + service_yaml=f"type: service\nname: {run.run_name}\n", + base=base, + model=model, + ), + ) + + +def _get_agent_run_name(suffix: str = "1") -> str: + return f"qwen-endpoint-{suffix}" + + +def _make_preset_plan(run_name: str = "qwen-endpoint-serving") -> Mock: + preset_plan = Mock() + preset_plan.preset.base = "Qwen/Qwen3-0.6B" + preset_plan.recipe.id = "vllm-t4" + preset_plan.recipe.model = "Qwen/Qwen3-0.6B" + preset_plan.run_plan.run_spec = RunSpec( + run_name=run_name, + configuration=ServiceConfiguration.parse_obj( + { + "type": "service", + "name": run_name, + "commands": [ + "vllm serve Qwen/Qwen3-0.6B --host 0.0.0.0 --port 8000", + ], + "port": 8000, + "model": "Qwen/Qwen3-0.6B", + "resources": {"gpu": "16GB"}, + } + ), + ) + preset_plan.run_plan.current_resource = None + return preset_plan + + +async def _get_endpoint_run_submissions( + session: AsyncSession, + endpoint_model: EndpointModel, +) -> list[EndpointRunSubmissionModel]: + res = await session.execute( + select(EndpointRunSubmissionModel) + .where(EndpointRunSubmissionModel.endpoint_id == endpoint_model.id) + .order_by(EndpointRunSubmissionModel.submission_num) + ) + return list(res.scalars().all()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) +class TestRecordEndpointRunSubmission: + async def test_records_ordered_submissions_and_is_idempotent( + self, test_db, session: AsyncSession + ): + endpoint_model = await _create_endpoint_model(session=session) + repo = await create_repo( + session=session, + project_id=endpoint_model.project_id, + repo_name="submissions-repo", + ) + run_1 = await create_run( + session=session, + project=endpoint_model.project, + repo=repo, + user=endpoint_model.user, + run_name="qwen-endpoint-1", + ) + run_2 = await create_run( + session=session, + project=endpoint_model.project, + repo=repo, + user=endpoint_model.user, + run_name="qwen-endpoint-2", + ) + + submission_1 = await record_endpoint_run_submission( + session=session, + endpoint_id=endpoint_model.id, + run_id=run_1.id, + ) + duplicate_submission_1 = await record_endpoint_run_submission( + session=session, + endpoint_id=endpoint_model.id, + run_id=run_1.id, + ) + submission_2 = await record_endpoint_run_submission( + session=session, + endpoint_id=endpoint_model.id, + run_id=run_2.id, + ) + await session.commit() + + assert duplicate_submission_1 is submission_1 + assert submission_1.submission_num == 1 + assert submission_2.submission_num == 2 + submissions = await _get_endpoint_run_submissions( + session=session, + endpoint_model=endpoint_model, + ) + assert [submission.run_id for submission in submissions] == [run_1.id, run_2.id] + assert [submission.submission_num for submission in submissions] == [1, 2] + + async def test_rejects_run_already_recorded_for_another_endpoint( + self, test_db, session: AsyncSession + ): + endpoint_model = await _create_endpoint_model(session=session) + other_configuration = EndpointConfiguration( + name="other-qwen-endpoint", + model=EndpointModelRepo(repo="Qwen/Qwen3-0.6B"), + env=Env.parse_obj({"HF_TOKEN": "secret"}), + ) + other_endpoint_model = EndpointModel( + project=endpoint_model.project, + user=endpoint_model.user, + name=other_configuration.name, + status=EndpointStatus.SUBMITTED, + configuration=other_configuration.json(), + created_at=datetime(2023, 1, 2, 3, 5, tzinfo=timezone.utc), + last_processed_at=datetime(2023, 1, 2, 3, 5, tzinfo=timezone.utc), + ) + session.add(other_endpoint_model) + await session.commit() + repo = await create_repo( + session=session, + project_id=endpoint_model.project_id, + repo_name="recorded-run-repo", + ) + run = await create_run( + session=session, + project=endpoint_model.project, + repo=repo, + user=endpoint_model.user, + run_name="recorded-run", + ) + await record_endpoint_run_submission( + session=session, + endpoint_id=endpoint_model.id, + run_id=run.id, + ) + await session.commit() + + with pytest.raises(ServerClientError, match="Run is already recorded"): + await record_endpoint_run_submission( + session=session, + endpoint_id=other_endpoint_model.id, + run_id=run.id, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) +class TestEndpointWorkerSubmitted: + async def test_fails_when_preset_run_name_exists_without_link( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model(session=session) + await create_fleet(session=session, project=endpoint_model.project) + run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + status=RunStatus.PROVISIONING, + link_endpoint=False, + ) + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "find_preset_planning_result", + new=AsyncMock( + return_value=EndpointPresetPlanningResult(provisionable=_make_preset_plan()) + ), + ) as find_preset_planning_result_mock, + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.runs_services.apply_plan", + new=AsyncMock(), + ) as apply_plan_mock, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(run) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message == ( + "Run name 'qwen-endpoint-serving' is taken by an existing run" + ) + assert endpoint_model.service_run_id is None + assert run.status == RunStatus.PROVISIONING + assert run.deleted is False + find_preset_planning_result_mock.assert_awaited_once() + apply_plan_mock.assert_not_awaited() + events = await list_events(session) + assert len(events) == 1 + assert ( + events[0].message == "Endpoint status changed SUBMITTED -> FAILED " + "(Run name 'qwen-endpoint-serving' is taken by an existing run)" + ) + + async def test_fails_when_preset_run_name_is_taken_by_another_user( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model(session=session) + await create_fleet(session=session, project=endpoint_model.project) + repo = await create_repo(session=session, project_id=endpoint_model.project_id) + another_user = await create_user(session=session, name="another-user") + run = await create_run( + session=session, + project=endpoint_model.project, + repo=repo, + user=another_user, + run_name=get_endpoint_serving_run_name(endpoint_model.name), + status=RunStatus.RUNNING, + ) + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "find_preset_planning_result", + new=AsyncMock( + return_value=EndpointPresetPlanningResult(provisionable=_make_preset_plan()) + ), + ) as find_preset_planning_result_mock, + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.runs_services.apply_plan", + new=AsyncMock(), + ) as apply_plan_mock, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(run) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message == ( + "Run name 'qwen-endpoint-serving' is taken by an existing run" + ) + assert endpoint_model.service_run_id is None + assert run.status == RunStatus.RUNNING + assert run.deleted is False + find_preset_planning_result_mock.assert_awaited_once() + apply_plan_mock.assert_not_awaited() + events = await list_events(session) + assert len(events) == 1 + assert ( + events[0].message == "Endpoint status changed SUBMITTED -> FAILED " + "(Run name 'qwen-endpoint-serving' is taken by an existing run)" + ) + + async def test_fails_when_preset_run_name_is_taken_by_pre_existing_run( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model(session=session) + await create_fleet(session=session, project=endpoint_model.project) + repo = await create_repo(session=session, project_id=endpoint_model.project_id) + run = await create_run( + session=session, + project=endpoint_model.project, + repo=repo, + user=endpoint_model.user, + run_name=get_endpoint_serving_run_name(endpoint_model.name), + status=RunStatus.RUNNING, + submitted_at=datetime(2023, 1, 2, 3, 3, tzinfo=timezone.utc), + ) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "find_preset_planning_result", + new=AsyncMock( + return_value=EndpointPresetPlanningResult(provisionable=_make_preset_plan()) + ), + ) as find_preset_planning_result_mock: + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(run) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message == ( + "Run name 'qwen-endpoint-serving' is taken by an existing run" + ) + assert endpoint_model.service_run_id is None + assert run.status == RunStatus.RUNNING + assert run.deleted is False + find_preset_planning_result_mock.assert_awaited_once() + + async def test_preset_run_name_conflict_does_not_block_prototyping_without_preset( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model(session=session) + await create_fleet(session=session, project=endpoint_model.project) + run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + status=RunStatus.RUNNING, + link_endpoint=False, + ) + agent_service = _FakeAgentService() + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "find_preset_planning_result", + new=AsyncMock(return_value=EndpointPresetPlanningResult()), + ) as find_preset_planning_result_mock, + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ), + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(run) + assert endpoint_model.status == EndpointStatus.PROTOTYPING + assert endpoint_model.status_message is None + assert endpoint_model.service_run_id is None + assert endpoint_model.provisioning_method == "agent" + assert run.status == RunStatus.RUNNING + assert run.deleted is False + find_preset_planning_result_mock.assert_awaited_once() + agent_service.provision_endpoint.assert_not_awaited() + + async def test_finished_same_name_run_does_not_block_preset_submission( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + user_ssh_public_key="ssh-rsa test", + ) + await create_fleet(session=session, project=endpoint_model.project) + run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + status=RunStatus.DONE, + link_endpoint=False, + ) + preset_plan = _make_preset_plan() + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "find_preset_planning_result", + new=AsyncMock( + return_value=EndpointPresetPlanningResult(provisionable=preset_plan) + ), + ) as find_preset_planning_result_mock, + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.runs_services.apply_plan", + new=AsyncMock(return_value=Mock(id=run.id)), + ) as apply_plan_mock, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(run) + assert endpoint_model.status == EndpointStatus.PROVISIONING + assert endpoint_model.status_message is None + assert endpoint_model.provisioning_method == "preset:Qwen/Qwen3-0.6B#vllm-t4" + assert run.status == RunStatus.DONE + assert run.deleted is False + find_preset_planning_result_mock.assert_awaited_once() + apply_plan_mock.assert_awaited_once() + + async def test_submitted_to_failed_without_provisioning_path( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model(session=session) + await create_fleet(session=session, project=endpoint_model.project) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "find_preset_planning_result", + new=AsyncMock(return_value=EndpointPresetPlanningResult()), + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.FAILED + assert ( + endpoint_model.status_message == "No matching endpoint presets found. " + "Creating a preset requires the server agent, but " + "DSTACK_AGENT_ANTHROPIC_API_KEY is not set." + ) + events = await list_events(session) + assert len(events) == 1 + assert ( + events[0].message == "Endpoint status changed SUBMITTED -> FAILED " + "(No matching endpoint presets found. " + "Creating a preset requires the server agent, but " + "DSTACK_AGENT_ANTHROPIC_API_KEY is not set.)" + ) + + async def test_submitted_to_prototyping_with_agent( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model(session=session) + await create_fleet(session=session, project=endpoint_model.project) + agent_service = _FakeAgentService() + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "find_preset_planning_result", + new=AsyncMock(return_value=EndpointPresetPlanningResult()), + ), + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ), + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.PROTOTYPING + assert endpoint_model.status_message is None + assert endpoint_model.service_run_id is None + assert endpoint_model.provisioning_method == "agent" + agent_service.provision_endpoint.assert_not_awaited() + events = await list_events(session) + assert len(events) == 1 + assert events[0].message == "Endpoint status changed SUBMITTED -> PROTOTYPING" + + async def test_submitted_to_failed_without_fleets_even_when_agent_enabled( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model(session=session) + agent_service = _FakeAgentService() + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "find_preset_planning_result", + new=AsyncMock(return_value=EndpointPresetPlanningResult()), + ), + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ), + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message == ( + "The project has no fleets. Create one before submitting an endpoint." + ) + assert endpoint_model.service_run_id is None + assert endpoint_model.provisioning_method is None + agent_service.provision_endpoint.assert_not_awaited() + + async def test_submitted_reuse_to_failed_without_fleets( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model(session=session) + configuration = EndpointConfiguration.__response__.parse_raw(endpoint_model.configuration) + configuration.preset_policy = EndpointPresetPolicy.REUSE + endpoint_model.configuration = configuration.json() + await session.commit() + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "find_preset_planning_result", + new=AsyncMock(return_value=EndpointPresetPlanningResult()), + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message == ( + "The project has no fleets. Create one before submitting an endpoint." + ) + assert endpoint_model.service_run_id is None + assert endpoint_model.provisioning_method is None + + async def test_project_user_cannot_start_agent( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + user_global_role=GlobalRole.USER, + project_role=ProjectRole.USER, + ) + await create_fleet(session=session, project=endpoint_model.project) + agent_service = _FakeAgentService() + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "find_preset_planning_result", + new=AsyncMock(return_value=EndpointPresetPlanningResult()), + ), + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ), + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message == ( + "No matching endpoint presets found. " + "Creating endpoint presets with the server agent requires project admin permissions." + ) + assert endpoint_model.service_run_id is None + assert endpoint_model.provisioning_method is None + agent_service.provision_endpoint.assert_not_awaited() + + async def test_submitted_to_provisioning_with_matching_preset( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + user_ssh_public_key="ssh-rsa test", + ) + await create_fleet(session=session, project=endpoint_model.project) + repo = await create_repo(session=session, project_id=endpoint_model.project_id) + run = await create_run( + session=session, + project=endpoint_model.project, + repo=repo, + user=endpoint_model.user, + run_name="qwen-endpoint-submitted", + ) + preset_plan = _make_preset_plan() + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "find_preset_planning_result", + new=AsyncMock( + return_value=EndpointPresetPlanningResult(provisionable=preset_plan) + ), + ) as find_preset_planning_result_mock, + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.runs_services.apply_plan", + new=AsyncMock(return_value=Mock(id=run.id)), + ) as apply_plan_mock, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.PROVISIONING + assert endpoint_model.status_message is None + assert endpoint_model.service_run_id == run.id + assert endpoint_model.provisioning_method == "preset:Qwen/Qwen3-0.6B#vllm-t4" + submissions = await _get_endpoint_run_submissions( + session=session, + endpoint_model=endpoint_model, + ) + assert len(submissions) == 1 + assert submissions[0].run_id == run.id + assert submissions[0].submission_num == 1 + find_preset_planning_result_mock.assert_awaited_once() + apply_plan_mock.assert_awaited_once() + events = await list_events(session) + assert len(events) == 1 + assert events[0].message == "Endpoint status changed SUBMITTED -> PROVISIONING" + + async def test_submitted_to_failed_when_preset_submission_fails( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + user_ssh_public_key="ssh-rsa test", + ) + await create_fleet(session=session, project=endpoint_model.project) + preset_plan = _make_preset_plan() + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "find_preset_planning_result", + new=AsyncMock( + return_value=EndpointPresetPlanningResult(provisionable=preset_plan) + ), + ), + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.runs_services.apply_plan", + new=AsyncMock(side_effect=ServerClientError("Cannot override running run")), + ) as apply_plan_mock, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message == "Cannot override running run" + apply_plan_mock.assert_awaited_once() + events = await list_events(session) + assert len(events) == 1 + assert ( + events[0].message + == "Endpoint status changed SUBMITTED -> FAILED (Cannot override running run)" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) +class TestEndpointWorkerProvisioning: + async def test_prototyping_does_not_fail_on_legacy_same_name_run_conflict( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + legacy_run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + status=RunStatus.PROVISIONING, + link_endpoint=False, + ) + agent_run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + link_endpoint=False, + run_name=_get_agent_run_name(), + ) + agent_service = _FakeAgentService(result=_get_verified_agent_result(agent_run)) + preset_service = Mock() + preset_service.save_preset = AsyncMock( + side_effect=lambda project_name, preset, comments: preset + ) + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ), + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "get_endpoint_preset_service", + return_value=preset_service, + ), + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(legacy_run) + await session.refresh(agent_run) + assert endpoint_model.status == EndpointStatus.RUNNING + assert endpoint_model.status_message is None + assert endpoint_model.service_run_id == agent_run.id + assert legacy_run.status == RunStatus.PROVISIONING + assert legacy_run.deleted is False + assert agent_run.deleted is False + events = await list_events(session) + assert len(events) == 1 + assert events[0].message == "Endpoint status changed PROTOTYPING -> RUNNING" + + async def test_fails_when_agent_reports_foreign_run( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + repo = await create_repo(session=session, project_id=endpoint_model.project_id) + another_user = await create_user(session=session, name="another-user") + run = await create_run( + session=session, + project=endpoint_model.project, + repo=repo, + user=another_user, + run_name=_get_agent_run_name(), + status=RunStatus.RUNNING, + ) + agent_service = _FakeAgentService(result=_get_verified_agent_result(run)) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(run) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message == ( + "Run 'qwen-endpoint-1' is not owned by the endpoint user" + ) + assert endpoint_model.service_run_id is None + assert run.status == RunStatus.RUNNING + assert run.deleted is False + events = await list_events(session) + assert len(events) == 1 + assert ( + events[0].message == "Endpoint status changed PROTOTYPING -> FAILED " + "(Run 'qwen-endpoint-1' is not owned by the endpoint user)" + ) + + async def test_fails_when_agent_final_run_name_is_not_submission_number( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + link_endpoint=False, + run_name="agent-run", + ) + await session.commit() + agent_service = _FakeAgentService(result=_get_verified_agent_result(run)) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message == ( + "Server agent final service run name must be 'qwen-endpoint-'" + ) + assert endpoint_model.service_run_id is None + assert run.status == RunStatus.RUNNING + + async def test_agent_creates_ready_service_and_endpoint_becomes_running( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model(session=session) + await create_fleet(session=session, project=endpoint_model.project) + preset_service = Mock() + preset_service.save_preset = AsyncMock( + side_effect=lambda project_name, preset, comments: preset + ) + + async def provision_endpoint(endpoint_model, pipeline_hinter): + return await _create_ready_backing_service_run_for_agent(endpoint_model.id) + + agent_service = _FakeAgentService(side_effect=provision_endpoint) + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "find_preset_planning_result", + new=AsyncMock(return_value=EndpointPresetPlanningResult()), + ), + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ), + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "get_endpoint_preset_service", + return_value=preset_service, + ), + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.PROTOTYPING + await _lock_endpoint_model(session=session, endpoint_model=endpoint_model) + + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.RUNNING + assert endpoint_model.status_message is None + assert endpoint_model.service_run_id is not None + submissions = await _get_endpoint_run_submissions( + session=session, + endpoint_model=endpoint_model, + ) + assert len(submissions) == 1 + assert submissions[0].run_id == endpoint_model.service_run_id + assert submissions[0].submission_num == 1 + agent_service.provision_endpoint.assert_awaited_once() + preset_service.save_preset.assert_awaited_once() + + async def test_agent_success_after_stop_keeps_endpoint_stopping_and_links_run( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + endpoint_model.provisioning_method = "agent" + await session.commit() + created_run_id = None + + async def provision_endpoint(endpoint_model, pipeline_hinter): + nonlocal created_run_id + async with get_session_ctx() as provisioning_session: + res = await provisioning_session.execute( + select(EndpointModel) + .where(EndpointModel.id == endpoint_model.id) + .options(joinedload(EndpointModel.project)) + .options(joinedload(EndpointModel.user)) + ) + provisioning_endpoint = res.unique().scalar_one() + run = await _create_backing_service_run( + session=provisioning_session, + endpoint_model=provisioning_endpoint, + link_endpoint=False, + run_name=_get_agent_run_name(), + ) + provisioning_endpoint.status = EndpointStatus.STOPPING + await provisioning_session.commit() + created_run_id = run.id + return _get_verified_agent_result(run) + + agent_service = _FakeAgentService(side_effect=provision_endpoint) + preset_service = Mock() + preset_service.save_preset = AsyncMock( + side_effect=lambda project_name, preset, comments: preset + ) + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ), + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "get_endpoint_preset_service", + return_value=preset_service, + ), + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + assert created_run_id is not None + await session.refresh(endpoint_model) + run = await session.get(RunModel, created_run_id) + assert run is not None + assert endpoint_model.status == EndpointStatus.STOPPING + assert endpoint_model.service_run_id == run.id + assert run.status == RunStatus.RUNNING + assert preset_service.save_preset.await_count == 1 + events = await list_events(session) + assert events == [] + + await _lock_endpoint_model(session=session, endpoint_model=endpoint_model) + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(run) + assert endpoint_model.status == EndpointStatus.STOPPING + assert run.status == RunStatus.TERMINATING + + async def test_agent_reported_run_id_links_backing_service_run( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + endpoint_model.provisioning_method = "agent" + await session.commit() + + async def provision_endpoint(endpoint_model, pipeline_hinter): + async with get_session_ctx() as provisioning_session: + res = await provisioning_session.execute( + select(EndpointModel) + .where(EndpointModel.id == endpoint_model.id) + .options(joinedload(EndpointModel.project)) + .options(joinedload(EndpointModel.user)) + ) + provisioning_endpoint = res.unique().scalar_one() + run = await _create_backing_service_run( + session=provisioning_session, + endpoint_model=provisioning_endpoint, + link_endpoint=False, + run_name=_get_agent_run_name(), + ) + return AgentProvisioningResult( + run_id=run.id, + final_report=AgentFinalReport( + success=True, + run_id=run.id, + run_name=run.run_name, + service_yaml=f"type: service\nname: {run.run_name}\n", + base="Qwen/Qwen3-0.6B", + model="Qwen/Qwen3-0.6B", + ), + ) + + agent_service = _FakeAgentService(side_effect=provision_endpoint) + preset_service = Mock() + preset_service.save_preset = AsyncMock( + side_effect=lambda project_name, preset, comments: preset + ) + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ), + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "get_endpoint_preset_service", + return_value=preset_service, + ), + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.RUNNING + assert endpoint_model.service_run_id is not None + submissions = await _get_endpoint_run_submissions( + session=session, + endpoint_model=endpoint_model, + ) + assert len(submissions) == 1 + assert submissions[0].run_id == endpoint_model.service_run_id + preset_service.save_preset.assert_awaited_once() + + async def test_agent_base_model_report_sets_endpoint_model_repo_and_preset_recipe_model( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + model=EndpointModelBase(base="Qwen/Qwen3-0.6B"), + ) + endpoint_model.provisioning_method = "agent" + await session.commit() + selected_model = "groxaxo/Qwen3-0.6B-GPTQ-4Bit" + + async def provision_endpoint(endpoint_model, pipeline_hinter): + async with get_session_ctx() as provisioning_session: + res = await provisioning_session.execute( + select(EndpointModel) + .where(EndpointModel.id == endpoint_model.id) + .options(joinedload(EndpointModel.project)) + .options(joinedload(EndpointModel.user)) + ) + provisioning_endpoint = res.unique().scalar_one() + run = await _create_backing_service_run( + session=provisioning_session, + endpoint_model=provisioning_endpoint, + link_endpoint=False, + run_name=_get_agent_run_name(), + ) + return AgentProvisioningResult( + run_id=run.id, + final_report=AgentFinalReport( + success=True, + run_id=run.id, + run_name=run.run_name, + service_yaml=f"type: service\nname: {run.run_name}\n", + base="Qwen/Qwen3-0.6B", + model=selected_model, + ), + ) + + agent_service = _FakeAgentService(side_effect=provision_endpoint) + preset_service = Mock() + preset_service.save_preset = AsyncMock( + side_effect=lambda project_name, preset, comments: preset + ) + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ), + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "get_endpoint_preset_service", + return_value=preset_service, + ), + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.RUNNING + assert endpoint_model.model_base == "Qwen/Qwen3-0.6B" + assert endpoint_model.model_repo == selected_model + saved_preset = preset_service.save_preset.await_args.args[1] + assert saved_preset.base == "Qwen/Qwen3-0.6B" + assert saved_preset.recipes[0].model == selected_model + + async def test_agent_exact_model_report_mismatch_fails_endpoint( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + endpoint_model.provisioning_method = "agent" + await session.commit() + + async def provision_endpoint(endpoint_model, pipeline_hinter): + async with get_session_ctx() as provisioning_session: + res = await provisioning_session.execute( + select(EndpointModel) + .where(EndpointModel.id == endpoint_model.id) + .options(joinedload(EndpointModel.project)) + .options(joinedload(EndpointModel.user)) + ) + provisioning_endpoint = res.unique().scalar_one() + run = await _create_backing_service_run( + session=provisioning_session, + endpoint_model=provisioning_endpoint, + link_endpoint=False, + run_name=_get_agent_run_name(), + ) + return AgentProvisioningResult( + run_id=run.id, + final_report=AgentFinalReport( + success=True, + run_id=run.id, + run_name=run.run_name, + service_yaml=f"type: service\nname: {run.run_name}\n", + base="Qwen/Qwen3-0.6B", + model="groxaxo/Qwen3-0.6B-GPTQ-4Bit", + ), + ) + + agent_service = _FakeAgentService(side_effect=provision_endpoint) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.service_run_id is None + assert endpoint_model.model_repo is None + assert endpoint_model.status_message == ( + "Server agent final report model does not match the requested model repo" + ) + + async def test_agent_success_stops_non_final_submitted_runs( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + endpoint_model.provisioning_method = "agent" + probe_run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + link_endpoint=False, + run_name=f"{endpoint_model.name}-1", + ) + final_run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + link_endpoint=False, + run_name=f"{endpoint_model.name}-2", + ) + endpoint_model.service_run_id = None + await session.commit() + agent_service = _FakeAgentService( + result=AgentProvisioningResult( + run_id=final_run.id, + run_name=final_run.run_name, + submitted_run_ids=(probe_run.id, final_run.id), + final_report=AgentFinalReport( + success=True, + run_id=final_run.id, + run_name=final_run.run_name, + service_yaml=f"type: service\nname: {final_run.run_name}\n", + base="Qwen/Qwen3-0.6B", + model="Qwen/Qwen3-0.6B", + ), + ) + ) + preset_service = Mock() + preset_service.save_preset = AsyncMock( + side_effect=lambda project_name, preset, comments: preset + ) + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ), + patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "get_endpoint_preset_service", + return_value=preset_service, + ), + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(probe_run) + await session.refresh(final_run) + assert endpoint_model.status == EndpointStatus.RUNNING + assert endpoint_model.service_run_id == final_run.id + assert probe_run.status == RunStatus.TERMINATING + assert final_run.status == RunStatus.RUNNING + submissions = await _get_endpoint_run_submissions( + session=session, + endpoint_model=endpoint_model, + ) + assert [submission.run_id for submission in submissions] == [probe_run.id, final_run.id] + assert [submission.submission_num for submission in submissions] == [1, 2] + preset_service.save_preset.assert_awaited_once() + + async def test_prototyping_linked_run_waits_without_showing_provisioning( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + endpoint_model.provisioning_method = "agent" + await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + status=RunStatus.PROVISIONING, + ) + + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.PROTOTYPING + assert endpoint_model.status_message is None + events = await list_events(session) + assert events == [] + + async def test_agent_reported_run_without_verification_report_fails_endpoint( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + endpoint_model.provisioning_method = "agent" + await session.commit() + + async def provision_endpoint(endpoint_model, pipeline_hinter): + async with get_session_ctx() as provisioning_session: + res = await provisioning_session.execute( + select(EndpointModel) + .where(EndpointModel.id == endpoint_model.id) + .options(joinedload(EndpointModel.project)) + .options(joinedload(EndpointModel.user)) + ) + provisioning_endpoint = res.unique().scalar_one() + run = await _create_backing_service_run( + session=provisioning_session, + endpoint_model=provisioning_endpoint, + link_endpoint=False, + ) + return AgentProvisioningResult(run_id=run.id) + + agent_service = _FakeAgentService(side_effect=provision_endpoint) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.service_run_id is None + assert endpoint_model.status_message == "Server agent did not return a final report" + + async def test_agent_failure_fails_endpoint( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + endpoint_model.provisioning_method = "agent" + await session.commit() + agent_service = _FakeAgentService( + result=AgentProvisioningResult(error="agent could not find a deployable recipe") + ) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message == "agent could not find a deployable recipe" + agent_service.provision_endpoint.assert_awaited_once() + + async def test_agent_in_progress_keeps_endpoint_prototyping( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + endpoint_model.provisioning_method = "agent" + await session.commit() + agent_service = _FakeAgentService(result=AgentProvisioningResult(in_progress=True)) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.PROTOTYPING + assert endpoint_model.status_message is None + assert endpoint_model.service_run_id is None + agent_service.provision_endpoint.assert_awaited_once() + + async def test_agent_in_progress_records_submitted_run( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + endpoint_model.provisioning_method = "agent" + submitted_run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + link_endpoint=False, + run_name=_get_agent_run_name("1"), + ) + endpoint_model.service_run_id = None + await session.commit() + agent_service = _FakeAgentService( + result=AgentProvisioningResult( + in_progress=True, + submitted_run_ids=(submitted_run.id,), + ) + ) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.PROTOTYPING + assert endpoint_model.service_run_id is None + submissions = await _get_endpoint_run_submissions( + session=session, + endpoint_model=endpoint_model, + ) + assert len(submissions) == 1 + assert submissions[0].run_id == submitted_run.id + + async def test_agent_in_progress_records_submitted_run_by_name( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + endpoint_model.provisioning_method = "agent" + submitted_run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + link_endpoint=False, + run_name=f"{endpoint_model.name}-1", + ) + endpoint_model.service_run_id = None + await session.commit() + agent_service = _FakeAgentService( + result=AgentProvisioningResult( + in_progress=True, + submitted_run_names=(submitted_run.run_name,), + ) + ) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.PROTOTYPING + assert endpoint_model.service_run_id is None + submissions = await _get_endpoint_run_submissions( + session=session, + endpoint_model=endpoint_model, + ) + assert len(submissions) == 1 + assert submissions[0].run_id == submitted_run.id + + async def test_agent_in_progress_ignores_non_strict_submitted_run_name( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + endpoint_model.provisioning_method = "agent" + submitted_run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + link_endpoint=False, + run_name="agent-run", + ) + endpoint_model.service_run_id = None + await session.commit() + agent_service = _FakeAgentService( + result=AgentProvisioningResult( + in_progress=True, + submitted_run_names=(submitted_run.run_name,), + ) + ) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + submissions = await _get_endpoint_run_submissions( + session=session, + endpoint_model=endpoint_model, + ) + assert submissions == [] + + async def test_agent_error_status_message_is_compact( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + endpoint_model.provisioning_method = "agent" + await session.commit() + long_error = "agent failed before report\n" + "\n".join( + f"offer {i:04d} gpu=A5000 price=0.27" for i in range(200) + ) + agent_service = _FakeAgentService(result=AgentProvisioningResult(error=long_error)) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message is not None + assert len(endpoint_model.status_message) <= 500 + assert "\n" not in endpoint_model.status_message + assert endpoint_model.status_message.startswith("agent failed before report offer 0000") + assert "offer 0199" not in endpoint_model.status_message + + async def test_agent_failure_report_status_message_is_compact( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + endpoint_model.provisioning_method = "agent" + await session.commit() + long_failure_summary = "agent could not verify service\n" + "\n".join( + f"line {i:04d} with detailed output" for i in range(200) + ) + agent_service = _FakeAgentService( + result=AgentProvisioningResult( + final_report=AgentFinalReport( + success=False, + failure_summary=long_failure_summary, + ) + ) + ) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message is not None + assert len(endpoint_model.status_message) <= 500 + assert "\n" not in endpoint_model.status_message + assert endpoint_model.status_message.startswith("agent could not verify service line 0000") + assert "line 0199" not in endpoint_model.status_message + + async def test_agent_failure_does_not_stop_unlinked_same_name_run( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + endpoint_model.provisioning_method = "agent" + await session.commit() + created_run_id = None + + async def provision_endpoint(endpoint_model, pipeline_hinter): + nonlocal created_run_id + async with get_session_ctx() as agent_session: + res = await agent_session.execute( + select(EndpointModel) + .where(EndpointModel.id == endpoint_model.id) + .options(joinedload(EndpointModel.project)) + .options(joinedload(EndpointModel.user)) + ) + endpoint = res.unique().scalar_one() + run = await _create_backing_service_run( + session=agent_session, + endpoint_model=endpoint, + link_endpoint=False, + ) + created_run_id = run.id + return AgentProvisioningResult(error="agent could not verify the service") + + agent_service = _FakeAgentService(side_effect=provision_endpoint) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + assert created_run_id is not None + await session.refresh(endpoint_model) + run = await session.get(RunModel, created_run_id) + assert run is not None + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message == "agent could not verify the service" + assert endpoint_model.service_run_id is None + assert run.status == RunStatus.RUNNING + agent_service.provision_endpoint.assert_awaited_once() + + async def test_agent_failure_stops_recorded_submitted_run( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROTOTYPING, + ) + endpoint_model.provisioning_method = "agent" + submitted_run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + link_endpoint=False, + run_name=_get_agent_run_name("1"), + ) + endpoint_model.service_run_id = None + await session.commit() + agent_service = _FakeAgentService( + result=AgentProvisioningResult( + error="agent could not verify the service", + submitted_run_ids=(submitted_run.id,), + ) + ) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.get_agent_service", + return_value=agent_service, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(submitted_run) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message == "agent could not verify the service" + assert endpoint_model.service_run_id is None + assert submitted_run.status == RunStatus.TERMINATING + submissions = await _get_endpoint_run_submissions( + session=session, + endpoint_model=endpoint_model, + ) + assert len(submissions) == 1 + assert submissions[0].run_id == submitted_run.id + + async def test_waits_when_backing_run_is_not_ready( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROVISIONING, + ) + await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + status=RunStatus.PROVISIONING, + ) + + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.PROVISIONING + assert endpoint_model.status_message is None + events = await list_events(session) + assert events == [] + + async def test_moves_to_running_when_backing_service_is_ready( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROVISIONING, + ) + await _create_backing_service_run(session=session, endpoint_model=endpoint_model) + + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.RUNNING + assert endpoint_model.status_message is None + events = await list_events(session) + assert len(events) == 1 + assert events[0].message == "Endpoint status changed PROVISIONING -> RUNNING" + + async def test_saves_preset_when_agent_backing_service_becomes_running( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROVISIONING, + ) + endpoint_model.provisioning_method = "agent" + await _create_backing_service_run(session=session, endpoint_model=endpoint_model) + preset_service = Mock() + preset_service.save_preset = AsyncMock( + side_effect=lambda project_name, preset, comments: preset + ) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "get_endpoint_preset_service", + return_value=preset_service, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.RUNNING + preset_service.save_preset.assert_awaited_once() + assert preset_service.save_preset.await_args.args[0] == "test_project" + saved_preset = preset_service.save_preset.await_args.args[1] + assert saved_preset.base == "Qwen/Qwen3-0.6B" + assert len(saved_preset.recipes) == 1 + assert saved_preset.recipes[0].service.resources.gpu is not None + assert saved_preset.recipes[0].validations[0].replicas[0].resources[0].gpu is not None + comments = preset_service.save_preset.await_args.kwargs["comments"] + assert f"endpoint: {endpoint_model.name}" in comments + + async def test_preset_save_failure_does_not_block_agent_endpoint_running( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROVISIONING, + ) + endpoint_model.provisioning_method = "agent" + await _create_backing_service_run(session=session, endpoint_model=endpoint_model) + preset_service = Mock() + preset_service.save_preset = AsyncMock(side_effect=OSError("read-only preset dir")) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints." + "get_endpoint_preset_service", + return_value=preset_service, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.RUNNING + preset_service.save_preset.assert_awaited_once() + + async def test_waits_when_backing_service_has_no_registered_jobs( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROVISIONING, + ) + await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + registered=False, + ) + + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.PROVISIONING + assert endpoint_model.status_message is None + events = await list_events(session) + assert events == [] + + async def test_fails_when_backing_run_finished( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROVISIONING, + ) + run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + status=RunStatus.FAILED, + ) + + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(run) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message == "Backing service run finished with status failed" + assert run.status == RunStatus.FAILED + events = await list_events(session) + assert len(events) == 1 + assert ( + events[0].message == "Endpoint status changed PROVISIONING -> FAILED " + "(Backing service run finished with status failed)" + ) + + async def test_stops_running_backing_run_when_endpoint_fails( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.PROVISIONING, + ) + run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + ) + run.service_spec = "invalid" + await session.commit() + + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(run) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message == "Backing service spec is invalid" + assert run.status == RunStatus.TERMINATING + events = await list_events(session) + messages = [event.message for event in events] + assert any( + message.startswith("Run status changed RUNNING -> TERMINATING") for message in messages + ) + assert ( + "Endpoint status changed PROVISIONING -> FAILED " + "(Backing service spec is invalid)" in messages + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) +class TestEndpointWorkerRunning: + async def test_ready_backing_service_keeps_endpoint_running( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.RUNNING, + ) + endpoint_model.status_message = "previous failure" + await _create_backing_service_run(session=session, endpoint_model=endpoint_model) + + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.RUNNING + assert endpoint_model.status_message is None + events = await list_events(session) + assert events == [] + + async def test_not_ready_backing_service_keeps_endpoint_running( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.RUNNING, + ) + await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + registered=False, + ) + + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.RUNNING + assert endpoint_model.status_message is None + events = await list_events(session) + assert events == [] + + async def test_finished_backing_run_fails_endpoint( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.RUNNING, + ) + run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + status=RunStatus.FAILED, + ) + + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(run) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message == "Backing service run finished with status failed" + assert run.status == RunStatus.FAILED + events = await list_events(session) + assert len(events) == 1 + assert ( + events[0].message == "Endpoint status changed RUNNING -> FAILED " + "(Backing service run finished with status failed)" + ) + + async def test_stops_running_backing_run_when_endpoint_fails( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.RUNNING, + ) + run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + ) + run.service_spec = "invalid" + await session.commit() + + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(run) + assert endpoint_model.status == EndpointStatus.FAILED + assert endpoint_model.status_message == "Backing service spec is invalid" + assert run.status == RunStatus.TERMINATING + events = await list_events(session) + messages = [event.message for event in events] + assert any( + message.startswith("Run status changed RUNNING -> TERMINATING") for message in messages + ) + assert ( + "Endpoint status changed RUNNING -> FAILED (Backing service spec is invalid)" + in messages + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) +class TestEndpointWorkerStopping: + async def test_marks_endpoint_stopped_without_backing_run( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.STOPPING, + ) + + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.STOPPED + assert endpoint_model.status_message is None + events = await list_events(session) + assert events[0].message == "Endpoint status changed STOPPING -> STOPPED" + + async def test_aborts_agent_before_marking_endpoint_stopped( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.STOPPING, + ) + endpoint_model.provisioning_method = "agent" + await session.commit() + abort_agent_endpoint = AsyncMock(return_value=True) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.abort_agent_endpoint", + abort_agent_endpoint, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + abort_agent_endpoint.assert_awaited_once() + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.STOPPED + events = await list_events(session) + assert events[0].message == "Endpoint status changed STOPPING -> STOPPED" + + async def test_waits_when_agent_abort_is_still_pending( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.STOPPING, + ) + endpoint_model.provisioning_method = "agent" + await session.commit() + abort_agent_endpoint = AsyncMock(return_value=False) + + with patch( + "dstack._internal.server.background.pipeline_tasks.endpoints.abort_agent_endpoint", + abort_agent_endpoint, + ): + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + abort_agent_endpoint.assert_awaited_once() + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.STOPPING + events = await list_events(session) + assert events == [] + + async def test_stops_backing_run_before_stopping_endpoint( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.STOPPING, + ) + run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + status=RunStatus.RUNNING, + ) + + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(run) + assert endpoint_model.status == EndpointStatus.STOPPING + assert run.status == RunStatus.TERMINATING + assert run.deleted is False + + async def test_stops_recorded_submitted_run_before_stopping_endpoint( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.STOPPING, + ) + run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + status=RunStatus.RUNNING, + link_endpoint=False, + run_name=_get_agent_run_name("1"), + ) + await record_endpoint_run_submission( + session=session, + endpoint_id=endpoint_model.id, + run_id=run.id, + ) + await session.commit() + + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(run) + assert endpoint_model.status == EndpointStatus.STOPPING + assert endpoint_model.service_run_id is None + assert run.status == RunStatus.TERMINATING + + async def test_waits_for_terminating_recorded_submitted_run_before_stopping_endpoint( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.STOPPING, + ) + run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + status=RunStatus.TERMINATING, + link_endpoint=False, + run_name=_get_agent_run_name("1"), + ) + await record_endpoint_run_submission( + session=session, + endpoint_id=endpoint_model.id, + run_id=run.id, + ) + await session.commit() + + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(run) + assert endpoint_model.status == EndpointStatus.STOPPING + assert run.status == RunStatus.TERMINATING + events = await list_events(session) + assert events == [] + + async def test_marks_endpoint_stopped_after_backing_run_finishes( + self, test_db, session: AsyncSession, worker: EndpointWorker + ): + endpoint_model = await _create_endpoint_model( + session=session, + status=EndpointStatus.STOPPING, + ) + run = await _create_backing_service_run( + session=session, + endpoint_model=endpoint_model, + status=RunStatus.TERMINATED, + ) + + await worker.process(_endpoint_to_pipeline_item(endpoint_model)) + + await session.refresh(endpoint_model) + await session.refresh(run) + assert endpoint_model.status == EndpointStatus.STOPPED + assert run.deleted is False + events = await list_events(session) + assert events[0].message == "Endpoint status changed STOPPING -> STOPPED" diff --git a/src/tests/_internal/server/routers/test_endpoints.py b/src/tests/_internal/server/routers/test_endpoints.py new file mode 100644 index 0000000000..436a341a9f --- /dev/null +++ b/src/tests/_internal/server/routers/test_endpoints.py @@ -0,0 +1,1073 @@ +from datetime import datetime, timezone +from unittest.mock import AsyncMock, Mock, patch +from uuid import UUID + +import pytest +from freezegun import freeze_time +from httpx import AsyncClient +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.configurations import ServiceConfiguration +from dstack._internal.core.models.endpoints import ( + EndpointConfiguration, + EndpointStatus, +) +from dstack._internal.core.models.instances import ( + InstanceAvailability, + InstanceOfferWithAvailability, +) +from dstack._internal.core.models.runs import RunSpec, RunStatus +from dstack._internal.core.models.users import GlobalRole, ProjectRole +from dstack._internal.server.models import EndpointModel +from dstack._internal.server.services.endpoints.agent import AgentPlan +from dstack._internal.server.services.endpoints.planning import ( + EndpointPresetPlan, + EndpointPresetPlanningResult, +) +from dstack._internal.server.services.endpoints.presets import ( + EndpointPreset, + EndpointPresetRecipe, + EndpointPresetValidation, + EndpointPresetValidationReplica, +) +from dstack._internal.server.services.projects import add_project_member +from dstack._internal.server.testing.common import ( + create_fleet, + create_project, + create_repo, + create_run, + create_user, + get_auth_headers, + list_events, +) + + +class TestEndpointPlan: + @pytest.mark.asyncio + async def test_returns_40x_if_not_authenticated(self, client: AsyncClient): + response = await client.post("/api/project/main/endpoints/get_plan") + assert response.status_code in [401, 403] + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_returns_none_provisioning_plan( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.USER + ) + await create_fleet(session=session, project=project) + + with patch( + "dstack._internal.server.services.endpoints.find_preset_planning_result", + new=AsyncMock(return_value=EndpointPresetPlanningResult()), + ): + response = await client.post( + f"/api/project/{project.name}/endpoints/get_plan", + headers=get_auth_headers(user.token), + json={ + "configuration": { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + "env": {"HF_TOKEN": "secret"}, + }, + "configuration_path": "endpoint.dstack.yml", + }, + ) + + assert response.status_code == 200, response.json() + body = response.json() + assert body["project_name"] == project.name + assert body["user"] == user.name + assert body["configuration"]["model"] == { + "repo": "Qwen/Qwen3-0.6B", + "name": None, + } + assert body["configuration_path"] == "endpoint.dstack.yml" + assert body["current_resource"] is None + assert body["action"] == "create" + assert body["preset_policy"] == "reuse-or-create" + assert body["provisioning_plan"] == { + "type": "none", + "reason": ( + "No matching endpoint presets found. " + "Creating a preset requires the server agent, but " + "DSTACK_AGENT_ANTHROPIC_API_KEY is not set." + ), + } + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_returns_create_plan_when_existing_endpoint_is_terminal( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.USER + ) + await _create_endpoint_model( + session=session, + project=project, + user=user, + status=EndpointStatus.FAILED, + ) + + with patch( + "dstack._internal.server.services.endpoints.find_preset_planning_result", + new=AsyncMock(return_value=EndpointPresetPlanningResult()), + ): + response = await client.post( + f"/api/project/{project.name}/endpoints/get_plan", + headers=get_auth_headers(user.token), + json={ + "configuration": { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + }, + }, + ) + + assert response.status_code == 200, response.json() + body = response.json() + assert body["current_resource"] is None + assert body["action"] == "create" + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_returns_agent_provisioning_plan( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.ADMIN + ) + await create_fleet(session=session, project=project) + agent_service = Mock() + agent_service.is_enabled.return_value = True + agent_service.get_plan.return_value = AgentPlan(model="test-agent") + + with ( + patch( + "dstack._internal.server.services.endpoints.find_preset_planning_result", + new=AsyncMock(return_value=EndpointPresetPlanningResult()), + ), + patch( + "dstack._internal.server.services.endpoints.get_agent_service", + return_value=agent_service, + ), + ): + response = await client.post( + f"/api/project/{project.name}/endpoints/get_plan", + headers=get_auth_headers(user.token), + json={ + "configuration": { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + }, + "configuration_path": "endpoint.dstack.yml", + }, + ) + + assert response.status_code == 200, response.json() + body = response.json() + assert body["preset_policy"] == "reuse-or-create" + assert body["provisioning_plan"] == { + "type": "agent", + "agent_model": "test-agent", + "reason": None, + } + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_returns_no_agent_plan_without_existing_fleets( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.ADMIN + ) + agent_service = Mock() + agent_service.is_enabled.return_value = True + agent_service.get_plan.return_value = AgentPlan(model="test-agent") + + with ( + patch( + "dstack._internal.server.services.endpoints.find_preset_planning_result", + new=AsyncMock(return_value=EndpointPresetPlanningResult()), + ), + patch( + "dstack._internal.server.services.endpoints.get_agent_service", + return_value=agent_service, + ), + ): + response = await client.post( + f"/api/project/{project.name}/endpoints/get_plan", + headers=get_auth_headers(user.token), + json={ + "configuration": { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + }, + "configuration_path": "endpoint.dstack.yml", + }, + ) + + assert response.status_code == 200, response.json() + assert response.json()["provisioning_plan"] == { + "type": "none", + "reason": "The project has no fleets. Create one before submitting an endpoint.", + } + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_returns_no_fleets_plan_for_reuse_without_existing_fleets( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.ADMIN + ) + + response = await client.post( + f"/api/project/{project.name}/endpoints/get_plan", + headers=get_auth_headers(user.token), + json={ + "configuration": { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + "preset_policy": "reuse", + }, + "configuration_path": "endpoint.dstack.yml", + }, + ) + + assert response.status_code == 200, response.json() + assert response.json()["provisioning_plan"] == { + "type": "none", + "reason": "The project has no fleets. Create one before submitting an endpoint.", + } + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_returns_no_agent_plan_when_configured_fleet_does_not_match( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.ADMIN + ) + await create_fleet(session=session, project=project, name="existing-fleet") + agent_service = Mock() + agent_service.is_enabled.return_value = True + agent_service.get_plan.return_value = AgentPlan(model="test-agent") + + with ( + patch( + "dstack._internal.server.services.endpoints.find_preset_planning_result", + new=AsyncMock(return_value=EndpointPresetPlanningResult()), + ), + patch( + "dstack._internal.server.services.endpoints.get_agent_service", + return_value=agent_service, + ), + ): + response = await client.post( + f"/api/project/{project.name}/endpoints/get_plan", + headers=get_auth_headers(user.token), + json={ + "configuration": { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + "fleets": ["missing-fleet"], + }, + "configuration_path": "endpoint.dstack.yml", + }, + ) + + assert response.status_code == 200, response.json() + assert response.json()["provisioning_plan"] == { + "type": "none", + "reason": ( + "No fleets match the endpoint configuration. Create a fleet or update `fleets` " + "before submitting an endpoint." + ), + } + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_returns_no_agent_plan_for_project_user( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.USER + ) + await create_fleet(session=session, project=project) + agent_service = Mock() + agent_service.is_enabled.return_value = True + + with ( + patch( + "dstack._internal.server.services.endpoints.find_preset_planning_result", + new=AsyncMock(return_value=EndpointPresetPlanningResult()), + ), + patch( + "dstack._internal.server.services.endpoints.get_agent_service", + return_value=agent_service, + ), + ): + response = await client.post( + f"/api/project/{project.name}/endpoints/get_plan", + headers=get_auth_headers(user.token), + json={ + "configuration": { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + }, + }, + ) + + assert response.status_code == 200, response.json() + assert response.json()["provisioning_plan"] == { + "type": "none", + "reason": ( + "No matching endpoint presets found. " + "Creating endpoint presets with the server agent requires project admin " + "permissions." + ), + } + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_returns_preset_provisioning_plan( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.USER + ) + await create_fleet(session=session, project=project) + + with patch( + "dstack._internal.server.services.endpoints.find_preset_planning_result", + new=AsyncMock( + return_value=EndpointPresetPlanningResult(provisionable=_endpoint_preset_plan()) + ), + ): + response = await client.post( + f"/api/project/{project.name}/endpoints/get_plan", + headers=get_auth_headers(user.token), + json={ + "configuration": { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + }, + }, + ) + + assert response.status_code == 200, response.json() + body = response.json() + assert body["provisioning_plan"]["type"] == "preset" + assert body["preset_policy"] == "reuse-or-create" + assert body["provisioning_plan"]["preset_base"] == "Qwen/Qwen3-0.6B" + assert body["provisioning_plan"]["recipe_id"] == "vllm-t4" + assert body["provisioning_plan"]["service_name"] == "qwen-endpoint-serving" + assert body["provisioning_plan"]["job_offers"][0]["replica_group"] == "0" + assert body["provisioning_plan"]["job_offers"][0]["resources"]["gpu"] is not None + assert body["provisioning_plan"]["job_offers"][0]["spot"] is None + assert body["provisioning_plan"]["job_offers"][0]["max_price"] is None + assert body["provisioning_plan"]["job_offers"][0]["offers"][0]["backend"] == "aws" + assert body["provisioning_plan"]["job_offers"][0]["total_offers"] == 2 + + +class TestCreateEndpoint: + @pytest.mark.asyncio + async def test_returns_40x_if_not_authenticated(self, client: AsyncClient): + response = await client.post("/api/project/main/endpoints/create") + assert response.status_code in [401, 403] + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc)) + async def test_creates_endpoint(self, test_db, session: AsyncSession, client: AsyncClient): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.ADMIN + ) + await create_fleet(session=session, project=project) + + response = await client.post( + f"/api/project/{project.name}/endpoints/create", + headers=get_auth_headers(user.token), + json={ + "configuration": { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + "env": {"HF_TOKEN": "secret"}, + } + }, + ) + + assert response.status_code == 200, response.json() + body = response.json() + UUID(body["id"]) + assert body["name"] == "qwen-endpoint" + assert body["project_name"] == project.name + assert body["user"] == user.name + assert body["configuration"]["model"] == { + "repo": "Qwen/Qwen3-0.6B", + "name": None, + } + assert body["configuration"]["env"] == {"HF_TOKEN": "secret"} + assert body["created_at"] == "2023-01-02T03:04:00+00:00" + assert body["last_processed_at"] == "2023-01-02T03:04:00+00:00" + assert body["status"] == "submitted" + assert body["status_message"] is None + assert body["run_name"] is None + assert body["url"] is None + assert body["error"] is None + + res = await session.execute(select(EndpointModel)) + assert res.scalar_one() + events = await list_events(session) + assert len(events) == 1 + assert events[0].message == "Endpoint created. Status: SUBMITTED" + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_rejects_create_without_existing_fleets( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.ADMIN + ) + + with patch( + "dstack._internal.server.services.endpoints.find_preset_planning_result", + new=AsyncMock(return_value=EndpointPresetPlanningResult()), + ): + response = await client.post( + f"/api/project/{project.name}/endpoints/create", + headers=get_auth_headers(user.token), + json={ + "configuration": { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + } + }, + ) + + assert response.status_code == 400 + assert response.json()["detail"][0]["msg"] == ( + "The project has no fleets. Create one before submitting an endpoint." + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_rejects_reuse_without_existing_fleets( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.ADMIN + ) + + response = await client.post( + f"/api/project/{project.name}/endpoints/create", + headers=get_auth_headers(user.token), + json={ + "configuration": { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + "preset_policy": "reuse", + } + }, + ) + + assert response.status_code == 400 + assert response.json()["detail"][0]["msg"] == ( + "The project has no fleets. Create one before submitting an endpoint." + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_resubmits_terminal_endpoint( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.ADMIN + ) + await create_fleet(session=session, project=project) + previous_endpoint = await _create_endpoint_model( + session=session, + project=project, + user=user, + status=EndpointStatus.FAILED, + ) + previous_endpoint.status_message = "old failure" + await session.commit() + + response = await client.post( + f"/api/project/{project.name}/endpoints/create", + headers=get_auth_headers(user.token), + json={ + "configuration": { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + } + }, + ) + + assert response.status_code == 200, response.json() + endpoint_id = UUID(response.json()["id"]) + assert endpoint_id == previous_endpoint.id + await session.refresh(previous_endpoint) + assert previous_endpoint.status == EndpointStatus.SUBMITTED + assert previous_endpoint.status_message is None + assert previous_endpoint.service_run_id is None + assert previous_endpoint.configuration != "" + res = await session.execute(select(EndpointModel)) + assert len(res.scalars().all()) == 1 + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_rejects_resubmit_terminal_endpoint_without_existing_fleets( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.ADMIN + ) + previous_endpoint = await _create_endpoint_model( + session=session, + project=project, + user=user, + status=EndpointStatus.FAILED, + ) + previous_endpoint.status_message = "old failure" + await session.commit() + + response = await client.post( + f"/api/project/{project.name}/endpoints/create", + headers=get_auth_headers(user.token), + json={ + "configuration": { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + } + }, + ) + + assert response.status_code == 400 + assert response.json()["detail"][0]["msg"] == ( + "The project has no fleets. Create one before submitting an endpoint." + ) + await session.refresh(previous_endpoint) + assert previous_endpoint.status == EndpointStatus.FAILED + assert previous_endpoint.status_message == "old failure" + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_rejects_project_user_when_create_requires_agent( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.USER + ) + await create_fleet(session=session, project=project) + + with patch( + "dstack._internal.server.services.endpoints.find_preset_planning_result", + new=AsyncMock(return_value=EndpointPresetPlanningResult()), + ): + response = await client.post( + f"/api/project/{project.name}/endpoints/create", + headers=get_auth_headers(user.token), + json={ + "configuration": { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + } + }, + ) + + assert response.status_code == 403 + assert response.json()["msg"] == ( + "Creating endpoint presets with the server agent requires project admin permissions." + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_allows_project_user_when_matching_preset_exists( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.USER + ) + await create_fleet(session=session, project=project) + + with patch( + "dstack._internal.server.services.endpoints.find_preset_planning_result", + new=AsyncMock( + return_value=EndpointPresetPlanningResult(provisionable=_endpoint_preset_plan()) + ), + ): + response = await client.post( + f"/api/project/{project.name}/endpoints/create", + headers=get_auth_headers(user.token), + json={ + "configuration": { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + } + }, + ) + + assert response.status_code == 200, response.json() + assert response.json()["status"] == "submitted" + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_rejects_duplicate_non_terminal_endpoint( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.ADMIN + ) + await _create_endpoint_model( + session=session, + project=project, + user=user, + status=EndpointStatus.SUBMITTED, + ) + + response = await client.post( + f"/api/project/{project.name}/endpoints/create", + headers=get_auth_headers(user.token), + json={ + "configuration": { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + } + }, + ) + + assert response.status_code == 400 + assert response.json()["detail"][0]["code"] == "resource_exists" + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_rejects_unresolved_env_sentinel( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.USER + ) + + response = await client.post( + f"/api/project/{project.name}/endpoints/create", + headers=get_auth_headers(user.token), + json={ + "configuration": { + "type": "endpoint", + "name": "qwen-endpoint", + "model": "Qwen/Qwen3-0.6B", + "env": ["HF_TOKEN"], + } + }, + ) + + assert response.status_code == 400 + assert "Endpoint env is unresolved" in response.json()["detail"][0]["msg"] + + +class TestGetEndpoint: + @pytest.mark.asyncio + async def test_returns_40x_if_not_authenticated(self, client: AsyncClient): + response = await client.post("/api/project/main/endpoints/get") + assert response.status_code in [401, 403] + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_stopped_endpoint_hides_linked_run_name( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.USER + ) + endpoint_model = await _create_endpoint_model( + session=session, + project=project, + user=user, + status=EndpointStatus.STOPPED, + ) + repo = await create_repo( + session=session, + project_id=project.id, + repo_name="qwen-endpoint-serving-repo", + ) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_name="qwen-endpoint-serving", + status=RunStatus.TERMINATED, + ) + endpoint_model.service_run_id = run.id + await session.commit() + + response = await client.post( + f"/api/project/{project.name}/endpoints/get", + headers=get_auth_headers(user.token), + json={"name": endpoint_model.name}, + ) + + assert response.status_code == 200, response.json() + assert response.json()["status"] == "stopped" + assert response.json()["run_name"] is None + + +class TestStopEndpoint: + @pytest.mark.asyncio + async def test_returns_40x_if_not_authenticated(self, client: AsyncClient): + response = await client.post("/api/project/main/endpoints/stop") + assert response.status_code in [401, 403] + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_marks_endpoint_for_stopping( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.USER + ) + endpoint_model = await _create_endpoint_model( + session=session, + project=project, + user=user, + status=EndpointStatus.RUNNING, + ) + + response = await client.post( + f"/api/project/{project.name}/endpoints/stop", + headers=get_auth_headers(user.token), + json={"names": [endpoint_model.name]}, + ) + + assert response.status_code == 200, response.json() + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.STOPPING + events = await list_events(session) + assert [e.message for e in events] == ["Endpoint marked for stopping"] + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_does_not_stop_finished_endpoint( + self, test_db, session: AsyncSession, client: AsyncClient + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.USER + ) + endpoint_model = await _create_endpoint_model( + session=session, + project=project, + user=user, + status=EndpointStatus.FAILED, + ) + + response = await client.post( + f"/api/project/{project.name}/endpoints/stop", + headers=get_auth_headers(user.token), + json={"names": [endpoint_model.name]}, + ) + + assert response.status_code == 200, response.json() + await session.refresh(endpoint_model) + assert endpoint_model.status == EndpointStatus.FAILED + events = await list_events(session) + assert events == [] + + +class TestEndpointPresets: + @pytest.mark.asyncio + async def test_list_returns_40x_if_not_authenticated(self, client: AsyncClient): + response = await client.post("/api/project/main/endpoints/presets/list") + assert response.status_code in [401, 403] + + @pytest.mark.asyncio + async def test_delete_returns_40x_if_not_authenticated(self, client: AsyncClient): + response = await client.post("/api/project/main/endpoints/presets/delete") + assert response.status_code in [401, 403] + + @pytest.mark.asyncio + async def test_get_returns_40x_if_not_authenticated(self, client: AsyncClient): + response = await client.post("/api/project/main/endpoints/presets/get") + assert response.status_code in [401, 403] + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_lists_endpoint_presets( + self, + test_db, + session: AsyncSession, + client: AsyncClient, + monkeypatch: pytest.MonkeyPatch, + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.USER + ) + preset_service = _FakeEndpointPresetService([_endpoint_preset_plan().preset]) + monkeypatch.setattr( + "dstack._internal.server.routers.endpoints.get_endpoint_preset_service", + lambda: preset_service, + ) + + response = await client.post( + f"/api/project/{project.name}/endpoints/presets/list", + headers=get_auth_headers(user.token), + ) + + assert response.status_code == 200, response.json() + body = response.json() + assert len(body) == 1 + assert body[0]["base"] == "Qwen/Qwen3-0.6B" + assert body[0]["recipes"][0]["id"] == "vllm-t4" + assert body[0]["recipes"][0]["service"]["resources"]["gpu"] is not None + assert body[0]["recipes"][0]["validations"][0]["replicas"][0]["resources"][0]["gpu"] + assert preset_service.list_project_names == [project.name] + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_deletes_endpoint_preset( + self, + test_db, + session: AsyncSession, + client: AsyncClient, + monkeypatch: pytest.MonkeyPatch, + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.USER + ) + preset_service = _FakeEndpointPresetService([_endpoint_preset_plan().preset]) + monkeypatch.setattr( + "dstack._internal.server.routers.endpoints.get_endpoint_preset_service", + lambda: preset_service, + ) + + response = await client.post( + f"/api/project/{project.name}/endpoints/presets/delete", + headers=get_auth_headers(user.token), + json={"models": ["Qwen/Qwen3-0.6B"]}, + ) + + assert response.status_code == 200, response.json() + assert preset_service.deleted_models == ["Qwen/Qwen3-0.6B"] + assert preset_service.delete_project_names == [project.name] + + @pytest.mark.asyncio + @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True) + async def test_gets_endpoint_preset( + self, + test_db, + session: AsyncSession, + client: AsyncClient, + monkeypatch: pytest.MonkeyPatch, + ): + user = await create_user(session, global_role=GlobalRole.USER) + project = await create_project(session) + await add_project_member( + session=session, project=project, user=user, project_role=ProjectRole.USER + ) + preset_service = _FakeEndpointPresetService([_endpoint_preset_plan().preset]) + monkeypatch.setattr( + "dstack._internal.server.routers.endpoints.get_endpoint_preset_service", + lambda: preset_service, + ) + + response = await client.post( + f"/api/project/{project.name}/endpoints/presets/get", + headers=get_auth_headers(user.token), + json={"model": "Qwen/Qwen3-0.6B"}, + ) + + assert response.status_code == 200, response.json() + body = response.json() + assert body["base"] == "Qwen/Qwen3-0.6B" + assert body["recipes"][0]["service"]["type"] == "service" + assert body["recipes"][0]["service"]["resources"]["gpu"] is not None + assert preset_service.get_project_names == [project.name] + + +class _FakeEndpointPresetService: + def __init__(self, presets): + self._presets = presets + self.list_project_names = [] + self.get_project_names = [] + self.delete_project_names = [] + self.deleted_models = [] + + async def list_presets(self, project_name): + self.list_project_names.append(project_name) + return self._presets + + async def get_preset(self, project_name, model): + self.get_project_names.append(project_name) + for preset in self._presets: + if preset.base == model: + return preset + return None + + async def delete_preset(self, project_name, model): + if model not in {preset.base for preset in self._presets}: + raise FileNotFoundError(model) + self.delete_project_names.append(project_name) + self.deleted_models.append(model) + + +async def _create_endpoint_model( + session: AsyncSession, + project, + user, + status: EndpointStatus = EndpointStatus.SUBMITTED, + name: str = "qwen-endpoint", +) -> EndpointModel: + configuration = EndpointConfiguration(name=name, model="Qwen/Qwen3-0.6B") + endpoint_model = EndpointModel( + name=name, + project=project, + user=user, + configuration=configuration.json(), + status=status, + created_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), + last_processed_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), + ) + session.add(endpoint_model) + await session.commit() + return endpoint_model + + +def _endpoint_preset_plan() -> EndpointPresetPlan: + service_configuration = ServiceConfiguration.parse_obj( + { + "type": "service", + "name": "qwen-endpoint-serving", + "commands": [ + "vllm serve Qwen/Qwen3-0.6B --host 0.0.0.0 --port 8000", + ], + "port": 8000, + "model": "Qwen/Qwen3-0.6B", + "resources": {"gpu": "16GB"}, + } + ) + recipe = EndpointPresetRecipe( + id="vllm-t4", + model="Qwen/Qwen3-0.6B", + service=service_configuration, + validations=[ + EndpointPresetValidation( + replicas=[ + EndpointPresetValidationReplica.parse_obj( + { + "resources": [ + { + "cpu": 4, + "memory": "16GB", + "disk": "100GB", + "gpu": {"name": "T4", "memory": "16GB", "count": 1}, + } + ] + } + ) + ] + ) + ], + ) + preset = EndpointPreset(base="Qwen/Qwen3-0.6B", recipes=[recipe]) + run_plan = Mock() + run_plan.get_effective_run_spec.return_value = RunSpec( + run_name="qwen-endpoint-serving", + configuration=service_configuration, + ) + job_plan = Mock() + job_plan.job_spec.replica_group = "0" + job_plan.job_spec.requirements.resources = service_configuration.resources + job_plan.job_spec.requirements.spot = None + job_plan.job_spec.requirements.max_price = None + job_plan.offers = [_instance_offer()] + job_plan.total_offers = 2 + job_plan.max_price = 1.25 + run_plan.job_plans = [job_plan] + return EndpointPresetPlan(preset=preset, recipe=recipe, run_plan=run_plan) + + +def _instance_offer() -> InstanceOfferWithAvailability: + return InstanceOfferWithAvailability.parse_obj( + { + "backend": BackendType.AWS, + "region": "us-east-1", + "price": 1.25, + "availability": InstanceAvailability.AVAILABLE, + "instance": { + "name": "g5.xlarge", + "resources": { + "cpus": 4, + "memory_mib": 16384, + "gpus": [], + "spot": False, + }, + }, + } + ) diff --git a/src/tests/_internal/server/services/endpoints/test_agent_report.py b/src/tests/_internal/server/services/endpoints/test_agent_report.py new file mode 100644 index 0000000000..a144ed387b --- /dev/null +++ b/src/tests/_internal/server/services/endpoints/test_agent_report.py @@ -0,0 +1,79 @@ +import uuid + +import pytest +from pydantic import ValidationError + +from dstack._internal.server.services.endpoints.agent.report import AgentFinalReport + + +class TestAgentFinalReport: + def test_accepts_success_report(self): + run_id = uuid.uuid4() + + report = AgentFinalReport.parse_obj( + { + "success": True, + "run_id": str(run_id), + "run_name": "qwen-serving", + "service_yaml": "type: service\nname: qwen-serving\n", + "base": "Qwen/Qwen3-0.6B", + "model": "Qwen/Qwen3-0.6B", + } + ) + + assert report.run_id == run_id + assert report.run_name == "qwen-serving" + assert report.base == "Qwen/Qwen3-0.6B" + assert report.model == "Qwen/Qwen3-0.6B" + + def test_rejects_success_without_run_id(self): + with pytest.raises(ValidationError, match="run_id"): + AgentFinalReport.parse_obj( + { + "success": True, + "run_name": "qwen-serving", + "service_yaml": "type: service\nname: qwen-serving\n", + "base": "Qwen/Qwen3-0.6B", + "model": "Qwen/Qwen3-0.6B", + } + ) + + def test_rejects_success_with_empty_base(self): + with pytest.raises(ValidationError, match="base"): + AgentFinalReport.parse_obj( + { + "success": True, + "run_id": str(uuid.uuid4()), + "run_name": "qwen-serving", + "service_yaml": "type: service\nname: qwen-serving\n", + "base": "", + "model": "Qwen/Qwen3-0.6B", + } + ) + + def test_rejects_success_with_empty_model(self): + with pytest.raises(ValidationError, match="model"): + AgentFinalReport.parse_obj( + { + "success": True, + "run_id": str(uuid.uuid4()), + "run_name": "qwen-serving", + "service_yaml": "type: service\nname: qwen-serving\n", + "base": "Qwen/Qwen3-0.6B", + "model": "", + } + ) + + def test_accepts_failure_report(self): + report = AgentFinalReport.parse_obj( + { + "success": False, + "failure_summary": "No deployable recipe matched the budget.", + } + ) + + assert report.failure_summary == "No deployable recipe matched the budget." + + def test_rejects_failure_without_summary(self): + with pytest.raises(ValidationError, match="failure_summary"): + AgentFinalReport.parse_obj({"success": False}) diff --git a/src/tests/_internal/server/services/endpoints/test_claude_agent.py b/src/tests/_internal/server/services/endpoints/test_claude_agent.py new file mode 100644 index 0000000000..90a8522c6f --- /dev/null +++ b/src/tests/_internal/server/services/endpoints/test_claude_agent.py @@ -0,0 +1,1475 @@ +import json +import signal +import stat +import subprocess +import uuid +from asyncio import StreamReader +from datetime import datetime, timezone +from pathlib import Path + +import pytest +import yaml +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +import dstack._internal.server.services.endpoints.agent.claude as claude_module +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.endpoints import EndpointConfiguration, EndpointStatus +from dstack._internal.core.models.envs import Env +from dstack._internal.core.models.profiles import SpotPolicy +from dstack._internal.server import settings +from dstack._internal.server.models import ( + EndpointAgentSessionModel, + EndpointAgentSessionStatus, + EndpointModel, +) +from dstack._internal.server.services.endpoints.agent.claude import ( + ClaudeAgentService, + _AgentRunnerResult, + _AgentWorkspace, + _build_claude_command, + _load_submissions, + _read_agent_stdout, + _run_agent_in_subprocess, + _write_trace_record, + get_claude_agent_unavailable_reason, +) +from dstack._internal.server.services.endpoints.agent.report import AgentFinalReport +from dstack._internal.server.testing.common import ( + create_fleet, + create_project, + create_user, +) + + +async def _create_endpoint_model( + session: AsyncSession, + model="Qwen/Qwen3-0.6B", + max_price: float | None = None, + spot_policy: SpotPolicy | None = None, + backends: list[BackendType] | None = None, + fleets: list[str] | None = None, + create_default_fleet: bool = True, +) -> EndpointModel: + user = await create_user(session=session, name="admin", token="user-token") + project = await create_project(session=session, owner=user, name="main") + if create_default_fleet: + await create_fleet(session=session, project=project) + configuration = EndpointConfiguration( + name="qwen-endpoint", + model=model, + env=Env.parse_obj({"HF_TOKEN": "hf-secret"}), + max_price=max_price, + spot_policy=spot_policy, + backends=backends, + fleets=fleets, + ) + endpoint_model = EndpointModel( + id=uuid.uuid4(), + name=configuration.name, + project=project, + user=user, + status=EndpointStatus.PROVISIONING, + provisioning_method="agent", + configuration=configuration.json(), + created_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), + last_processed_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), + ) + session.add(endpoint_model) + await session.commit() + return endpoint_model + + +def _successful_agent_report( + *, + run_id: uuid.UUID | None = None, + run_name: str = "qwen-endpoint-1", + base: str = "Qwen/Qwen3-0.6B", + model: str = "Qwen/Qwen3-0.6B", + service_yaml: str | None = None, +) -> AgentFinalReport: + if run_id is None: + run_id = uuid.uuid4() + if service_yaml is None: + service_yaml = f"type: service\nname: {run_name}\n" + return AgentFinalReport( + success=True, + run_id=run_id, + run_name=run_name, + service_yaml=service_yaml, + base=base, + model=model, + ) + + +def _configure_fake_claude(tmp_path, monkeypatch: pytest.MonkeyPatch) -> str: + claude_path = tmp_path / "claude" + claude_path.write_text("#!/bin/sh\nexit 0\n", encoding="utf-8") + claude_path.chmod(0o755) + monkeypatch.setattr(settings, "AGENT_CLAUDE_PATH", str(claude_path)) + return str(claude_path) + + +@pytest.fixture(autouse=True) +def _clear_claude_effort(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(settings, "AGENT_CLAUDE_EFFORT", None) + + +class TestClaudeAgentService: + @pytest.mark.asyncio + async def test_default_runner_starts_agent_detached( + self, + session: AsyncSession, + tmp_path, + monkeypatch, + ): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "agent-secret") + _configure_fake_claude(tmp_path, monkeypatch) + monkeypatch.setattr(settings, "SERVER_URL", "http://127.0.0.1:8000") + started = {} + written_logs = [] + + def write_logs(**kwargs): + written_logs.extend(log.message.decode().rstrip("\n") for log in kwargs["job_logs"]) + + def start_agent(workspace, request): + started["workspace"] = workspace + started["request"] = request + return 12345 + + monkeypatch.setattr(claude_module.logs_services, "write_logs", write_logs) + monkeypatch.setattr(claude_module, "_start_agent_subprocess_detached", start_agent) + endpoint_model = await _create_endpoint_model(session) + service = ClaudeAgentService(workspace_base_dir=tmp_path) + + result = await service.provision_endpoint(endpoint_model, pipeline_hinter=None) + + assert result.in_progress is True + assert result.error is None + session_root = tmp_path / str(endpoint_model.id) / "1" + assert started["workspace"].root_dir == session_root + assert started["request"]["cwd"] == str(session_root / "workspace") + assert "max_budget" not in started["request"]["options"] + res = await session.execute( + select(EndpointAgentSessionModel).where( + EndpointAgentSessionModel.endpoint_id == endpoint_model.id + ) + ) + agent_session = res.scalar_one() + assert agent_session.session_num == 1 + assert agent_session.status == EndpointAgentSessionStatus.RUNNING + assert agent_session.pid == 12345 + assert agent_session.workspace_path == str(session_root) + progress_path = session_root / "workspace" / "progress.jsonl" + progress_text = progress_path.read_text() + assert "Starting endpoint prototyping agent for Qwen/Qwen3-0.6B" in progress_text + assert agent_session.progress_log_offset == progress_path.stat().st_size + assert written_logs == [ + "Starting endpoint prototyping agent for Qwen/Qwen3-0.6B. " + "Allowed fleets: test-fleet. The agent will inspect offers, choose a service " + "recipe, deploy it, and verify the model API before the endpoint becomes active." + ] + + @pytest.mark.asyncio + async def test_restart_resumes_same_session_when_previous_process_exited( + self, + session: AsyncSession, + tmp_path, + monkeypatch, + ): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "agent-secret") + _configure_fake_claude(tmp_path, monkeypatch) + monkeypatch.setattr(settings, "SERVER_URL", "http://127.0.0.1:8000") + endpoint_model = await _create_endpoint_model(session) + session_root = tmp_path / str(endpoint_model.id) / "1" + work_dir = session_root / "workspace" + work_dir.mkdir(parents=True) + (work_dir / "submissions.jsonl").write_text( + json.dumps( + { + "name": "qwen-endpoint-1", + "type": "service", + "status": "submitted", + "run_id": str(uuid.uuid4()), + } + ) + + "\n", + encoding="utf-8", + ) + agent_session = EndpointAgentSessionModel( + endpoint_id=endpoint_model.id, + session_num=1, + status=EndpointAgentSessionStatus.RUNNING, + workspace_path=str(session_root), + pid=12345, + process_host=claude_module._get_process_host(), + progress_log_offset=0, + stdout_log_offset=0, + stderr_log_offset=0, + created_at=endpoint_model.created_at, + updated_at=endpoint_model.created_at, + ) + session.add(agent_session) + await session.commit() + monkeypatch.setattr(claude_module, "_is_process_group_running", lambda pgid: False) + started = {} + + def start_agent(workspace, request): + started["workspace"] = workspace + started["request"] = request + return 23456 + + monkeypatch.setattr(claude_module, "_start_agent_subprocess_detached", start_agent) + service = ClaudeAgentService(workspace_base_dir=tmp_path) + + result = await service.provision_endpoint(endpoint_model, pipeline_hinter=None) + + assert result.in_progress is True + assert result.error is None + assert started["workspace"].root_dir == session_root + assert "On startup or resume" in started["request"]["prompt"] + assert "previous_sessions.md" not in started["request"]["prompt"] + await session.refresh(agent_session) + assert agent_session.session_num == 1 + assert agent_session.status == EndpointAgentSessionStatus.RUNNING + assert agent_session.pid == 23456 + assert agent_session.workspace_path == str(session_root) + + @pytest.mark.asyncio + async def test_new_endpoint_lifecycle_starts_new_session_without_previous_context( + self, + session: AsyncSession, + tmp_path, + monkeypatch, + ): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "agent-secret") + _configure_fake_claude(tmp_path, monkeypatch) + endpoint_model = await _create_endpoint_model(session) + previous_root = tmp_path / str(endpoint_model.id) / "1" + previous_work_dir = previous_root / "workspace" + previous_work_dir.mkdir(parents=True) + previous_run_id = uuid.uuid4() + (previous_work_dir / "final_report.json").write_text( + json.dumps( + { + "success": False, + "run_id": str(previous_run_id), + "run_name": "qwen-endpoint-1", + "service_yaml": "type: service\nname: qwen-endpoint-1\n", + "failure_summary": "Previous image required a newer driver.", + } + ) + + "\n", + encoding="utf-8", + ) + (previous_work_dir / "submissions.jsonl").write_text( + json.dumps( + { + "name": "qwen-endpoint-1", + "status": "failed", + "run_id": str(previous_run_id), + } + ) + + "\n", + encoding="utf-8", + ) + previous_session = EndpointAgentSessionModel( + endpoint_id=endpoint_model.id, + session_num=1, + status=EndpointAgentSessionStatus.FAILED, + workspace_path=str(previous_root), + progress_log_offset=0, + stdout_log_offset=0, + stderr_log_offset=0, + status_message="Previous image required a newer driver.", + created_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), + updated_at=datetime(2023, 1, 2, 3, 5, tzinfo=timezone.utc), + finished_at=datetime(2023, 1, 2, 3, 6, tzinfo=timezone.utc), + ) + session.add(previous_session) + endpoint_model.created_at = datetime(2023, 1, 3, 3, 4, tzinfo=timezone.utc) + await session.commit() + captured = {} + + async def runner(workspace, request): + captured["workspace"] = workspace + captured["request"] = request + return _AgentRunnerResult(report=_successful_agent_report()) + + service = ClaudeAgentService(runner=runner, workspace_base_dir=tmp_path) + + result = await service.provision_endpoint(endpoint_model, pipeline_hinter=None) + + assert result.error is None + assert captured["workspace"].root_dir == tmp_path / str(endpoint_model.id) / "2" + prompt = captured["request"]["prompt"] + assert "previous_sessions.md" not in prompt + assert "Previous image required a newer driver." not in prompt + assert not ( + tmp_path / str(endpoint_model.id) / "2" / "workspace" / "previous_sessions.md" + ).exists() + res = await session.execute( + select(EndpointAgentSessionModel) + .where(EndpointAgentSessionModel.endpoint_id == endpoint_model.id) + .order_by(EndpointAgentSessionModel.session_num) + ) + sessions = list(res.scalars().all()) + assert [agent_session.session_num for agent_session in sessions] == [1, 2] + + @pytest.mark.asyncio + async def test_abort_endpoint_stops_agent_process_group( + self, + session: AsyncSession, + tmp_path, + monkeypatch, + ): + monkeypatch.setattr(settings, "SERVER_URL", "http://127.0.0.1:8000") + endpoint_model = await _create_endpoint_model(session) + session_root = tmp_path / str(endpoint_model.id) / "1" + work_dir = session_root / "workspace" + work_dir.mkdir(parents=True) + process_state_path = work_dir / "agent_process.json" + process_state_path.write_text( + json.dumps( + { + "pid": 12345, + "pgid": 12345, + "host": claude_module._get_process_host(), + } + ) + + "\n", + encoding="utf-8", + ) + agent_session = EndpointAgentSessionModel( + endpoint_id=endpoint_model.id, + session_num=1, + status=EndpointAgentSessionStatus.RUNNING, + workspace_path=str(session_root), + pid=12345, + process_host=claude_module._get_process_host(), + progress_log_offset=0, + stdout_log_offset=0, + stderr_log_offset=0, + created_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), + updated_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), + ) + session.add(agent_session) + await session.commit() + + group_running = iter([True, False]) + signals = [] + + monkeypatch.setattr( + claude_module, + "_is_process_group_running", + lambda pgid: next(group_running), + ) + monkeypatch.setattr(claude_module, "_AGENT_PROCESS_ABORT_GRACE_SECONDS", 0) + monkeypatch.setattr( + claude_module.os, + "killpg", + lambda pgid, sig: signals.append((pgid, sig)), + ) + + service = ClaudeAgentService(workspace_base_dir=tmp_path) + + aborted = await service.abort_endpoint(endpoint_model) + + assert aborted is True + assert signals == [(12345, signal.SIGTERM)] + assert not process_state_path.exists() + await session.refresh(agent_session) + assert agent_session.status == EndpointAgentSessionStatus.FAILED + assert agent_session.status_message == "Endpoint stop requested" + + @pytest.mark.asyncio + async def test_abort_endpoint_waits_for_agent_process_on_another_host( + self, + session: AsyncSession, + tmp_path, + monkeypatch, + ): + monkeypatch.setattr(settings, "SERVER_URL", "http://127.0.0.1:8000") + endpoint_model = await _create_endpoint_model(session) + session_root = tmp_path / str(endpoint_model.id) / "1" + work_dir = session_root / "workspace" + work_dir.mkdir(parents=True) + process_state_path = work_dir / "agent_process.json" + process_state_path.write_text( + json.dumps({"pid": 12345, "pgid": 12345, "host": "other-host"}) + "\n", + encoding="utf-8", + ) + agent_session = EndpointAgentSessionModel( + endpoint_id=endpoint_model.id, + session_num=1, + status=EndpointAgentSessionStatus.RUNNING, + workspace_path=str(session_root), + pid=12345, + process_host="other-host", + progress_log_offset=0, + stdout_log_offset=0, + stderr_log_offset=0, + created_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), + updated_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), + ) + session.add(agent_session) + await session.commit() + monkeypatch.setattr( + claude_module.os, + "killpg", + lambda pgid, sig: pytest.fail("remote process group must not be signaled locally"), + ) + + service = ClaudeAgentService(workspace_base_dir=tmp_path) + + aborted = await service.abort_endpoint(endpoint_model) + + assert aborted is False + assert process_state_path.exists() + await session.refresh(agent_session) + assert agent_session.status == EndpointAgentSessionStatus.RUNNING + + def test_process_group_liveness_ignores_zombie_only_group(self, monkeypatch): + class _PsResult: + returncode = 0 + stdout = "64979 Z\n64979 Z+\n" + + def killpg(pgid, sig): + raise PermissionError + + monkeypatch.setattr(claude_module.os, "killpg", killpg) + monkeypatch.setattr( + claude_module.subprocess, + "run", + lambda *args, **kwargs: _PsResult(), + ) + + assert claude_module._is_process_group_running(64979) is False + + def test_process_group_liveness_detects_non_zombie_member(self, monkeypatch): + class _PsResult: + returncode = 0 + stdout = "64979 Z\n64979 S+\n" + + def killpg(pgid, sig): + raise PermissionError + + monkeypatch.setattr(claude_module.os, "killpg", killpg) + monkeypatch.setattr( + claude_module.subprocess, + "run", + lambda *args, **kwargs: _PsResult(), + ) + + assert claude_module._is_process_group_running(64979) is True + + @pytest.mark.asyncio + async def test_invokes_agent_with_isolated_dstack_cli_context( + self, + session: AsyncSession, + tmp_path, + monkeypatch, + ): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "agent-secret") + claude_path = _configure_fake_claude(tmp_path, monkeypatch) + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_MODEL", "test-claude-model") + monkeypatch.setattr(settings, "SERVER_URL", "http://127.0.0.1:8000") + monkeypatch.setenv("DATABASE_URL", "must-not-leak") + captured = {} + run_id = uuid.uuid4() + + async def runner(workspace, request): + captured["workspace"] = workspace + captured["request"] = request + return _AgentRunnerResult(report=_successful_agent_report(run_id=run_id)) + + endpoint_model = await _create_endpoint_model(session) + service = ClaudeAgentService(runner=runner, workspace_base_dir=tmp_path) + + result = await service.provision_endpoint(endpoint_model, pipeline_hinter=None) + + assert result.error is None + assert result.run_id == run_id + assert result.run_name == "qwen-endpoint-1" + assert result.final_report is not None + assert result.final_report.success is True + assert ( + "`final_report.json` must contain only the schema fields" + in captured["request"]["prompt"] + ) + session_root = tmp_path / str(endpoint_model.id) / "1" + assert captured["request"]["cwd"] == str(session_root / "workspace") + env = captured["request"]["env"] + assert env["ANTHROPIC_API_KEY"] == "agent-secret" + agent_home = Path(env["HOME"]) + assert agent_home != session_root / "home" + assert agent_home.resolve() == session_root / "home" + assert captured["workspace"].home_dir == session_root / "home" + assert captured["workspace"].dstack_home_dir == agent_home + control_sock_path = agent_home / ".dstack" / "ssh" / f"{'x' * 60}.control.sock" + assert len(str(control_sock_path)) < 104 + assert stat.S_IMODE((session_root / "home" / ".ssh").stat().st_mode) == 0o700 + assert not (session_root / "bin" / "dstack").exists() + assert env["DSTACK_PROJECT"] == "main" + assert env["DSTACK_ENDPOINT_SERVER_URL"] == "http://127.0.0.1:8000" + assert env["DSTACK_ENDPOINT_BEARER_TOKEN"] == "user-token" + assert env["HF_TOKEN"] == "hf-secret" + assert "DATABASE_URL" not in env + assert captured["request"]["options"]["model"] == "test-claude-model" + assert "max_budget" not in captured["request"]["options"] + assert "StructuredOutput" in captured["request"]["options"]["tools"] + assert "Edit" in captured["request"]["options"]["allowed_tools"] + assert "StructuredOutput" in captured["request"]["options"]["allowed_tools"] + schema_text = json.dumps(captured["request"]["options"]["json_schema"]) + assert '"title"' not in schema_text + assert '"format"' not in schema_text + command = _build_claude_command(captured["request"]) + assert command[0] == claude_path + assert "--bare" in command + assert "--setting-sources" not in command + assert "--effort" not in command + assert command[command.index("--tools") + 1] == captured["request"]["options"]["tools"] + assert "--permission-mode" in command + assert command[command.index("--permission-mode") + 1] == "bypassPermissions" + assert "--max-budget-usd" not in command + config = yaml.safe_load((session_root / "home" / ".dstack" / "config.yml").read_text()) + assert config == { + "projects": [ + { + "name": "main", + "url": "http://127.0.0.1:8000", + "token": "user-token", + "default": True, + } + ], + } + work_dir = session_root / "workspace" + state = json.loads((work_dir / "agent_state.json").read_text()) + assert state["endpoint_name"] == "qwen-endpoint" + assert state["model"] == "Qwen/Qwen3-0.6B" + assert state["phase"] == "success" + assert "max_agent_budget" not in state + assert not (work_dir / "sources.jsonl").exists() + assert (work_dir / "submissions.jsonl").exists() + assert (work_dir / "commands.jsonl").exists() + assert (work_dir / "progress.jsonl").exists() + progress_helper = session_root / "bin" / "progress" + assert progress_helper.exists() + assert not (session_root / "bin" / "dstack").exists() + assert not (session_root / "bin" / "ssh").exists() + progress_result = subprocess.run( + [str(progress_helper), "Checked model config."], + cwd=work_dir, + env=captured["request"]["env"], + capture_output=True, + text=True, + check=False, + ) + assert progress_result.returncode == 0 + assert json.loads((work_dir / "progress.jsonl").read_text().splitlines()[-1]) == { + "message": "Checked model config." + } + assert not (work_dir / "hardware_reasoning.md").exists() + assert (work_dir / ".claude" / "skills" / "dstack" / "SKILL.md").exists() + assert (work_dir / ".claude" / "skills" / "dstack-prototyping" / "SKILL.md").exists() + final_report = json.loads((work_dir / "final_report.json").read_text()) + assert final_report["success"] is True + assert final_report["run_name"] == "qwen-endpoint-1" + + @pytest.mark.asyncio + async def test_prompt_includes_base_model_endpoint_context( + self, + session: AsyncSession, + tmp_path, + monkeypatch, + ): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "agent-secret") + _configure_fake_claude(tmp_path, monkeypatch) + captured = {} + + async def runner(workspace, request): + captured["request"] = request + return _AgentRunnerResult( + report=AgentFinalReport( + success=False, + failure_summary="not needed", + ) + ) + + endpoint_model = await _create_endpoint_model( + session, + model={"base": "Qwen/Qwen3.6-27B"}, + ) + service = ClaudeAgentService(runner=runner, workspace_base_dir=tmp_path) + + result = await service.provision_endpoint(endpoint_model, pipeline_hinter=None) + + assert result.final_report is not None + prompt = captured["request"]["prompt"] + assert "Endpoint context:" in prompt + assert "- endpoint_name: qwen-endpoint" in prompt + assert "- service_model_name: Qwen/Qwen3.6-27B" in prompt + assert "- base_model: Qwen/Qwen3.6-27B" in prompt + assert "- model_repo:" not in prompt + assert "- final_report.json.model:" not in prompt + + @pytest.mark.asyncio + async def test_prompt_includes_exact_model_endpoint_context( + self, + session: AsyncSession, + tmp_path, + monkeypatch, + ): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "agent-secret") + _configure_fake_claude(tmp_path, monkeypatch) + captured = {} + + async def runner(workspace, request): + captured["request"] = request + return _AgentRunnerResult( + report=AgentFinalReport( + success=False, + failure_summary="not needed", + ) + ) + + endpoint_model = await _create_endpoint_model( + session, + model={ + "repo": "groxaxo/Qwen3.6-27B-GPTQ-Pro-4Bit", + "name": "Qwen/Qwen3.6-27B", + }, + ) + service = ClaudeAgentService(runner=runner, workspace_base_dir=tmp_path) + + result = await service.provision_endpoint(endpoint_model, pipeline_hinter=None) + + assert result.final_report is not None + prompt = captured["request"]["prompt"] + assert "Endpoint context:" in prompt + assert "- endpoint_name: qwen-endpoint" in prompt + assert "- service_model_name: Qwen/Qwen3.6-27B" in prompt + assert "- model_repo: groxaxo/Qwen3.6-27B-GPTQ-Pro-4Bit" in prompt + assert "- base_model:" not in prompt + assert "- final_report.json.model:" not in prompt + + @pytest.mark.asyncio + async def test_existing_claude_auth_uses_real_home_for_claude_and_short_home_for_dstack( + self, + session: AsyncSession, + tmp_path, + monkeypatch, + ): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", None) + monkeypatch.setattr(settings, "AGENT_CLAUDE_USE_EXISTING_AUTH", True) + monkeypatch.setattr(claude_module, "_existing_auth_warning_logged", False) + _configure_fake_claude(tmp_path, monkeypatch) + monkeypatch.setattr(settings, "SERVER_URL", "http://127.0.0.1:8000") + captured = {} + run_id = uuid.uuid4() + + async def runner(workspace, request): + captured["workspace"] = workspace + captured["request"] = request + return _AgentRunnerResult(report=_successful_agent_report(run_id=run_id)) + + endpoint_model = await _create_endpoint_model(session) + service = ClaudeAgentService(runner=runner, workspace_base_dir=tmp_path) + + result = await service.provision_endpoint(endpoint_model, pipeline_hinter=None) + + assert result.error is None + session_root = tmp_path / str(endpoint_model.id) / "1" + env = captured["request"]["env"] + workspace = captured["workspace"] + assert env["HOME"] == str(Path.home()) + assert "ANTHROPIC_API_KEY" not in env + assert workspace.home_dir == session_root / "home" + assert workspace.dstack_home_dir != workspace.home_dir + assert workspace.dstack_home_dir.resolve() == workspace.home_dir + wrapper = session_root / "bin" / "dstack" + assert wrapper.exists() + wrapper_text = wrapper.read_text() + assert f'os.environ["HOME"] = "{workspace.dstack_home_dir}"' in wrapper_text + assert str(workspace.home_dir) not in wrapper_text + ssh_wrapper = session_root / "bin" / "ssh" + assert ssh_wrapper.exists() + ssh_wrapper_text = ssh_wrapper.read_text() + assert f'os.environ["HOME"] = "{workspace.dstack_home_dir}"' in ssh_wrapper_text + assert str(workspace.home_dir) not in ssh_wrapper_text + command = _build_claude_command(captured["request"]) + assert "--bare" not in command + assert command[command.index("--setting-sources") + 1] == "project,local" + assert "--effort" not in command + + @pytest.mark.asyncio + async def test_includes_endpoint_profile_constraints_in_prompt( + self, + session: AsyncSession, + tmp_path, + monkeypatch, + ): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "agent-secret") + _configure_fake_claude(tmp_path, monkeypatch) + captured = {} + + async def runner(workspace, request): + captured["request"] = request + return _AgentRunnerResult(report=_successful_agent_report()) + + endpoint_model = await _create_endpoint_model( + session, + max_price=0.3, + spot_policy=SpotPolicy.ONDEMAND, + ) + service = ClaudeAgentService(runner=runner, workspace_base_dir=tmp_path) + + result = await service.provision_endpoint(endpoint_model, pipeline_hinter=None) + + assert result.error is None + prompt = captured["request"]["prompt"] + assert "- max_price: 0.3" in prompt + assert "- spot_policy: on-demand" in prompt + assert "Reuse these CLI flags" not in prompt + assert "--max-price 0.3 --on-demand" not in prompt + assert "--availability-zone" not in prompt + assert "--instance " not in prompt + + @pytest.mark.asyncio + async def test_reuses_existing_final_report_without_invoking_runner( + self, + session: AsyncSession, + tmp_path, + monkeypatch, + ): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "agent-secret") + _configure_fake_claude(tmp_path, monkeypatch) + endpoint_model = await _create_endpoint_model(session) + work_dir = tmp_path / str(endpoint_model.id) / "workspace" + work_dir.mkdir(parents=True) + run_id = uuid.uuid4() + (work_dir / "final_report.json").write_text( + json.dumps( + { + "success": True, + "run_id": str(run_id), + "run_name": "qwen-endpoint-1", + "service_yaml": "type: service\nname: qwen-endpoint-1\n", + "base": "Qwen/Qwen3-0.6B", + "model": "Qwen/Qwen3-0.6B", + } + ), + encoding="utf-8", + ) + + def start_agent(workspace, request): + raise AssertionError("agent process must not be started") + + monkeypatch.setattr(claude_module, "_start_agent_subprocess_detached", start_agent) + service = ClaudeAgentService(workspace_base_dir=tmp_path) + + result = await service.provision_endpoint(endpoint_model, pipeline_hinter=None) + + assert result.error is None + assert result.run_id == run_id + assert result.run_name == "qwen-endpoint-1" + assert result.final_report is not None + assert result.final_report.success is True + + @pytest.mark.asyncio + async def test_waits_for_claude_result_after_final_report_while_process_runs( + self, + session: AsyncSession, + tmp_path, + monkeypatch, + ): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "agent-secret") + _configure_fake_claude(tmp_path, monkeypatch) + monkeypatch.setattr( + claude_module, + "_get_running_agent_process_pid", + lambda workspace, agent_session=None: 123, + ) + endpoint_model = await _create_endpoint_model(session) + work_dir = tmp_path / str(endpoint_model.id) / "workspace" + work_dir.mkdir(parents=True) + (work_dir / "final_report.json").write_text( + json.dumps( + { + "success": True, + "run_id": str(uuid.uuid4()), + "run_name": "qwen-endpoint-1", + "service_yaml": "type: service\nname: qwen-endpoint-1\n", + "base": "Qwen/Qwen3-0.6B", + "model": "Qwen/Qwen3-0.6B", + } + ), + encoding="utf-8", + ) + + service = ClaudeAgentService(workspace_base_dir=tmp_path) + + result = await service.provision_endpoint(endpoint_model, pipeline_hinter=None) + + assert result.in_progress is True + assert result.final_report is None + + @pytest.mark.asyncio + async def test_returns_in_progress_when_agent_process_is_still_running( + self, + session: AsyncSession, + tmp_path, + monkeypatch, + ): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "agent-secret") + _configure_fake_claude(tmp_path, monkeypatch) + monkeypatch.setattr( + claude_module, + "_get_running_agent_process_pid", + lambda workspace, agent_session=None: 123, + ) + endpoint_model = await _create_endpoint_model(session) + + def start_agent(workspace, request): + raise AssertionError("agent process must not be started") + + monkeypatch.setattr(claude_module, "_start_agent_subprocess_detached", start_agent) + service = ClaudeAgentService(workspace_base_dir=tmp_path) + + result = await service.provision_endpoint(endpoint_model, pipeline_hinter=None) + + assert result.in_progress is True + assert result.error is None + assert result.final_report is None + + @pytest.mark.asyncio + async def test_writes_redacted_trace_in_debug_mode( + self, + session: AsyncSession, + tmp_path, + monkeypatch, + ): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "agent-secret") + _configure_fake_claude(tmp_path, monkeypatch) + monkeypatch.setattr(settings, "LOG_LEVEL", "DEBUG") + run_id = uuid.uuid4() + + async def runner(workspace, request): + _write_trace_record( + workspace, + {"type": "FakeMessage", "text": "hf-secret agent-secret"}, + ) + return _AgentRunnerResult( + report=_successful_agent_report(run_id=run_id, service_yaml="token: hf-secret\n") + ) + + endpoint_model = await _create_endpoint_model(session) + service = ClaudeAgentService(runner=runner, workspace_base_dir=tmp_path) + + result = await service.provision_endpoint(endpoint_model, pipeline_hinter=None) + + assert result.error is None + trace_path = tmp_path / str(endpoint_model.id) / "1" / "trace.jsonl" + trace = trace_path.read_text() + assert "agent-secret" not in trace + assert "hf-secret" not in trace + assert "[redacted]" in trace + event = json.loads(trace.splitlines()[0]) + assert event["type"] == "FakeMessage" + + @pytest.mark.asyncio + async def test_returns_error_when_sdk_fails_before_report( + self, + session: AsyncSession, + tmp_path, + monkeypatch, + ): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "agent-secret") + _configure_fake_claude(tmp_path, monkeypatch) + + async def runner(workspace, request): + return _AgentRunnerResult( + error="Server agent failed before returning a final report: bad key" + ) + + endpoint_model = await _create_endpoint_model(session) + service = ClaudeAgentService(runner=runner, workspace_base_dir=tmp_path) + + result = await service.provision_endpoint(endpoint_model, pipeline_hinter=None) + + assert result.error == "Server agent failed before returning a final report: bad key" + assert result.final_report is None + work_dir = tmp_path / str(endpoint_model.id) / "1" / "workspace" + state = json.loads((work_dir / "agent_state.json").read_text()) + assert state["phase"] == "failure" + agent_error = json.loads((work_dir / "agent_error.json").read_text()) + assert agent_error["success"] is False + assert agent_error["failure_summary"] == ( + "Server agent failed before returning a final report: bad key" + ) + + def test_returns_error_when_configured_claude_path_is_not_executable( + self, + tmp_path, + monkeypatch: pytest.MonkeyPatch, + ): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "agent-secret") + monkeypatch.setattr(settings, "AGENT_CLAUDE_PATH", str(tmp_path / "missing-claude")) + + assert get_claude_agent_unavailable_reason() == ( + "DSTACK_AGENT_CLAUDE_PATH does not resolve to an executable: " + f"{tmp_path / 'missing-claude'}" + ) + + def test_returns_error_when_agent_key_is_blank(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "") + monkeypatch.setattr(settings, "AGENT_CLAUDE_USE_EXISTING_AUTH", False) + + assert get_claude_agent_unavailable_reason() == ( + "DSTACK_AGENT_ANTHROPIC_API_KEY is not set. Set it or opt in to existing " + "Claude CLI auth with DSTACK_AGENT_CLAUDE_USE_EXISTING_AUTH=1." + ) + + @pytest.mark.parametrize("effort", ["low", "medium", "high", "xhigh", "max"]) + def test_claude_command_includes_configured_effort( + self, tmp_path, monkeypatch: pytest.MonkeyPatch, effort: str + ): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "agent-secret") + monkeypatch.setattr(settings, "AGENT_CLAUDE_USE_EXISTING_AUTH", False) + monkeypatch.setattr(settings, "AGENT_CLAUDE_EFFORT", effort) + _configure_fake_claude(tmp_path, monkeypatch) + + command = _build_claude_command( + { + "prompt": "prompt", + "options": { + "allowed_tools": "Bash", + "disallowed_tools": "", + "model": "test-model", + "max_turns": None, + "json_schema": {}, + }, + } + ) + + assert get_claude_agent_unavailable_reason() is None + assert command[command.index("--effort") + 1] == effort + + def test_returns_error_when_claude_effort_is_invalid( + self, tmp_path, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "agent-secret") + monkeypatch.setattr(settings, "AGENT_CLAUDE_USE_EXISTING_AUTH", False) + monkeypatch.setattr(settings, "AGENT_CLAUDE_EFFORT", "turbo") + _configure_fake_claude(tmp_path, monkeypatch) + + assert get_claude_agent_unavailable_reason() == ( + "DSTACK_AGENT_CLAUDE_EFFORT must be one of: low, medium, high, xhigh, max." + ) + + def test_existing_claude_auth_opt_in_enables_agent_without_api_key( + self, tmp_path, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setattr(settings, "AGENT_CLAUDE_PATH", None) + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", None) + monkeypatch.setattr(settings, "AGENT_CLAUDE_USE_EXISTING_AUTH", True) + monkeypatch.setenv("USER", "server-user") + monkeypatch.setattr( + claude_module.shutil, + "which", + lambda name: f"/usr/local/bin/{name}" if name in {"claude", "dstack", "ssh"} else None, + ) + + command = _build_claude_command( + { + "prompt": "prompt", + "options": { + "allowed_tools": "Bash", + "disallowed_tools": "", + "model": "test-model", + "max_turns": None, + "json_schema": {}, + }, + } + ) + env = claude_module._build_agent_env( + home_dir=claude_module._get_claude_process_home_dir(Path("/tmp/agent-home")), + project_name="main", + server_url="http://127.0.0.1:8000", + token="token", + endpoint_env={}, + bin_dir=Path("/tmp/agent-bin"), + ) + dstack_home_dir = tmp_path / "short-home" + workspace = _AgentWorkspace( + root_dir=tmp_path, + home_dir=tmp_path / "home", + work_dir=tmp_path / "work", + trace_path=None, + env=env, + redacted_values=[], + endpoint_name="qwen-endpoint", + model="Qwen/Qwen3-0.6B", + dstack_home_dir=dstack_home_dir, + ) + + assert get_claude_agent_unavailable_reason() is None + assert command[0] == "/usr/local/bin/claude" + assert "--bare" not in command + assert command[command.index("--setting-sources") + 1] == "project,local" + assert env["HOME"] == str(Path.home()) + assert env["USER"] == "server-user" + assert "ANTHROPIC_API_KEY" not in env + claude_module._install_agent_helper_scripts(workspace) + wrapper = tmp_path / "bin" / "dstack" + assert wrapper.exists() + wrapper_text = wrapper.read_text() + assert str(dstack_home_dir) in wrapper_text + assert str(tmp_path / "home") not in wrapper_text + ssh_wrapper = tmp_path / "bin" / "ssh" + assert ssh_wrapper.exists() + ssh_wrapper_text = ssh_wrapper.read_text() + assert str(dstack_home_dir) in ssh_wrapper_text + assert str(tmp_path / "home") not in ssh_wrapper_text + + def test_returns_error_when_api_key_and_existing_claude_auth_are_both_set( + self, monkeypatch: pytest.MonkeyPatch + ): + monkeypatch.setattr(settings, "AGENT_CLAUDE_PATH", None) + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "agent-secret") + monkeypatch.setattr(settings, "AGENT_CLAUDE_USE_EXISTING_AUTH", True) + + assert get_claude_agent_unavailable_reason() == ( + "DSTACK_AGENT_ANTHROPIC_API_KEY and " + "DSTACK_AGENT_CLAUDE_USE_EXISTING_AUTH cannot both be set." + ) + + def test_existing_claude_auth_logs_development_warning_once( + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + ): + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", None) + monkeypatch.setattr(settings, "AGENT_CLAUDE_USE_EXISTING_AUTH", True) + monkeypatch.setattr(claude_module, "_existing_auth_warning_logged", False) + + claude_module._warn_if_using_existing_claude_auth() + claude_module._warn_if_using_existing_claude_auth() + + messages = [ + record.message + for record in caplog.records + if "DSTACK_AGENT_CLAUDE_USE_EXISTING_AUTH=1 is enabled" in record.message + ] + assert len(messages) == 1 + assert "only for local development" in messages[0] + assert "Production servers must set DSTACK_AGENT_ANTHROPIC_API_KEY" in messages[0] + + def test_falls_back_to_claude_in_path(self, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(settings, "AGENT_CLAUDE_PATH", None) + monkeypatch.setattr(settings, "AGENT_ANTHROPIC_API_KEY", "agent-secret") + monkeypatch.setattr( + claude_module.shutil, + "which", + lambda name: "/usr/local/bin/claude" if name == "claude" else None, + ) + + command = _build_claude_command( + { + "prompt": "prompt", + "options": { + "allowed_tools": "Bash", + "disallowed_tools": "", + "model": "test-model", + "max_turns": 1, + "json_schema": {}, + }, + } + ) + + assert get_claude_agent_unavailable_reason() is None + assert command[0] == "/usr/local/bin/claude" + + @pytest.mark.asyncio + async def test_reads_claude_stream_incrementally(self, tmp_path): + reader = StreamReader() + workspace = _AgentWorkspace( + root_dir=tmp_path, + home_dir=tmp_path / "home", + work_dir=tmp_path / "work", + trace_path=tmp_path / "trace.jsonl", + env={}, + redacted_values=["secret"], + endpoint_name="qwen-endpoint", + model="Qwen/Qwen3-0.6B", + ) + report = { + "success": True, + "run_id": str(uuid.uuid4()), + "run_name": "qwen-endpoint-1", + "service_yaml": "env: secret\n", + "base": "Qwen/Qwen3-0.6B", + "model": "Qwen/Qwen3-0.6B", + } + reader.feed_data( + json.dumps({"type": "assistant", "message": {"content": "working"}}).encode() + b"\n" + ) + reader.feed_data( + json.dumps({"type": "result", "result": json.dumps(report)}).encode() + b"\n" + ) + reader.feed_eof() + + output = await _read_agent_stdout(reader, workspace) + + assert output.report_data == report + trace = (tmp_path / "trace.jsonl").read_text() + assert "secret" not in trace + assert "[redacted]" in trace + + def test_loads_submitted_run_names_without_run_ids(self, tmp_path): + workspace = _AgentWorkspace( + root_dir=tmp_path, + home_dir=tmp_path / "home", + work_dir=tmp_path / "work", + trace_path=tmp_path / "trace.jsonl", + env={}, + redacted_values=[], + endpoint_name="qwen-endpoint", + model="Qwen/Qwen3-0.6B", + ) + workspace.work_dir.mkdir(parents=True) + run_id = uuid.uuid4() + (workspace.work_dir / "submissions.jsonl").write_text( + "\n".join( + [ + json.dumps({"name": "qwen-endpoint-1", "run_id": None}), + json.dumps({"name": "qwen-endpoint-1", "run_id": str(run_id)}), + json.dumps({"name": "qwen-endpoint-2"}), + ] + ) + + "\n", + encoding="utf-8", + ) + + submissions = _load_submissions(workspace) + + assert submissions.run_ids == (run_id,) + assert submissions.run_names == ("qwen-endpoint-1", "qwen-endpoint-2") + + @pytest.mark.asyncio + async def test_stores_assistant_text_without_copying_stream_to_endpoint_logs( + self, + tmp_path, + ): + class FakeLogWriter: + def __init__(self): + self.messages = [] + + async def write(self, message): + self.messages.append(message) + + async def flush(self): + pass + + reader = StreamReader() + log_writer = FakeLogWriter() + workspace = _AgentWorkspace( + root_dir=tmp_path, + home_dir=tmp_path / "home", + work_dir=tmp_path / "work", + trace_path=None, + env={}, + redacted_values=["secret"], + endpoint_name="qwen-endpoint", + model="Qwen/Qwen3-0.6B", + log_writer=log_writer, + ) + reader.feed_data( + json.dumps( + { + "type": "assistant", + "message": { + "content": [ + { + "type": "text", + "text": "Checking offers without exposing secret.", + }, + { + "type": "tool_use", + "id": "tool-1", + "name": "Bash", + "input": { + "description": "List offers", + "command": "dstack offer --gpu 1 --max-price 0.3", + }, + }, + ] + }, + } + ).encode() + + b"\n" + ) + reader.feed_data( + json.dumps( + { + "type": "user", + "message": { + "content": [ + { + "type": "tool_result", + "tool_use_id": "tool-1", + "content": "No offers secret", + "is_error": False, + } + ] + }, + } + ).encode() + + b"\n" + ) + reader.feed_eof() + + await _read_agent_stdout(reader, workspace) + + assert log_writer.messages == [] + command_records = [ + json.loads(line) + for line in (tmp_path / "work" / "commands.jsonl").read_text().splitlines() + ] + assert command_records[0]["event"] == "tool_use" + assert command_records[0]["command"] == "dstack offer --gpu 1 --max-price 0.3" + assert command_records[1]["event"] == "tool_result" + assert command_records[1]["tool_use_id"] == "tool-1" + output = (tmp_path / "work" / command_records[1]["output_path"]).read_text() + assert output == "No offers [redacted]" + + @pytest.mark.asyncio + async def test_stores_large_tool_outputs_without_endpoint_log_preview(self, tmp_path): + class FakeLogWriter: + def __init__(self): + self.messages = [] + + async def write(self, message): + self.messages.append(message) + + async def flush(self): + pass + + reader = StreamReader() + log_writer = FakeLogWriter() + workspace = _AgentWorkspace( + root_dir=tmp_path, + home_dir=tmp_path / "home", + work_dir=tmp_path / "work", + trace_path=None, + env={}, + redacted_values=[], + endpoint_name="qwen-endpoint", + model="Qwen/Qwen3-0.6B", + log_writer=log_writer, + ) + long_output = "\n".join(f"offer {i:04d} gpu=A5000 price=0.27" for i in range(200)) + reader.feed_data( + json.dumps( + { + "type": "user", + "message": { + "content": [ + { + "type": "tool_result", + "tool_use_id": "tool-1", + "content": long_output, + "is_error": False, + } + ] + }, + } + ).encode() + + b"\n" + ) + reader.feed_eof() + + await _read_agent_stdout(reader, workspace) + + assert log_writer.messages == [] + command_records = [ + json.loads(line) + for line in (tmp_path / "work" / "commands.jsonl").read_text().splitlines() + ] + output = (tmp_path / "work" / command_records[0]["output_path"]).read_text() + assert output == long_output + + @pytest.mark.asyncio + async def test_subprocess_streams_agent_progress_jsonl_to_endpoint_logs( + self, + tmp_path, + monkeypatch, + ): + class FakeLogWriter: + def __init__(self): + self.messages = [] + + async def write(self, message): + self.messages.append(message) + + async def flush(self): + pass + + claude_path = tmp_path / "claude" + claude_path.write_text( + """#!/bin/sh +printf '%s\\n' 'Checking recipes with secret' >> progress.jsonl +printf '%s\\n' '{"phase":"submit","message":"Submitted service run"}' >> progress.jsonl +cat > final_report.json <<'JSON' +{ + "success": true, + "run_id": "6e578748-d597-4fde-a3a4-203587cad5a2", + "run_name": "qwen-endpoint-1", + "service_yaml": "type: service\\nname: qwen-endpoint-1\\n", + "base": "Qwen/Qwen3-0.6B", + "model": "Qwen/Qwen3-0.6B" +} +JSON +printf '%s\\n' '{"type":"result","is_error":false,"result":"done"}' +""", + encoding="utf-8", + ) + claude_path.chmod(0o755) + monkeypatch.setattr(settings, "AGENT_CLAUDE_PATH", str(claude_path)) + work_dir = tmp_path / "work" + work_dir.mkdir() + log_writer = FakeLogWriter() + workspace = _AgentWorkspace( + root_dir=tmp_path, + home_dir=tmp_path / "home", + work_dir=work_dir, + trace_path=None, + env={}, + redacted_values=["secret"], + endpoint_name="qwen-endpoint", + model="Qwen/Qwen3-0.6B", + log_writer=log_writer, + ) + request = { + "prompt": "prompt", + "env": {}, + "cwd": str(work_dir), + "options": { + "allowed_tools": "Bash", + "disallowed_tools": "", + "model": "test-model", + "max_turns": 1, + "json_schema": {}, + }, + } + (work_dir / "progress.jsonl").write_text( + '{"phase":"old","message":"Old progress must not replay"}\n', + encoding="utf-8", + ) + + result = await _run_agent_in_subprocess(workspace, request) + + assert result.error is None + assert log_writer.messages == [ + "Checking recipes with [redacted]", + "Submitted service run", + ] + + @pytest.mark.asyncio + async def test_subprocess_uses_final_report_artifact_when_stream_result_is_missing( + self, + tmp_path, + monkeypatch, + ): + claude_path = tmp_path / "claude" + claude_path.write_text( + """#!/bin/sh +cat > final_report.json <<'JSON' +{ + "success": true, + "run_id": "6e578748-d597-4fde-a3a4-203587cad5a2", + "run_name": "qwen-endpoint-1", + "service_yaml": "type: service\\nname: qwen-endpoint-1\\n", + "base": "Qwen/Qwen3-0.6B", + "model": "Qwen/Qwen3-0.6B" +} +JSON +printf '%s\\n' '{"type":"result","is_error":false,"result":"done"}' +""", + encoding="utf-8", + ) + claude_path.chmod(0o755) + monkeypatch.setattr(settings, "AGENT_CLAUDE_PATH", str(claude_path)) + (tmp_path / "work").mkdir() + workspace = _AgentWorkspace( + root_dir=tmp_path, + home_dir=tmp_path / "home", + work_dir=tmp_path / "work", + trace_path=None, + env={}, + redacted_values=[], + endpoint_name="qwen-endpoint", + model="Qwen/Qwen3-0.6B", + ) + request = { + "prompt": "prompt", + "env": {}, + "cwd": str(tmp_path / "work"), + "options": { + "allowed_tools": "Bash", + "disallowed_tools": "", + "model": "test-model", + "max_turns": 1, + "json_schema": {}, + }, + } + + result = await _run_agent_in_subprocess(workspace, request) + + assert result.error is None + assert result.report is not None + assert result.report.success is True + assert result.report.run_name == "qwen-endpoint-1" + + @pytest.mark.asyncio + async def test_subprocess_no_report_error_does_not_include_stream_output( + self, + tmp_path, + monkeypatch, + ): + claude_path = tmp_path / "claude" + claude_path.write_text( + "#!/bin/sh\n" + 'printf \'%s\\n\' \'{"type":"user","message":{"content":[{"type":"tool_result","content":"huge offer table","is_error":false}]}}\' >&2\n' + "exit 143\n", + encoding="utf-8", + ) + claude_path.chmod(0o755) + monkeypatch.setattr(settings, "AGENT_CLAUDE_PATH", str(claude_path)) + (tmp_path / "work").mkdir() + workspace = _AgentWorkspace( + root_dir=tmp_path, + home_dir=tmp_path / "home", + work_dir=tmp_path / "work", + trace_path=None, + env={}, + redacted_values=[], + endpoint_name="qwen-endpoint", + model="Qwen/Qwen3-0.6B", + ) + request = { + "prompt": "prompt", + "env": {}, + "options": { + "allowed_tools": "Bash", + "disallowed_tools": "", + "model": "test-model", + "max_turns": 1, + "json_schema": {}, + }, + } + + result = await _run_agent_in_subprocess(workspace, request) + + assert result.error == ( + "Server agent process exited without a final report (return code 143)" + ) + assert "huge offer table" not in result.error diff --git a/src/tests/_internal/server/services/test_endpoint_presets.py b/src/tests/_internal/server/services/test_endpoint_presets.py new file mode 100644 index 0000000000..93f9e6bc9f --- /dev/null +++ b/src/tests/_internal/server/services/test_endpoint_presets.py @@ -0,0 +1,1346 @@ +from unittest.mock import AsyncMock, Mock, patch + +import pytest +from sqlalchemy.ext.asyncio import AsyncSession + +from dstack._internal.core.models.backends.base import BackendType +from dstack._internal.core.models.configurations import ServiceConfiguration +from dstack._internal.core.models.endpoints import EndpointConfiguration +from dstack._internal.core.models.envs import Env +from dstack._internal.core.models.instances import ( + Disk, + Gpu, + InstanceAvailability, + InstanceOfferWithAvailability, + InstanceRuntime, + InstanceType, + Resources, +) +from dstack._internal.core.models.profiles import CreationPolicy +from dstack._internal.core.models.runs import JobStatus, RunSpec, RunStatus +from dstack._internal.core.models.users import GlobalRole +from dstack._internal.server.services.endpoints.names import get_endpoint_serving_run_name +from dstack._internal.server.services.endpoints.planning import ( + build_preset_run_spec, + build_preset_service_configuration, + find_matching_preset_plan, + find_preset_planning_result, +) +from dstack._internal.server.services.endpoints.preset_building import ( + build_endpoint_preset_from_run, +) +from dstack._internal.server.services.endpoints.presets import ( + EndpointPreset, + EndpointPresetRecipe, + EndpointPresetValidation, + EndpointPresetValidationReplica, + LocalDirEndpointPresetService, + make_endpoint_preset_recipe_id, +) +from dstack._internal.server.testing.common import ( + create_job, + create_project, + create_repo, + create_run, + create_user, + get_job_runtime_data, +) + +PROJECT_NAME = "test_project" + + +class TestLocalDirEndpointPresetService: + @pytest.mark.asyncio + async def test_lists_valid_endpoint_presets(self, tmp_path): + presets_dir = _project_presets_dir(tmp_path) + presets_dir.mkdir(parents=True) + (presets_dir / "qwen.dstack.yml").write_text( + """\ +model: Qwen/Qwen3-4B +type: endpoint-preset +recipes: + - id: vllm-t4 + service: + image: vllm/vllm-openai:latest + commands: + - vllm serve Qwen/Qwen3-4B --host 0.0.0.0 --port 8000 + port: 8000 + resources: + gpu: 16GB + validations: + - replicas: + - resources: + - cpu: 4 + memory: 16GB + disk: 100GB + gpu: + name: T4 + memory: 16GB + count: 1 +""" + ) + + presets = await LocalDirEndpointPresetService(tmp_path).list_presets(PROJECT_NAME) + + assert len(presets) == 1 + assert presets[0].base == "Qwen/Qwen3-4B" + assert [recipe.id for recipe in presets[0].recipes] == ["vllm-t4"] + recipe = presets[0].recipes[0] + assert recipe.service.resources.gpu is not None + assert recipe.service.resources.gpu.vendor is None + assert recipe.validations[0].replicas[0].resources[0].gpu is not None + + @pytest.mark.asyncio + async def test_lists_replica_group_presets_in_group_order(self, tmp_path): + presets_dir = _project_presets_dir(tmp_path) + presets_dir.mkdir(parents=True) + (presets_dir / "qwen.yml").write_text( + """\ +type: endpoint-preset +model: Qwen/Qwen3-4B +recipes: + - id: grouped + service: + port: 8000 + model: Qwen/Qwen3-4B + replicas: + - name: router + count: 1 + image: ghcr.io/example/router:latest + commands: + - python router.py + resources: + cpu: 4 + - name: worker + count: 2 + image: vllm/vllm-openai:latest + commands: + - vllm serve Qwen/Qwen3-4B --host 0.0.0.0 --port 8000 + resources: + gpu: 24GB + validations: + - replicas: + - resources: + - cpu: 8 + memory: 16GB + disk: 100GB + gpu: 0 + - resources: + - cpu: 14 + memory: 64GB + disk: 200GB + gpu: + name: L4 + memory: 24GB + count: 1 + - cpu: 14 + memory: 64GB + disk: 200GB + gpu: + name: L4 + memory: 24GB + count: 1 +""" + ) + + presets = await LocalDirEndpointPresetService(tmp_path).list_presets(PROJECT_NAME) + + assert len(presets) == 1 + replica_groups = presets[0].recipes[0].service.replica_groups + assert [group.name for group in replica_groups] == ["router", "worker"] + assert replica_groups[0].resources.gpu is not None + assert replica_groups[0].resources.gpu.count.min == 0 + assert replica_groups[0].resources.cpu.count.min == 4 + assert replica_groups[1].resources.gpu is not None + + @pytest.mark.asyncio + async def test_loads_legacy_replica_spec_groups_as_recipe_validation(self, tmp_path): + presets_dir = _project_presets_dir(tmp_path) + presets_dir.mkdir(parents=True) + (presets_dir / "qwen-legacy.yml").write_text( + """\ +model: Qwen/Qwen3-4B +type: endpoint-preset +service: + image: vllm/vllm-openai:latest + commands: + - vllm serve Qwen/Qwen3-4B --host 0.0.0.0 --port 8000 + port: 8000 +replica_spec_groups: + - resources: + gpu: 16GB + tested_resources: + - cpu: 4 + memory: 16GB + disk: 100GB + gpu: + name: T4 + memory: 16GB + count: 1 +""" + ) + + presets = await LocalDirEndpointPresetService(tmp_path).list_presets(PROJECT_NAME) + + assert len(presets) == 1 + recipe = presets[0].recipes[0] + assert recipe.service.resources.gpu is not None + assert recipe.service.resources.gpu.vendor is None + assert recipe.validations[0].replicas[0].resources[0].gpu is not None + + @pytest.mark.asyncio + async def test_rejects_service_gpu_vendor_mismatch(self, tmp_path, caplog): + presets_dir = _project_presets_dir(tmp_path) + presets_dir.mkdir(parents=True) + (presets_dir / "qwen.yml").write_text( + """\ +type: endpoint-preset +model: Qwen/Qwen3-4B +recipes: + - id: vllm-amd + service: + commands: + - vllm serve Qwen/Qwen3-4B --host 0.0.0.0 --port 8000 + port: 8000 + resources: + gpu: amd:16GB + validations: + - replicas: + - resources: + - cpu: 4 + memory: 16GB + disk: 100GB + gpu: + name: T4 + memory: 16GB + count: 1 +""" + ) + + presets = await LocalDirEndpointPresetService(tmp_path).list_presets(PROJECT_NAME) + + assert presets == [] + assert "preset service GPU vendor does not match validation" in caplog.text + + @pytest.mark.asyncio + async def test_saves_preset_with_comments_and_merges_same_base(self, tmp_path): + preset = _qwen_preset(recipe_id="vllm-t4") + extra_preset = _qwen_preset(recipe_id="vllm-l4", resources={"gpu": "24GB"}) + service = LocalDirEndpointPresetService(tmp_path) + + saved = await service.save_preset( + PROJECT_NAME, + preset, + comments=[ + "endpoint: qwen-endpoint", + "run: 00000000-0000-0000-0000-000000000001", + ], + ) + assert saved.base == "Qwen/Qwen3-4B" + presets_dir = _project_presets_dir(tmp_path) + preset_files = list(presets_dir.glob("*.dstack.yml")) + assert len(preset_files) == 1 + text = preset_files[0].read_text() + assert text.startswith( + "# endpoint: qwen-endpoint\n# run: 00000000-0000-0000-0000-000000000001\n" + ) + assert "SECRET_VALUE" not in text + assert "preset-service-name" not in text + assert "creation_policy" not in text + assert "validations:" in text + assert saved.recipes[0].model == "Qwen/Qwen3-4B" + assert "HF_HOME=/root/.cache/huggingface" in text + assert "HF_TOKEN" in text + saved_again = await service.save_preset(PROJECT_NAME, extra_preset) + + assert [recipe.id for recipe in saved_again.recipes] == ["vllm-l4", "vllm-t4"] + assert len(list(presets_dir.glob("*.dstack.yml"))) == 1 + + @pytest.mark.asyncio + async def test_merges_duplicate_base_files_for_list_get_save_and_delete(self, tmp_path): + service = LocalDirEndpointPresetService(tmp_path) + presets_dir = _project_presets_dir(tmp_path) + presets_dir.mkdir(parents=True) + _write_qwen_preset_file(presets_dir / "qwen-a.yml", recipe_id="vllm-t4", gpu="16GB") + _write_qwen_preset_file(presets_dir / "qwen-b.yml", recipe_id="vllm-l4", gpu="24GB") + + presets = await service.list_presets(PROJECT_NAME) + preset = await service.get_preset(PROJECT_NAME, "Qwen/Qwen3-4B") + + assert len(presets) == 1 + assert preset is not None + assert [recipe.id for recipe in presets[0].recipes] == ["vllm-t4", "vllm-l4"] + assert [recipe.id for recipe in preset.recipes] == ["vllm-t4", "vllm-l4"] + + saved = await service.save_preset( + PROJECT_NAME, + _qwen_preset(recipe_id="vllm-a100", resources={"gpu": "80GB"}), + ) + + assert [recipe.id for recipe in saved.recipes] == ["vllm-a100", "vllm-t4", "vllm-l4"] + assert len(list(presets_dir.glob("*.yml"))) == 1 + + await service.delete_preset(PROJECT_NAME, "Qwen/Qwen3-4B") + + assert await service.list_presets(PROJECT_NAME) == [] + assert list(presets_dir.glob("*.yml")) == [] + + @pytest.mark.asyncio + async def test_saved_preset_round_trips(self, tmp_path): + saved = await LocalDirEndpointPresetService(tmp_path).save_preset( + PROJECT_NAME, + _qwen_preset(recipe_id="vllm-t4"), + ) + + presets = await LocalDirEndpointPresetService(tmp_path).list_presets(PROJECT_NAME) + + assert saved.base == "Qwen/Qwen3-4B" + assert len(presets) == 1 + assert presets[0].base == "Qwen/Qwen3-4B" + recipe = presets[0].recipes[0] + assert recipe.model == "Qwen/Qwen3-4B" + assert recipe.service.model is not None + assert recipe.service.model.name == "Qwen/Qwen3-4B" + assert set(recipe.service.env.keys()) == {"HF_HOME", "HF_TOKEN"} + assert recipe.service.env["HF_HOME"] == "/root/.cache/huggingface" + assert recipe.service.resources.gpu is not None + + @pytest.mark.asyncio + async def test_loads_recipe_model_that_differs_from_preset_base(self, tmp_path): + presets_dir = _project_presets_dir(tmp_path) + presets_dir.mkdir(parents=True) + (presets_dir / "qwen.yml").write_text( + """\ +type: endpoint-preset +model: Qwen/Qwen3-4B +recipes: + - id: vllm-t4 + model: groxaxo/Qwen3-4B-GPTQ-4Bit + service: + commands: + - vllm serve groxaxo/Qwen3-4B-GPTQ-4Bit --served-model-name Qwen/Qwen3-4B + port: 8000 + model: Qwen/Qwen3-4B + resources: + gpu: 16GB + validations: + - replicas: + - resources: + - cpu: 4 + memory: 16GB + disk: 100GB + gpu: + name: T4 + memory: 16GB + count: 1 +""" + ) + + presets = await LocalDirEndpointPresetService(tmp_path).list_presets(PROJECT_NAME) + + recipe = presets[0].recipes[0] + assert recipe.model == "groxaxo/Qwen3-4B-GPTQ-4Bit" + assert recipe.service.model is not None + assert recipe.service.model.name == "Qwen/Qwen3-4B" + + @pytest.mark.asyncio + async def test_legacy_recipe_without_model_uses_preset_base(self, tmp_path): + _write_qwen_preset(tmp_path) + + presets = await LocalDirEndpointPresetService(tmp_path).list_presets(PROJECT_NAME) + + assert presets[0].recipes[0].model == "Qwen/Qwen3-4B" + + @pytest.mark.asyncio + async def test_rejects_same_recipe_id_with_different_recipe_model(self, tmp_path): + service = LocalDirEndpointPresetService(tmp_path) + await service.save_preset( + PROJECT_NAME, + _qwen_preset(recipe_id="vllm-t4", recipe_model="Qwen/Qwen3-4B-GPTQ-A"), + ) + + with pytest.raises(ValueError, match="endpoint preset recipe id conflict"): + await service.save_preset( + PROJECT_NAME, + _qwen_preset(recipe_id="vllm-t4", recipe_model="Qwen/Qwen3-4B-GPTQ-B"), + ) + + @pytest.mark.asyncio + async def test_deletes_preset_by_model(self, tmp_path): + service = LocalDirEndpointPresetService(tmp_path) + saved = await service.save_preset(PROJECT_NAME, _qwen_preset()) + + await service.delete_preset(PROJECT_NAME, saved.base) + + assert await service.list_presets(PROJECT_NAME) == [] + assert list(_project_presets_dir(tmp_path).glob("*.dstack.yml")) == [] + + @pytest.mark.asyncio + async def test_delete_missing_preset_raises(self, tmp_path): + service = LocalDirEndpointPresetService(tmp_path) + + with pytest.raises(FileNotFoundError): + await service.delete_preset(PROJECT_NAME, "missing") + + @pytest.mark.asyncio + async def test_ignores_non_yaml_files(self, tmp_path): + presets_dir = _project_presets_dir(tmp_path) + presets_dir.mkdir(parents=True) + (presets_dir / "notes.txt").write_text("ignored") + + presets = await LocalDirEndpointPresetService(tmp_path).list_presets(PROJECT_NAME) + + assert presets == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("filename", "content", "message"), + [ + ( + "task.yml", + "type: task\ncommands:\n - echo nope\n", + "preset must be an endpoint preset", + ), + ( + "missing-base.yml", + """\ +type: endpoint-preset +recipes: + - id: vllm + service: + commands: + - python -m http.server 8000 + port: 8000 + resources: + gpu: 16GB + validations: + - replicas: + - resources: + - cpu: 4 + memory: 16GB + disk: 100GB + gpu: 0 +""", + "preset must specify a base", + ), + ( + "base-model-mismatch.yml", + """\ +type: endpoint-preset +base: Qwen/Qwen3-4B +model: Qwen/Qwen3-8B +recipes: + - id: vllm + service: + commands: + - python -m http.server 8000 + port: 8000 + resources: + gpu: 16GB + validations: + - replicas: + - resources: + - cpu: 4 + memory: 16GB + disk: 100GB + gpu: 0 +""", + "preset base must match legacy model", + ), + ( + "missing-recipes.yml", + """\ +type: endpoint-preset +model: Qwen/Qwen3-4B +""", + "preset recipe must specify a service object", + ), + ( + "missing-service.yml", + """\ +type: endpoint-preset +model: Qwen/Qwen3-4B +recipes: + - id: vllm + validations: + - replicas: + - resources: + - cpu: 4 + memory: 16GB + disk: 100GB + gpu: 0 +""", + "preset recipe must specify a service object", + ), + ( + "service-profile.yml", + """\ +type: endpoint-preset +model: Qwen/Qwen3-4B +recipes: + - id: vllm + service: + commands: + - python -m http.server 8000 + port: 8000 + resources: + gpu: 16GB + creation_policy: reuse + validations: + - replicas: + - resources: + - cpu: 4 + memory: 16GB + disk: 100GB + gpu: + name: T4 + memory: 16GB + count: 1 +""", + "preset service object must not specify profile fields", + ), + ( + "missing-service-resources.yml", + """\ +type: endpoint-preset +model: Qwen/Qwen3-4B +recipes: + - id: vllm + service: + commands: + - python -m http.server 8000 + port: 8000 + validations: + - replicas: + - resources: + - cpu: 4 + memory: 16GB + disk: 100GB + gpu: + name: T4 + memory: 16GB + count: 1 +""", + "preset service object must specify resources", + ), + ( + "group-missing-resources.yml", + """\ +type: endpoint-preset +model: Qwen/Qwen3-4B +recipes: + - id: vllm + service: + port: 8000 + replicas: + - name: worker + count: 1 + image: vllm/vllm-openai:latest + validations: + - replicas: + - resources: + - cpu: 4 + memory: 16GB + disk: 100GB + gpu: + name: T4 + memory: 16GB + count: 1 +""", + "preset service replica groups must specify resources", + ), + ( + "loose-resources.yml", + """\ +type: endpoint-preset +model: Qwen/Qwen3-4B +recipes: + - id: vllm + service: + commands: + - vllm serve Qwen/Qwen3-4B --host 0.0.0.0 --port 8000 + port: 8000 + resources: + gpu: 16GB + validations: + - replicas: + - resources: + - gpu: 16GB +""", + "preset validations must use exact replica resources", + ), + ( + "mismatched-validation-replicas.yml", + """\ +type: endpoint-preset +model: Qwen/Qwen3-4B +recipes: + - id: grouped + service: + port: 8000 + replicas: + - name: router + count: 1 + image: ghcr.io/example/router:latest + resources: + cpu: 4 + - name: worker + count: 1 + image: vllm/vllm-openai:latest + resources: + gpu: 16GB + validations: + - replicas: + - resources: + - cpu: 4 + memory: 16GB + disk: 100GB + gpu: 0 +""", + "preset validation replicas must match service replica group order", + ), + ( + "bad-replica-group-shape.yml", + """\ +type: endpoint-preset +model: Qwen/Qwen3-4B +recipes: + - id: grouped + service: + port: 8000 + replicas: + - worker + validations: + - replicas: + - resources: + - cpu: 4 + memory: 16GB + disk: 100GB + gpu: + name: T4 + memory: 16GB + count: 1 +""", + "preset service replica groups must be objects", + ), + ("not-yaml.yml", ":\n", "while parsing a block mapping"), + ], + ) + async def test_invalid_preset_is_skipped_and_logged( + self, tmp_path, caplog, filename, content, message + ): + presets_dir = _project_presets_dir(tmp_path) + presets_dir.mkdir(parents=True) + (presets_dir / filename).write_text(content) + + presets = await LocalDirEndpointPresetService(tmp_path).list_presets(PROJECT_NAME) + + assert presets == [] + assert filename in caplog.text + assert message in caplog.text + + +class TestBuildEndpointPresetFromRun: + @pytest.mark.asyncio + async def test_builds_implicit_replica_group_from_running_service( + self, session: AsyncSession, tmp_path + ): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_name="qwen-endpoint", + status=RunStatus.RUNNING, + run_spec=RunSpec( + run_name="qwen-endpoint", + configuration=ServiceConfiguration.parse_obj( + { + "type": "service", + "name": "qwen-endpoint", + "commands": [ + "vllm serve Qwen/Qwen3-4B --host 0.0.0.0 --port 8000", + ], + "port": 8000, + "model": "Qwen/Qwen3-4B", + "resources": {"gpu": "16GB"}, + } + ), + ), + ) + await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + registered=True, + job_runtime_data=get_job_runtime_data( + offer=_instance_offer(gpu_name="L4", gpu_memory_gib=24, cpu_count=14) + ), + ) + await session.refresh(run, attribute_names=["jobs"]) + + preset = build_endpoint_preset_from_run(run) + variant_preset = build_endpoint_preset_from_run( + run, + recipe_model="groxaxo/Qwen3-4B-GPTQ-4Bit", + ) + saved = await LocalDirEndpointPresetService(tmp_path).save_preset(PROJECT_NAME, preset) + + assert saved.base == "Qwen/Qwen3-4B" + assert variant_preset.base == "Qwen/Qwen3-4B" + assert variant_preset.recipes[0].model == "groxaxo/Qwen3-4B-GPTQ-4Bit" + recipe = saved.recipes[0] + assert recipe.model == "Qwen/Qwen3-4B" + assert recipe.id == make_endpoint_preset_recipe_id(recipe.service) + assert recipe.service.resources.gpu is not None + assert recipe.service.resources.gpu.vendor is not None + assert recipe.service.resources.gpu.vendor.value == "nvidia" + assert "gpu=nvidia:16GB" in recipe.service.resources.pretty_format() + resources = recipe.validations[0].replicas[0].resources[0].pretty_format() + assert "cpu=14" in resources + assert "gpu=L4:24GB:1" in resources + + @pytest.mark.asyncio + async def test_requires_actual_instance_resources(self, session: AsyncSession): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_name="qwen-endpoint", + status=RunStatus.RUNNING, + run_spec=RunSpec( + run_name="qwen-endpoint", + configuration=ServiceConfiguration.parse_obj( + { + "type": "service", + "name": "qwen-endpoint", + "commands": [ + "vllm serve Qwen/Qwen3-4B --host 0.0.0.0 --port 8000", + ], + "port": 8000, + "model": "Qwen/Qwen3-4B", + "resources": {"gpu": "16GB"}, + } + ), + ), + ) + await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + registered=True, + ) + await session.refresh(run, attribute_names=["jobs"]) + + with pytest.raises(ValueError, match="actual instance resources"): + build_endpoint_preset_from_run(run) + + @pytest.mark.asyncio + async def test_builds_replica_groups_in_service_order(self, session: AsyncSession, tmp_path): + project = await create_project(session=session) + user = await create_user(session=session) + repo = await create_repo(session=session, project_id=project.id) + run = await create_run( + session=session, + project=project, + repo=repo, + user=user, + run_name="qwen-endpoint", + status=RunStatus.RUNNING, + run_spec=RunSpec( + run_name="qwen-endpoint", + configuration=ServiceConfiguration.parse_obj( + { + "type": "service", + "name": "qwen-endpoint", + "port": 8000, + "model": "Qwen/Qwen3-4B", + "replicas": [ + { + "name": "router", + "count": 1, + "commands": ["python router.py"], + "resources": {"cpu": 4}, + }, + { + "name": "worker", + "count": 2, + "commands": [ + "vllm serve Qwen/Qwen3-4B --host 0.0.0.0 --port 8000", + ], + "resources": {"gpu": "24GB"}, + }, + ], + } + ), + ), + ) + await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + registered=True, + replica_num=0, + replica_group_name="router", + job_runtime_data=get_job_runtime_data(offer=_instance_offer(cpu_count=8, gpu_count=0)), + ) + await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + registered=True, + replica_num=1, + replica_group_name="worker", + job_runtime_data=get_job_runtime_data( + offer=_instance_offer(gpu_name="L4", gpu_memory_gib=24) + ), + ) + await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + registered=True, + replica_num=2, + replica_group_name="worker", + job_runtime_data=get_job_runtime_data( + offer=_instance_offer(gpu_name="L4", gpu_memory_gib=24) + ), + ) + await session.refresh(run, attribute_names=["jobs"]) + + preset = build_endpoint_preset_from_run(run) + saved = await LocalDirEndpointPresetService(tmp_path).save_preset(PROJECT_NAME, preset) + + recipe = saved.recipes[0] + assert [group.name for group in recipe.service.replica_groups] == [ + "router", + "worker", + ] + assert "cpu=4" in recipe.service.replica_groups[0].resources.pretty_format() + worker_resources = recipe.service.replica_groups[1].resources + assert worker_resources.gpu is not None + assert worker_resources.gpu.vendor is not None + assert worker_resources.gpu.vendor.value == "nvidia" + assert "gpu=nvidia:24GB" in worker_resources.pretty_format() + assert len(recipe.validations[0].replicas[0].resources) == 1 + assert len(recipe.validations[0].replicas[1].resources) == 2 + + +class TestBuildPresetServiceConfiguration: + def test_get_endpoint_serving_run_name(self): + assert get_endpoint_serving_run_name("qwen-endpoint") == "qwen-endpoint-serving" + + def test_get_endpoint_serving_run_name_keeps_long_name_to_avoid_truncation(self): + endpoint_name = "a" * 41 + + assert get_endpoint_serving_run_name(endpoint_name) == endpoint_name + + def test_endpoint_name_env_and_profile_are_applied(self): + preset = _qwen_preset() + recipe = preset.recipes[0] + endpoint_configuration = EndpointConfiguration( + name="endpoint-name", + model="Qwen/Qwen3-4B", + env=Env.parse_obj({"HF_TOKEN": "endpoint"}), + creation_policy=CreationPolicy.REUSE_OR_CREATE, + max_price=1.5, + ) + + service_configuration = build_preset_service_configuration( + endpoint_name="endpoint-name", + endpoint_configuration=endpoint_configuration, + recipe=recipe, + ) + + assert service_configuration.name == "endpoint-name-serving" + assert service_configuration.env.as_dict() == { + "HF_HOME": "/root/.cache/huggingface", + "HF_TOKEN": "endpoint", + } + assert service_configuration.creation_policy == CreationPolicy.REUSE_OR_CREATE + assert service_configuration.max_price == 1.5 + assert service_configuration.resources.gpu is not None + + def test_builds_repo_less_run_spec(self): + preset = _qwen_preset() + recipe = preset.recipes[0] + endpoint_configuration = EndpointConfiguration( + name="endpoint-name", + model="Qwen/Qwen3-4B", + ) + + run_spec = build_preset_run_spec( + endpoint_name="endpoint-name", + endpoint_configuration=endpoint_configuration, + recipe=recipe, + ) + + assert run_spec.run_name == "endpoint-name-serving" + assert run_spec.repo_id is None + assert run_spec.repo_data is None + assert run_spec.ssh_key_pub is None + assert run_spec.configuration.name == "endpoint-name-serving" + + @pytest.mark.asyncio + async def test_missing_dir_returns_no_presets(self, tmp_path): + presets = await LocalDirEndpointPresetService(tmp_path / "missing").list_presets( + PROJECT_NAME + ) + + assert presets == [] + + +class TestFindMatchingPreset: + @pytest.mark.asyncio + async def test_returns_first_preset_with_available_offers( + self, session: AsyncSession, tmp_path + ): + _write_qwen_preset(tmp_path) + user = await create_user( + session=session, + global_role=GlobalRole.USER, + ssh_public_key="ssh-rsa test", + ) + project = await create_project(session=session, owner=user) + endpoint_configuration = EndpointConfiguration( + name="qwen-endpoint", + model="Qwen/Qwen3-4B", + ) + + with patch( + "dstack._internal.server.services.endpoints.planning.runs_services.get_plan", + new=AsyncMock(return_value=_run_plan_with_offer(available=True)), + ) as get_plan_mock: + match = await find_matching_preset_plan( + session=session, + project=project, + user=user, + endpoint_name="qwen-endpoint", + endpoint_configuration=endpoint_configuration, + preset_service=LocalDirEndpointPresetService(tmp_path), + ) + + assert match is not None + assert match.preset.base == "Qwen/Qwen3-4B" + assert match.recipe.id == "vllm-t4" + get_plan_mock.assert_awaited_once() + assert get_plan_mock.await_args is not None + assert get_plan_mock.await_args.kwargs["run_spec"].run_name == "qwen-endpoint-serving" + + @pytest.mark.asyncio + async def test_exact_repo_request_skips_other_recipe_models( + self, session: AsyncSession, tmp_path + ): + presets_dir = _project_presets_dir(tmp_path) + presets_dir.mkdir(parents=True) + _write_qwen_preset_file( + presets_dir / "qwen-a.yml", + recipe_id="variant", + recipe_model="groxaxo/Qwen3-4B-GPTQ-4Bit", + ) + _write_qwen_preset_file(presets_dir / "qwen-b.yml", recipe_id="base") + user = await create_user( + session=session, + global_role=GlobalRole.USER, + ssh_public_key="ssh-rsa test", + ) + project = await create_project(session=session, owner=user) + + with patch( + "dstack._internal.server.services.endpoints.planning.runs_services.get_plan", + new=AsyncMock(return_value=_run_plan_with_offer(available=True)), + ) as get_plan_mock: + match = await find_matching_preset_plan( + session=session, + project=project, + user=user, + endpoint_name="qwen-endpoint", + endpoint_configuration=EndpointConfiguration( + name="qwen-endpoint", + model={"repo": "Qwen/Qwen3-4B"}, + ), + preset_service=LocalDirEndpointPresetService(tmp_path), + ) + + assert match is not None + assert match.recipe.id == "base" + get_plan_mock.assert_awaited_once() + + @pytest.mark.asyncio + async def test_tracks_first_unprovisionable_preset_when_offers_are_unavailable( + self, session: AsyncSession, tmp_path + ): + _write_qwen_preset(tmp_path) + user = await create_user( + session=session, + global_role=GlobalRole.USER, + ssh_public_key="ssh-rsa test", + ) + project = await create_project(session=session, owner=user) + + with patch( + "dstack._internal.server.services.endpoints.planning.runs_services.get_plan", + new=AsyncMock(return_value=_run_plan_with_offer(available=False)), + ): + result = await find_preset_planning_result( + session=session, + project=project, + user=user, + endpoint_name="qwen-endpoint", + endpoint_configuration=EndpointConfiguration( + name="qwen-endpoint", + model="Qwen/Qwen3-4B", + ), + preset_service=LocalDirEndpointPresetService(tmp_path), + ) + + assert result.provisionable is None + assert result.unprovisionable is not None + assert result.unprovisionable.preset.base == "Qwen/Qwen3-4B" + assert result.unprovisionable.recipe.id == "vllm-t4" + + @pytest.mark.asyncio + async def test_uses_first_later_recipe_with_available_offers( + self, session: AsyncSession, tmp_path + ): + presets_dir = _project_presets_dir(tmp_path) + presets_dir.mkdir(parents=True) + _write_qwen_preset_file(presets_dir / "qwen-a.yml", recipe_id="vllm-t4", gpu="16GB") + _write_qwen_preset_file(presets_dir / "qwen-b.yml", recipe_id="vllm-l4", gpu="24GB") + user = await create_user( + session=session, + global_role=GlobalRole.USER, + ssh_public_key="ssh-rsa test", + ) + project = await create_project(session=session, owner=user) + + with patch( + "dstack._internal.server.services.endpoints.planning.runs_services.get_plan", + new=AsyncMock( + side_effect=[ + _run_plan_with_offer(available=False), + _run_plan_with_offer(available=True), + ] + ), + ) as get_plan_mock: + result = await find_preset_planning_result( + session=session, + project=project, + user=user, + endpoint_name="qwen-endpoint", + endpoint_configuration=EndpointConfiguration( + name="qwen-endpoint", + model="Qwen/Qwen3-4B", + ), + preset_service=LocalDirEndpointPresetService(tmp_path), + ) + + assert result.provisionable is not None + assert result.provisionable.recipe.id == "vllm-l4" + assert result.unprovisionable is not None + assert result.unprovisionable.recipe.id == "vllm-t4" + assert get_plan_mock.await_count == 2 + + @pytest.mark.asyncio + async def test_stops_at_first_recipe_with_available_offers( + self, session: AsyncSession, tmp_path + ): + presets_dir = _project_presets_dir(tmp_path) + presets_dir.mkdir(parents=True) + _write_qwen_preset_file(presets_dir / "qwen-a.yml", recipe_id="vllm-t4", gpu="16GB") + _write_qwen_preset_file(presets_dir / "qwen-b.yml", recipe_id="vllm-l4", gpu="24GB") + user = await create_user( + session=session, + global_role=GlobalRole.USER, + ssh_public_key="ssh-rsa test", + ) + project = await create_project(session=session, owner=user) + + with patch( + "dstack._internal.server.services.endpoints.planning.runs_services.get_plan", + new=AsyncMock(return_value=_run_plan_with_offer(available=True)), + ) as get_plan_mock: + match = await find_matching_preset_plan( + session=session, + project=project, + user=user, + endpoint_name="qwen-endpoint", + endpoint_configuration=EndpointConfiguration( + name="qwen-endpoint", + model="Qwen/Qwen3-4B", + ), + preset_service=LocalDirEndpointPresetService(tmp_path), + ) + + assert match is not None + assert match.recipe.id == "vllm-t4" + get_plan_mock.assert_awaited_once() + + @pytest.mark.asyncio + async def test_matching_plan_requires_available_offers(self, session: AsyncSession, tmp_path): + _write_qwen_preset(tmp_path) + user = await create_user( + session=session, + global_role=GlobalRole.USER, + ssh_public_key="ssh-rsa test", + ) + project = await create_project(session=session, owner=user) + + with patch( + "dstack._internal.server.services.endpoints.planning.runs_services.get_plan", + new=AsyncMock(return_value=_run_plan_with_offer(available=False)), + ): + match = await find_matching_preset_plan( + session=session, + project=project, + user=user, + endpoint_name="qwen-endpoint", + endpoint_configuration=EndpointConfiguration( + name="qwen-endpoint", + model="Qwen/Qwen3-4B", + ), + preset_service=LocalDirEndpointPresetService(tmp_path), + ) + + assert match is None + + @pytest.mark.asyncio + async def test_skips_preset_with_unresolved_env_sentinel( + self, session: AsyncSession, tmp_path + ): + presets_dir = _project_presets_dir(tmp_path) + presets_dir.mkdir(parents=True) + (presets_dir / "qwen.yml").write_text( + """\ +type: endpoint-preset +model: Qwen/Qwen3-4B +recipes: + - id: vllm-t4 + service: + env: + - HF_TOKEN + commands: + - vllm serve Qwen/Qwen3-4B --host 0.0.0.0 --port 8000 + port: 8000 + resources: + gpu: 16GB + validations: + - replicas: + - resources: + - cpu: 4 + memory: 16GB + disk: 100GB + gpu: + name: T4 + memory: 16GB + count: 1 +""" + ) + user = await create_user( + session=session, + global_role=GlobalRole.USER, + ssh_public_key="ssh-rsa test", + ) + project = await create_project(session=session, owner=user) + + with patch( + "dstack._internal.server.services.endpoints.planning.runs_services.get_plan", + new=AsyncMock(), + ) as get_plan_mock: + result = await find_preset_planning_result( + session=session, + project=project, + user=user, + endpoint_name="qwen-endpoint", + endpoint_configuration=EndpointConfiguration( + name="qwen-endpoint", + model="Qwen/Qwen3-4B", + ), + preset_service=LocalDirEndpointPresetService(tmp_path), + ) + + assert result.provisionable is None + assert result.unprovisionable is None + get_plan_mock.assert_not_awaited() + + @pytest.mark.asyncio + async def test_refreshes_user_ssh_key_before_planning(self, session: AsyncSession, tmp_path): + _write_qwen_preset(tmp_path) + user = await create_user(session=session, global_role=GlobalRole.USER) + project = await create_project(session=session, owner=user) + user.ssh_public_key = None + await session.commit() + + with ( + patch( + "dstack._internal.server.services.endpoints.planning.users_services.refresh_ssh_key", + new=AsyncMock(return_value=user), + ) as refresh_mock, + patch( + "dstack._internal.server.services.endpoints.planning.runs_services.get_plan", + new=AsyncMock(return_value=_run_plan_with_offer(available=True)), + ), + ): + match = await find_matching_preset_plan( + session=session, + project=project, + user=user, + endpoint_name="qwen-endpoint", + endpoint_configuration=EndpointConfiguration( + name="qwen-endpoint", + model="Qwen/Qwen3-4B", + ), + preset_service=LocalDirEndpointPresetService(tmp_path), + ) + + assert match is not None + refresh_mock.assert_awaited_once() + + +def _write_qwen_preset(tmp_path): + presets_dir = _project_presets_dir(tmp_path) + presets_dir.mkdir(parents=True) + _write_qwen_preset_file(presets_dir / "qwen.yml") + + +def _write_qwen_preset_file( + path, + recipe_id: str = "vllm-t4", + gpu: str = "16GB", + recipe_model: str | None = None, +): + recipe_model_line = f" model: {recipe_model}\n" if recipe_model is not None else "" + path.write_text( + """\ +type: endpoint-preset +model: Qwen/Qwen3-4B +recipes: + - id: {recipe_id} +{recipe_model_line}\ + service: + commands: + - vllm serve Qwen/Qwen3-4B --host 0.0.0.0 --port 8000 + port: 8000 + resources: + gpu: {gpu} + validations: + - replicas: + - resources: + - cpu: 4 + memory: 16GB + disk: 100GB + gpu: + name: T4 + memory: 16GB + count: 1 +""".format(recipe_id=recipe_id, gpu=gpu, recipe_model_line=recipe_model_line) + ) + + +def _project_presets_dir(projects_dir): + return projects_dir / PROJECT_NAME / "presets" + + +def _instance_offer( + gpu_name: str = "T4", + gpu_memory_gib: float = 16, + gpu_count: int = 1, + cpu_count: int = 4, + memory_gib: float = 16, + disk_gib: float = 100, +) -> InstanceOfferWithAvailability: + return InstanceOfferWithAvailability( + backend=BackendType.AWS, + instance=InstanceType( + name="test-instance", + resources=Resources( + cpus=cpu_count, + memory_mib=int(memory_gib * 1024), + gpus=[ + Gpu( + name=gpu_name, + memory_mib=int(gpu_memory_gib * 1024), + ) + for _ in range(gpu_count) + ], + spot=False, + disk=Disk(size_mib=int(disk_gib * 1024)), + ), + ), + region="us-east-1", + price=1.0, + availability=InstanceAvailability.AVAILABLE, + instance_runtime=InstanceRuntime.SHIM, + ) + + +def _qwen_preset( + recipe_id: str = "vllm-t4", + resources: dict | None = None, + recipe_model: str = "Qwen/Qwen3-4B", +) -> EndpointPreset: + service = ServiceConfiguration.parse_obj( + { + "type": "service", + "name": "preset-service-name", + "commands": [ + "vllm serve Qwen/Qwen3-4B --host 0.0.0.0 --port 8000", + ], + "port": 8000, + "model": "Qwen/Qwen3-4B", + "env": { + "HF_HOME": "/root/.cache/huggingface", + "HF_TOKEN": "SECRET_VALUE", + }, + "resources": resources or {"gpu": "16GB"}, + "creation_policy": "reuse", + } + ) + return EndpointPreset( + base="Qwen/Qwen3-4B", + recipes=[ + EndpointPresetRecipe( + id=recipe_id, + model=recipe_model, + service=service, + validations=[ + EndpointPresetValidation( + replicas=[ + EndpointPresetValidationReplica.parse_obj(_t4_validation_replica()) + ] + ) + ], + ) + ], + ) + + +def _t4_validation_replica() -> dict: + return { + "resources": [ + { + "cpu": 4, + "memory": "16GB", + "disk": "100GB", + "gpu": {"name": "T4", "memory": "16GB", "count": 1}, + } + ], + } + + +def _run_plan_with_offer(available: bool): + availability = Mock() + availability.is_available.return_value = available + offer = Mock(availability=availability) + job_plan = Mock(offers=[offer]) + job_plan.job_spec.requirements.resources = Mock() + job_plan.job_spec.requirements.spot = None + job_plan.job_spec.requirements.max_price = None + return Mock(job_plans=[job_plan])