dstack 0.0.9__py3-none-any.whl → 0.20.7__py3-none-any.whl

dstack/api/server/__init__.py

@@ -0,0 +1,206 @@
+import hashlib
+import os
+import pprint
+import time
+from typing import Dict, List, Optional, Type
+
+import requests
+
+from dstack import version
+from dstack._internal.core.errors import (
+    ClientError,
+    MethodNotAllowedError,
+    ServerClientError,
+    URLNotFoundError,
+)
+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._events import EventsAPIClient
+from dstack.api.server._files import FilesAPIClient
+from dstack.api.server._fleets import FleetsAPIClient
+from dstack.api.server._gateways import GatewaysAPIClient
+from dstack.api.server._gpus import GpusAPIClient
+from dstack.api.server._logs import LogsAPIClient
+from dstack.api.server._metrics import MetricsAPIClient
+from dstack.api.server._projects import ProjectsAPIClient
+from dstack.api.server._repos import ReposAPIClient
+from dstack.api.server._runs import RunsAPIClient
+from dstack.api.server._secrets import SecretsAPIClient
+from dstack.api.server._users import UsersAPIClient
+from dstack.api.server._volumes import VolumesAPIClient
+
+_MAX_RETRIES = 3
+_RETRY_INTERVAL = 1
+
+
+class APIClient:
+    """
+    Low-level API client for interacting with the `dstack` server.
+    Supports all HTTP API endpoints.
+
+    Attributes:
+        users: operations with users
+        projects: operations with projects
+        backends: operations with backends
+        fleets: operations with fleets
+        runs: operations with runs
+        gpus: operations with GPUs
+        metrics: operations with metrics
+        logs: operations with logs
+        gateways: operations with gateways
+        volumes: operations with volumes
+        files: operations with files
+    """
+
+    def __init__(self, base_url: str, token: Optional[str] = None):
+        """
+        Args:
+            base_url: The API endpoints prefix, e.g. `http://127.0.0.1:3000/`.
+            token: The API token.
+        """
+        self._base_url = base_url.rstrip("/")
+        self._s = requests.session()
+        self._token = None
+        if token is not None:
+            self._token = token
+            self._s.headers.update({"Authorization": f"Bearer {token}"})
+        client_api_version = os.getenv("DSTACK_CLIENT_API_VERSION", version.__version__)
+        if client_api_version is not None:
+            self._s.headers.update({"X-API-VERSION": client_api_version})
+        self._logger = get_logger(__name__)
+
+    @property
+    def base_url(self) -> str:
+        return self._base_url
+
+    @property
+    def auth(self) -> AuthAPIClient:
+        return AuthAPIClient(self._request, self._logger)
+
+    @property
+    def users(self) -> UsersAPIClient:
+        return UsersAPIClient(self._request, self._logger)
+
+    @property
+    def projects(self) -> ProjectsAPIClient:
+        return ProjectsAPIClient(self._request, self._logger)
+
+    @property
+    def backends(self) -> BackendsAPIClient:
+        return BackendsAPIClient(self._request, self._logger)
+
+    @property
+    def fleets(self) -> FleetsAPIClient:
+        return FleetsAPIClient(self._request, self._logger)
+
+    @property
+    def repos(self) -> ReposAPIClient:
+        return ReposAPIClient(self._request, self._logger)
+
+    @property
+    def runs(self) -> RunsAPIClient:
+        return RunsAPIClient(self._request, self._logger)
+
+    @property
+    def gpus(self) -> GpusAPIClient:
+        return GpusAPIClient(self._request, self._logger)
+
+    @property
+    def metrics(self) -> MetricsAPIClient:
+        return MetricsAPIClient(self._request, self._logger)
+
+    @property
+    def logs(self) -> LogsAPIClient:
+        return LogsAPIClient(self._request, self._logger)
+
+    @property
+    def secrets(self) -> SecretsAPIClient:
+        return SecretsAPIClient(self._request, self._logger)
+
+    @property
+    def gateways(self) -> GatewaysAPIClient:
+        return GatewaysAPIClient(self._request, self._logger)
+
+    @property
+    def volumes(self) -> VolumesAPIClient:
+        return VolumesAPIClient(self._request, self._logger)
+
+    @property
+    def files(self) -> FilesAPIClient:
+        return FilesAPIClient(self._request, self._logger)
+
+    @property
+    def events(self) -> EventsAPIClient:
+        return EventsAPIClient(self._request, self._logger)
+
+    def get_token_hash(self) -> str:
+        if self._token is None:
+            raise ValueError("Token not set")
+        return hashlib.sha1(self._token.encode()).hexdigest()[:8]
+
+    def _request(
+        self,
+        path: str,
+        body: Optional[str] = None,
+        raise_for_status: bool = True,
+        method: str = "POST",
+        **kwargs,
+    ) -> requests.Response:
+        path = path.lstrip("/")
+        if body is not None:
+            kwargs.setdefault("headers", {})["Content-Type"] = "application/json"
+            kwargs["data"] = body
+
+        self._logger.debug("POST /%s", path)
+        for _ in range(_MAX_RETRIES):
+            try:
+                # TODO: set adequate timeout here or everywhere the method is used
+                resp = self._s.request(method, f"{self._base_url}/{path}", **kwargs)
+                break
+            except requests.exceptions.ConnectionError as e:
+                self._logger.debug("Could not connect to server: %s", e)
+                time.sleep(_RETRY_INTERVAL)
+        else:
+            raise ClientError(f"Failed to connect to dstack server {self._base_url}")
+
+        if 400 <= resp.status_code < 600:
+            self._logger.debug(
+                "Error requesting %s. Status: %s. Headers: %s. Body: %s",
+                resp.request.url,
+                resp.status_code,
+                resp.headers,
+                resp.content,
+            )
+
+        if raise_for_status:
+            if resp.status_code == 400:  # raise ServerClientError
+                detail: List[Dict] = resp.json()["detail"]
+                if len(detail) == 1 and detail[0]["code"] in _server_client_errors:
+                    kwargs = detail[0]
+                    code = kwargs.pop("code")
+                    raise _server_client_errors[code](**kwargs)
+            if resp.status_code == 422:
+                formatted_error = pprint.pformat(resp.json())
+                raise ClientError(f"Server validation error: \n{formatted_error}")
+            if resp.status_code == 403:
+                raise ClientError(
+                    f"Access to {resp.request.url} is denied. Please check your access token"
+                )
+            if resp.status_code == 404:
+                raise URLNotFoundError(f"Status code 404 when requesting {resp.request.url}")
+            if resp.status_code == 405:
+                raise MethodNotAllowedError(f"Status code 405 when requesting {resp.request.url}")
+            if 400 <= resp.status_code < 600:
+                raise ClientError(
+                    f"Unexpected error: status code {resp.status_code}"
+                    f" when requesting {resp.request.url}."
+                    " Check the server logs for backend issues, and the CLI logs at (~/.dstack/logs/cli/latest.log) local CLI output"
+                )
+        return resp
+
+
+_server_client_errors: Dict[str, Type[ServerClientError]] = {
+    cls.code: cls for cls in ServerClientError.__subclasses__()
+}
+_server_client_errors[ServerClientError.code] = ServerClientError

dstack/api/server/_auth.py

@@ -0,0 +1,30 @@
+from typing import Optional
+
+from pydantic import parse_obj_as
+
+from dstack._internal.core.models.auth import OAuthProviderInfo
+from dstack._internal.core.models.users import UserWithCreds
+from dstack._internal.server.schemas.auth import (
+    OAuthAuthorizeRequest,
+    OAuthAuthorizeResponse,
+    OAuthCallbackRequest,
+)
+from dstack.api.server._group import APIClientGroup
+
+
+class AuthAPIClient(APIClientGroup):
+    def list_providers(self) -> list[OAuthProviderInfo]:
+        resp = self._request("/api/auth/list_providers")
+        return parse_obj_as(list[OAuthProviderInfo.__response__], resp.json())
+
+    def authorize(self, provider: str, local_port: Optional[int] = None) -> OAuthAuthorizeResponse:
+        body = OAuthAuthorizeRequest(local_port=local_port)
+        resp = self._request(f"/api/auth/{provider}/authorize", body=body.json())
+        return parse_obj_as(OAuthAuthorizeResponse.__response__, resp.json())
+
+    def callback(
+        self, provider: str, code: str, state: str, base_url: Optional[str] = None
+    ) -> UserWithCreds:
+        body = OAuthCallbackRequest(code=code, state=state, base_url=base_url)
+        resp = self._request(f"/api/auth/{provider}/callback", body=body.json())
+        return parse_obj_as(UserWithCreds.__response__, resp.json())

dstack/api/server/_backends.py

@@ -0,0 +1,38 @@
+from typing import List
+
+from pydantic import parse_obj_as
+
+from dstack._internal.core.backends.models import (
+    AnyBackendConfigWithCreds,
+)
+from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.server.schemas.backends import DeleteBackendsRequest
+from dstack.api.server._group import APIClientGroup
+
+
+class BackendsAPIClient(APIClientGroup):
+    def list_backend_types(self) -> List[BackendType]:
+        resp = self._request("/api/backends/list_types")
+        return parse_obj_as(List[BackendType], resp.json())
+
+    def create(
+        self, project_name: str, config: AnyBackendConfigWithCreds
+    ) -> AnyBackendConfigWithCreds:
+        resp = self._request(f"/api/project/{project_name}/backends/create", body=config.json())
+        return parse_obj_as(AnyBackendConfigWithCreds, resp.json())
+
+    def update(
+        self, project_name: str, config: AnyBackendConfigWithCreds
+    ) -> AnyBackendConfigWithCreds:
+        resp = self._request(f"/api/project/{project_name}/backends/update", body=config.json())
+        return parse_obj_as(AnyBackendConfigWithCreds, resp.json())
+
+    def delete(self, project_name: str, backends_names: List[BackendType]):
+        body = DeleteBackendsRequest(backends_names=backends_names)
+        self._request(f"/api/project/{project_name}/backends/delete", body=body.json())
+
+    def config_info(
+        self, project_name: str, backend_name: BackendType
+    ) -> AnyBackendConfigWithCreds:
+        resp = self._request(f"/api/project/{project_name}/backends/{backend_name}/config_info")
+        return parse_obj_as(AnyBackendConfigWithCreds, resp.json())

dstack/api/server/_events.py

@@ -0,0 +1,64 @@
+from datetime import datetime, timezone
+from typing import Optional
+from uuid import UUID
+
+from pydantic import parse_obj_as
+
+from dstack._internal.core.compatibility.events import get_list_events_excludes
+from dstack._internal.core.models.events import Event, EventTargetType
+from dstack._internal.server.schemas.events import LIST_EVENTS_DEFAULT_LIMIT, ListEventsRequest
+from dstack.api.server._group import APIClientGroup
+
+
+class EventsAPIClient(APIClientGroup):
+    def list(
+        self,
+        target_projects: Optional[list[UUID]] = None,
+        target_users: Optional[list[UUID]] = None,
+        target_fleets: Optional[list[UUID]] = None,
+        target_instances: Optional[list[UUID]] = None,
+        target_runs: Optional[list[UUID]] = None,
+        target_jobs: Optional[list[UUID]] = None,
+        within_projects: Optional[list[UUID]] = None,
+        within_fleets: Optional[list[UUID]] = None,
+        within_runs: Optional[list[UUID]] = None,
+        include_target_types: Optional[list[EventTargetType]] = None,
+        actors: Optional[list[Optional[UUID]]] = None,
+        prev_recorded_at: Optional[datetime] = None,
+        prev_id: Optional[UUID] = None,
+        limit: int = LIST_EVENTS_DEFAULT_LIMIT,
+        ascending: bool = False,
+        *,
+        # NOTE: New parameters go here. Avoid positional parameters, they can break compatibility.
+        target_volumes: Optional[list[UUID]] = None,
+        target_gateways: Optional[list[UUID]] = None,
+        target_secrets: Optional[list[UUID]] = None,
+    ) -> list[Event]:
+        if prev_recorded_at is not None:
+            # Time zones other than UTC are misinterpreted by the server:
+            # https://github.com/dstackai/dstack/issues/3354
+            prev_recorded_at = prev_recorded_at.astimezone(timezone.utc)
+        req = ListEventsRequest(
+            target_projects=target_projects,
+            target_users=target_users,
+            target_fleets=target_fleets,
+            target_instances=target_instances,
+            target_runs=target_runs,
+            target_jobs=target_jobs,
+            target_volumes=target_volumes,
+            target_gateways=target_gateways,
+            target_secrets=target_secrets,
+            within_projects=within_projects,
+            within_fleets=within_fleets,
+            within_runs=within_runs,
+            include_target_types=include_target_types,
+            actors=actors,
+            prev_recorded_at=prev_recorded_at,
+            prev_id=prev_id,
+            limit=limit,
+            ascending=ascending,
+        )
+        resp = self._request(
+            "/api/events/list", body=req.json(exclude=get_list_events_excludes(req))
+        )
+        return parse_obj_as(list[Event.__response__], resp.json())

dstack/api/server/_files.py

@@ -0,0 +1,18 @@
+from typing import BinaryIO
+
+from pydantic import parse_obj_as
+
+from dstack._internal.core.models.files import FileArchive
+from dstack._internal.server.schemas.files import GetFileArchiveByHashRequest
+from dstack.api.server._group import APIClientGroup
+
+
+class FilesAPIClient(APIClientGroup):
+    def get_archive_by_hash(self, hash: str) -> FileArchive:
+        body = GetFileArchiveByHashRequest(hash=hash)
+        resp = self._request("/api/files/get_archive_by_hash", body=body.json())
+        return parse_obj_as(FileArchive.__response__, resp.json())
+
+    def upload_archive(self, hash: str, fp: BinaryIO) -> FileArchive:
+        resp = self._request("/api/files/upload_archive", files={"file": (hash, fp)})
+        return parse_obj_as(FileArchive.__response__, resp.json())

dstack/api/server/_fleets.py

@@ -0,0 +1,82 @@
+from typing import List, Optional, Union
+from uuid import UUID
+
+from pydantic import parse_obj_as
+
+from dstack._internal.core.compatibility.fleets import (
+    get_apply_plan_excludes,
+    get_create_fleet_excludes,
+    get_get_plan_excludes,
+)
+from dstack._internal.core.models.fleets import ApplyFleetPlanInput, Fleet, FleetPlan, FleetSpec
+from dstack._internal.server.schemas.fleets import (
+    ApplyFleetPlanRequest,
+    CreateFleetRequest,
+    DeleteFleetInstancesRequest,
+    DeleteFleetsRequest,
+    GetFleetPlanRequest,
+    GetFleetRequest,
+)
+from dstack.api.server._group import APIClientGroup
+
+
+class FleetsAPIClient(APIClientGroup):
+    def list(self, project_name: str) -> List[Fleet]:
+        resp = self._request(f"/api/project/{project_name}/fleets/list")
+        return parse_obj_as(List[Fleet.__response__], resp.json())
+
+    def get(
+        self, project_name: str, name: Optional[str] = None, fleet_id: Optional[UUID] = None
+    ) -> Fleet:
+        if name is None and fleet_id is None:
+            raise ValueError("Either name or fleet_id must be provided")
+        if name is not None and fleet_id is not None:
+            raise ValueError("Cannot specify both name and fleet_id")
+        body = GetFleetRequest(name=name, id=fleet_id)
+        resp = self._request(
+            f"/api/project/{project_name}/fleets/get",
+            body=body.json(),
+        )
+        return parse_obj_as(Fleet.__response__, resp.json())
+
+    def get_plan(
+        self,
+        project_name: str,
+        spec: FleetSpec,
+    ) -> FleetPlan:
+        body = GetFleetPlanRequest(spec=spec)
+        body_json = body.json(exclude=get_get_plan_excludes(spec))
+        resp = self._request(f"/api/project/{project_name}/fleets/get_plan", body=body_json)
+        return parse_obj_as(FleetPlan.__response__, resp.json())
+
+    def apply_plan(
+        self,
+        project_name: str,
+        plan: Union[FleetPlan, ApplyFleetPlanInput],
+        force: bool = False,
+    ) -> Fleet:
+        plan_input = ApplyFleetPlanInput.__response__.parse_obj(plan)
+        body = ApplyFleetPlanRequest(plan=plan_input, force=force)
+        body_json = body.json(exclude=get_apply_plan_excludes(plan_input))
+        resp = self._request(f"/api/project/{project_name}/fleets/apply", body=body_json)
+        return parse_obj_as(Fleet.__response__, resp.json())
+
+    def delete(self, project_name: str, names: List[str]) -> None:
+        body = DeleteFleetsRequest(names=names)
+        self._request(f"/api/project/{project_name}/fleets/delete", body=body.json())
+
+    def delete_instances(self, project_name: str, name: str, instance_nums: List[int]) -> None:
+        body = DeleteFleetInstancesRequest(name=name, instance_nums=instance_nums)
+        self._request(f"/api/project/{project_name}/fleets/delete_instances", body=body.json())
+
+    # Deprecated
+    # TODO: Remove in 0.21
+    def create(
+        self,
+        project_name: str,
+        spec: FleetSpec,
+    ) -> Fleet:
+        body = CreateFleetRequest(spec=spec)
+        body_json = body.json(exclude=get_create_fleet_excludes(spec))
+        resp = self._request(f"/api/project/{project_name}/fleets/create", body=body_json)
+        return parse_obj_as(Fleet.__response__, resp.json())

dstack/api/server/_gateways.py

@@ -0,0 +1,54 @@
+from typing import List
+
+from pydantic import parse_obj_as
+
+from dstack._internal.core.compatibility.gateways import get_create_gateway_excludes
+from dstack._internal.core.models.gateways import Gateway, GatewayConfiguration
+from dstack._internal.server.schemas.gateways import (
+    CreateGatewayRequest,
+    DeleteGatewaysRequest,
+    GetGatewayRequest,
+    SetDefaultGatewayRequest,
+    SetWildcardDomainRequest,
+)
+from dstack.api.server._group import APIClientGroup
+
+
+class GatewaysAPIClient(APIClientGroup):
+    def list(self, project_name: str) -> List[Gateway]:
+        resp = self._request(f"/api/project/{project_name}/gateways/list")
+        return parse_obj_as(List[Gateway.__response__], resp.json())
+
+    def get(self, project_name: str, gateway_name: str) -> Gateway:
+        body = GetGatewayRequest(name=gateway_name)
+        resp = self._request(f"/api/project/{project_name}/gateways/get", body=body.json())
+        return parse_obj_as(Gateway.__response__, resp.json())
+
+    def create(
+        self,
+        project_name: str,
+        configuration: GatewayConfiguration,
+    ) -> Gateway:
+        body = CreateGatewayRequest(configuration=configuration)
+        resp = self._request(
+            f"/api/project/{project_name}/gateways/create",
+            body=body.json(exclude=get_create_gateway_excludes(configuration)),
+        )
+        return parse_obj_as(Gateway.__response__, resp.json())
+
+    def delete(self, project_name: str, gateways_names: List[str]) -> None:
+        body = DeleteGatewaysRequest(names=gateways_names)
+        self._request(f"/api/project/{project_name}/gateways/delete", body=body.json())
+
+    def set_default(self, project_name: str, gateway_name: str) -> None:
+        body = SetDefaultGatewayRequest(name=gateway_name)
+        self._request(f"/api/project/{project_name}/gateways/set_default", body=body.json())
+
+    def set_wildcard_domain(
+        self, project_name: str, gateway_name: str, wildcard_domain: str
+    ) -> Gateway:
+        body = SetWildcardDomainRequest(name=gateway_name, wildcard_domain=wildcard_domain)
+        resp = self._request(
+            f"/api/project/{project_name}/gateways/set_wildcard_domain", body=body.json()
+        )
+        return parse_obj_as(Gateway.__response__, resp.json())

dstack/api/server/_gpus.py

@@ -0,0 +1,27 @@
+from typing import List, Literal, Optional, cast
+
+from pydantic import parse_obj_as
+
+from dstack._internal.core.compatibility.gpus import get_list_gpus_excludes
+from dstack._internal.core.models.gpus import GpuGroup
+from dstack._internal.core.models.runs import RunSpec
+from dstack._internal.server.schemas.gpus import ListGpusRequest, ListGpusResponse
+from dstack.api.server._group import APIClientGroup
+
+
+class GpusAPIClient(APIClientGroup):
+    def list_gpus(
+        self,
+        project_name: str,
+        run_spec: RunSpec,
+        group_by: Optional[List[str]] = None,
+    ) -> List[GpuGroup]:
+        body = ListGpusRequest(
+            run_spec=run_spec,
+            group_by=cast(Optional[List[Literal["backend", "region", "count"]]], group_by),
+        )
+        resp = self._request(
+            f"/api/project/{project_name}/gpus/list",
+            body=body.json(exclude=get_list_gpus_excludes(body)),
+        )
+        return parse_obj_as(ListGpusResponse, resp.json()).gpus

dstack/api/server/_group.py

@@ -0,0 +1,22 @@
+from logging import Logger
+from typing import Optional
+
+import requests
+from typing_extensions import Protocol
+
+
+class APIRequest(Protocol):
+    def __call__(
+        self,
+        path: str,
+        body: Optional[str] = None,
+        raise_for_status: bool = True,
+        method: str = "POST",
+        **kwargs,
+    ) -> requests.Response: ...
+
+
+class APIClientGroup:
+    def __init__(self, _request: APIRequest, _logger: Logger):
+        self._request = _request
+        self._logger = _logger

dstack/api/server/_logs.py

@@ -0,0 +1,15 @@
+from pydantic import parse_obj_as
+
+from dstack._internal.core.compatibility.logs import get_poll_logs_excludes
+from dstack._internal.core.models.logs import JobSubmissionLogs
+from dstack._internal.server.schemas.logs import PollLogsRequest
+from dstack.api.server._group import APIClientGroup
+
+
+class LogsAPIClient(APIClientGroup):
+    def poll(self, project_name: str, body: PollLogsRequest) -> JobSubmissionLogs:
+        resp = self._request(
+            f"/api/project/{project_name}/logs/poll",
+            body=body.json(exclude=get_poll_logs_excludes(body)),
+        )
+        return parse_obj_as(JobSubmissionLogs.__response__, resp.json())

dstack/api/server/_metrics.py

@@ -0,0 +1,23 @@
+from pydantic import parse_obj_as
+
+from dstack._internal.core.models.metrics import JobMetrics
+from dstack.api.server._group import APIClientGroup
+
+
+class MetricsAPIClient(APIClientGroup):
+    def get_job_metrics(
+        self,
+        project_name: str,
+        run_name: str,
+        replica_num: int = 0,
+        job_num: int = 0,
+    ) -> JobMetrics:
+        resp = self._request(
+            f"/api/project/{project_name}/metrics/job/{run_name}",
+            method="GET",
+            params={
+                "replica_num": replica_num,
+                "job_num": job_num,
+            },
+        )
+        return parse_obj_as(JobMetrics.__response__, resp.json())