dstack 0.19.20__py3-none-any.whl → 0.19.22__py3-none-any.whl

dstack/_internal/server/services/logs/aws.py

@@ -1,5 +1,7 @@
 import itertools
 import operator
+import urllib
+import urllib.parse
 from contextlib import contextmanager
 from datetime import datetime, timedelta, timezone
 from typing import Iterator, List, Optional, Set, Tuple, TypedDict
@@ -64,6 +66,7 @@ class CloudWatchLogStorage(LogStorage):
             self._client = session.client("logs")
             self._check_group_exists(group)
         self._group = group
+        self._region = self._client.meta.region_name
         # Stores names of already created streams.
         # XXX: This set acts as an unbound cache. If this becomes a problem (in case of _very_ long
         # running server and/or lots of jobs, consider replacing it with an LRU cache, e.g.,
@@ -103,7 +106,11 @@ class CloudWatchLogStorage(LogStorage):
             )
             for cw_event in cw_events
         ]
-        return JobSubmissionLogs(logs=logs, next_token=next_token)
+        return JobSubmissionLogs(
+            logs=logs,
+            external_url=self._get_stream_external_url(stream),
+            next_token=next_token,
+        )
 
     def _get_log_events_with_retry(
         self, stream: str, request: PollLogsRequest
@@ -181,6 +188,11 @@ class CloudWatchLogStorage(LogStorage):
 
         return events, next_token
 
+    def _get_stream_external_url(self, stream: str) -> str:
+        quoted_group = urllib.parse.quote(self._group, safe="")
+        quoted_stream = urllib.parse.quote(stream, safe="")
+        return f"https://console.aws.amazon.com/cloudwatch/home?region={self._region}#logsV2:log-groups/log-group/{quoted_group}/log-events/{quoted_stream}"
+
     def write_logs(
         self,
         project: ProjectModel,

dstack/_internal/server/services/logs/gcp.py

@@ -1,3 +1,4 @@
+import urllib.parse
 from typing import List
 from uuid import UUID
 
@@ -48,6 +49,7 @@ class GCPLogStorage(LogStorage):
     # (https://cloud.google.com/logging/docs/analyze/custom-index).
 
     def __init__(self, project_id: str):
+        self.project_id = project_id
         try:
             self.client = logging_v2.Client(project=project_id)
             self.logger = self.client.logger(name=self.LOG_NAME)
@@ -106,7 +108,11 @@ class GCPLogStorage(LogStorage):
                 "GCP Logging read request limit exceeded."
                 " It's recommended to increase default entries.list request quota from 60 per minute."
             )
-        return JobSubmissionLogs(logs=logs, next_token=next_token if len(logs) > 0 else None)
+        return JobSubmissionLogs(
+            logs=logs,
+            external_url=self._get_stream_extrnal_url(stream_name),
+            next_token=next_token if len(logs) > 0 else None,
+        )
 
     def write_logs(
         self,
@@ -162,3 +168,12 @@ class GCPLogStorage(LogStorage):
         self, project_name: str, run_name: str, job_submission_id: UUID, producer: LogProducer
     ) -> str:
         return f"{project_name}-{run_name}-{job_submission_id}-{producer.value}"
+
+    def _get_stream_extrnal_url(self, stream_name: str) -> str:
+        log_name_resource_name = self._get_log_name_resource_name()
+        query = f'logName="{log_name_resource_name}" AND labels.stream="{stream_name}"'
+        quoted_query = urllib.parse.quote(query, safe="")
+        return f"https://console.cloud.google.com/logs/query;query={quoted_query}?project={self.project_id}"
+
+    def _get_log_name_resource_name(self) -> str:
+        return f"projects/{self.project_id}/logs/{self.LOG_NAME}"

dstack/_internal/server/services/offers.py

@@ -2,13 +2,13 @@ from typing import List, Literal, Optional, Tuple, Union
 
 import gpuhunt
 
-from dstack._internal.core.backends import (
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.base.compute import ComputeWithPlacementGroupSupport
+from dstack._internal.core.backends.features import (
     BACKENDS_WITH_CREATE_INSTANCE_SUPPORT,
     BACKENDS_WITH_MULTINODE_SUPPORT,
     BACKENDS_WITH_RESERVATION_SUPPORT,
 )
-from dstack._internal.core.backends.base.backend import Backend
-from dstack._internal.core.backends.base.compute import ComputeWithPlacementGroupSupport
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.instances import (
     InstanceOfferWithAvailability,

dstack/_internal/server/services/probes.py

@@ -0,0 +1,6 @@
+from dstack._internal.core.models.runs import Probe
+from dstack._internal.server.models import ProbeModel
+
+
+def probe_model_to_probe(probe_model: ProbeModel) -> Probe:
+    return Probe(success_streak=probe_model.success_streak)

dstack/_internal/server/services/projects.py

@@ -1,11 +1,10 @@
 import uuid
-from datetime import timezone
 from typing import Awaitable, Callable, List, Optional, Tuple
 
 from sqlalchemy import delete, select, update
 from sqlalchemy import func as safunc
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload
+from sqlalchemy.orm import QueryableAttribute, joinedload, load_only
 
 from dstack._internal.core.backends.configurators import get_configurator
 from dstack._internal.core.backends.dstack.models import (
@@ -54,13 +53,12 @@ async def list_user_projects(
     user: UserModel,
 ) -> List[Project]:
     """
-    Returns projects where the user is a member.
+    Returns projects where the user is a member or all projects for global admins.
     """
-    if user.global_role == GlobalRole.ADMIN:
-        projects = await list_project_models(session=session)
-    else:
-        projects = await list_user_project_models(session=session, user=user)
-
+    projects = await list_user_project_models(
+        session=session,
+        user=user,
+    )
     projects = sorted(projects, key=lambda p: p.created_at)
     return [
         project_model_to_project(p, include_backends=False, include_members=False)
@@ -80,7 +78,7 @@ async def list_user_accessible_projects(
     if user.global_role == GlobalRole.ADMIN:
         projects = await list_project_models(session=session)
     else:
-        member_projects = await list_user_project_models(session=session, user=user)
+        member_projects = await list_member_project_models(session=session, user=user)
         public_projects = await list_public_non_member_project_models(session=session, user=user)
         projects = member_projects + public_projects
 
@@ -167,7 +165,7 @@ async def delete_projects(
     projects_names: List[str],
 ):
     if user.global_role != GlobalRole.ADMIN:
-        user_projects = await list_user_project_models(
+        user_projects = await list_member_project_models(
             session=session, user=user, include_members=True
         )
         user_project_names = [p.name for p in user_projects]
@@ -199,6 +197,10 @@ async def set_project_members(
     project: ProjectModel,
     members: List[MemberSetting],
 ):
+    usernames = {m.username for m in members}
+    if len(usernames) != len(members):
+        raise ServerClientError("Cannot add same user multiple times")
+
     project = await get_project_model_by_name_or_error(
         session=session,
         project_name=project.name,
@@ -247,6 +249,10 @@ async def add_project_members(
     members: List[MemberSetting],
 ):
     """Add multiple members to a project."""
+    usernames = {m.username for m in members}
+    if len(usernames) != len(members):
+        raise ServerClientError("Cannot add same user multiple times")
+
     project = await get_project_model_by_name_or_error(
         session=session,
         project_name=project.name,
@@ -261,7 +267,10 @@ async def add_project_members(
     )
 
     if not is_self_join_to_public:
-        if requesting_user_role not in [ProjectRole.ADMIN, ProjectRole.MANAGER]:
+        if user.global_role != GlobalRole.ADMIN and requesting_user_role not in [
+            ProjectRole.ADMIN,
+            ProjectRole.MANAGER,
+        ]:
             raise ForbiddenError("Access denied: insufficient permissions to add members")
 
         if user.global_role != GlobalRole.ADMIN and requesting_user_role == ProjectRole.MANAGER:
@@ -274,8 +283,6 @@ async def add_project_members(
         if members[0].project_role != ProjectRole.USER:
             raise ForbiddenError("Access denied: can only join public projects as user role")
 
-    usernames = [member.username for member in members]
-
     res = await session.execute(
         select(UserModel).where((UserModel.name.in_(usernames)) | (UserModel.email.in_(usernames)))
     )
@@ -339,9 +346,25 @@ async def clear_project_members(
 
 
 async def list_user_project_models(
+    session: AsyncSession,
+    user: UserModel,
+    only_names: bool = False,
+) -> List[ProjectModel]:
+    load_only_attrs = []
+    if only_names:
+        load_only_attrs += [ProjectModel.id, ProjectModel.name]
+    if user.global_role == GlobalRole.ADMIN:
+        return await list_project_models(session=session, load_only_attrs=load_only_attrs)
+    return await list_member_project_models(
+        session=session, user=user, load_only_attrs=load_only_attrs
+    )
+
+
+async def list_member_project_models(
     session: AsyncSession,
     user: UserModel,
     include_members: bool = False,
+    load_only_attrs: Optional[List[QueryableAttribute]] = None,
 ) -> List[ProjectModel]:
     """
     List project models for a user where they are a member.
@@ -349,6 +372,8 @@ async def list_user_project_models(
     options = []
     if include_members:
         options.append(joinedload(ProjectModel.members))
+    if load_only_attrs:
+        options.append(load_only(*load_only_attrs))
     res = await session.execute(
         select(ProjectModel)
         .where(
@@ -395,13 +420,20 @@ async def list_user_owned_project_models(
 
 async def list_project_models(
     session: AsyncSession,
+    load_only_attrs: Optional[List[QueryableAttribute]] = None,
 ) -> List[ProjectModel]:
+    options = []
+    if load_only_attrs:
+        options.append(load_only(*load_only_attrs))
     res = await session.execute(
-        select(ProjectModel).where(ProjectModel.deleted == False),
+        select(ProjectModel).where(ProjectModel.deleted == False).options(*options)
     )
     return list(res.scalars().all())
 
 
+# TODO: Do not load ProjectModel.backends and ProjectModel.members by default when getting project
+
+
 async def get_project_model_by_name(
     session: AsyncSession, project_name: str, ignore_case: bool = True
 ) -> Optional[ProjectModel]:
@@ -415,7 +447,6 @@ async def get_project_model_by_name(
         .where(*filters)
         .options(joinedload(ProjectModel.backends))
         .options(joinedload(ProjectModel.members))
-        .options(joinedload(ProjectModel.default_gateway))
     )
     return res.unique().scalar()
 
@@ -432,7 +463,6 @@ async def get_project_model_by_name_or_error(
         )
         .options(joinedload(ProjectModel.backends))
         .options(joinedload(ProjectModel.members))
-        .options(joinedload(ProjectModel.default_gateway))
     )
     return res.unique().scalar_one()
 
@@ -449,7 +479,6 @@ async def get_project_model_by_id_or_error(
         )
         .options(joinedload(ProjectModel.backends))
         .options(joinedload(ProjectModel.members))
-        .options(joinedload(ProjectModel.default_gateway))
     )
     return res.unique().scalar_one()
 
@@ -537,7 +566,7 @@ def project_model_to_project(
         project_id=project_model.id,
         project_name=project_model.name,
         owner=users.user_model_to_user(project_model.owner),
-        created_at=project_model.created_at.replace(tzinfo=timezone.utc),
+        created_at=project_model.created_at,
         backends=backends,
         members=members,
         is_public=project_model.is_public,
@@ -608,7 +637,10 @@ async def remove_project_members(
     )
 
     if not is_self_leave:
-        if requesting_user_role not in [ProjectRole.ADMIN, ProjectRole.MANAGER]:
+        if user.global_role != GlobalRole.ADMIN and requesting_user_role not in [
+            ProjectRole.ADMIN,
+            ProjectRole.MANAGER,
+        ]:
             raise ForbiddenError("Access denied: insufficient permissions to remove members")
 
     res = await session.execute(

dstack/_internal/server/services/prometheus/client_metrics.py

@@ -5,6 +5,9 @@ class RunMetrics:
     """Wrapper class for run-related Prometheus metrics."""
 
     def __init__(self):
+        # submit_to_provision_duration reflects real provisioning time
+        # but does not reflect how quickly provisioning processing works
+        # since it includes scheduling time, retrying, etc.
         self._submit_to_provision_duration = Histogram(
             "dstack_submit_to_provision_duration_seconds",
             "Time from when a run has been submitted and first job provisioning",

dstack/_internal/server/services/prometheus/custom_metrics.py

@@ -2,7 +2,6 @@ import itertools
 import json
 from collections import defaultdict
 from collections.abc import Generator, Iterable
-from datetime import timezone
 from typing import ClassVar
 from uuid import UUID
 
@@ -80,7 +79,7 @@ async def get_instance_metrics(session: AsyncSession) -> Iterable[Metric]:
             "dstack_backend": instance.backend.value if instance.backend is not None else "",
             "dstack_gpu": gpu,
         }
-        duration = (now - instance.created_at.replace(tzinfo=timezone.utc)).total_seconds()
+        duration = (now - instance.created_at).total_seconds()
         metrics.add_sample(_INSTANCE_DURATION, labels, duration)
         metrics.add_sample(_INSTANCE_PRICE, labels, instance.price or 0.0)
         metrics.add_sample(_INSTANCE_GPU_COUNT, labels, gpu_count)
@@ -167,7 +166,7 @@ async def get_job_metrics(session: AsyncSession) -> Iterable[Metric]:
             "dstack_backend": jpd.get_base_backend().value,
             "dstack_gpu": gpus[0].name if gpus else "",
         }
-        duration = (now - job.submitted_at.replace(tzinfo=timezone.utc)).total_seconds()
+        duration = (now - job.submitted_at).total_seconds()
         metrics.add_sample(_JOB_DURATION, labels, duration)
         metrics.add_sample(_JOB_PRICE, labels, price)
         metrics.add_sample(_JOB_GPU_COUNT, labels, len(gpus))

dstack/_internal/server/services/runner/client.py

@@ -1,7 +1,6 @@
 import uuid
-from dataclasses import dataclass
 from http import HTTPStatus
-from typing import BinaryIO, Dict, List, Optional, TypeVar, Union
+from typing import BinaryIO, Dict, List, Literal, Optional, TypeVar, Union, overload
 
 import packaging.version
 import requests
@@ -14,9 +13,11 @@ from dstack._internal.core.models.repos.remote import RemoteRepoCreds
 from dstack._internal.core.models.resources import Memory
 from dstack._internal.core.models.runs import ClusterInfo, Job, Run
 from dstack._internal.core.models.volumes import InstanceMountPoint, Volume, VolumeMountPoint
+from dstack._internal.server.schemas.instances import InstanceCheck
 from dstack._internal.server.schemas.runner import (
     GPUDevice,
     HealthcheckResponse,
+    InstanceHealthResponse,
     LegacyPullResponse,
     LegacyStopBody,
     LegacySubmitBody,
@@ -37,15 +38,6 @@ UPLOAD_CODE_REQUEST_TIMEOUT = 60
 logger = get_logger(__name__)
 
 
-@dataclass
-class HealthStatus:
-    healthy: bool
-    reason: str
-
-    def __str__(self) -> str:
-        return self.reason
-
-
 class RunnerClient:
     def __init__(
         self,
@@ -193,6 +185,9 @@ class ShimClient:
     # API v1 (a.k.a. Legacy API) — `/api/{submit,pull,stop}`
     _API_V2_MIN_SHIM_VERSION = (0, 18, 34)
 
+    # `/api/instance/health`
+    _INSTANCE_HEALTH_MIN_SHIM_VERSION = (0, 19, 22)
+
     _shim_version: Optional["_Version"]
     _api_version: int
     _negotiated: bool = False
@@ -212,11 +207,25 @@ class ShimClient:
             self._negotiate()
         return self._api_version == 2
 
-    def healthcheck(self, unmask_exeptions: bool = False) -> Optional[HealthcheckResponse]:
+    def is_instance_health_supported(self) -> bool:
+        if not self._negotiated:
+            self._negotiate()
+        return (
+            self._shim_version is None
+            or self._shim_version >= self._INSTANCE_HEALTH_MIN_SHIM_VERSION
+        )
+
+    @overload
+    def healthcheck(self) -> Optional[HealthcheckResponse]: ...
+
+    @overload
+    def healthcheck(self, unmask_exceptions: Literal[True]) -> HealthcheckResponse: ...
+
+    def healthcheck(self, unmask_exceptions: bool = False) -> Optional[HealthcheckResponse]:
         try:
             resp = self._request("GET", "/api/healthcheck", raise_for_status=True)
         except requests.exceptions.RequestException:
-            if unmask_exeptions:
+            if unmask_exceptions:
                 raise
             return None
         if not self._negotiated:
@@ -225,6 +234,17 @@ class ShimClient:
 
     # API v2 methods
 
+    def get_instance_health(self) -> Optional[InstanceHealthResponse]:
+        if not self.is_instance_health_supported():
+            logger.debug("instance health is not supported: %s", self._shim_version)
+            return None
+        resp = self._request("GET", "/api/instance/health")
+        if resp.status_code == HTTPStatus.NOT_FOUND:
+            logger.warning("instance health: %s", resp.text)
+            return None
+        self._raise_for_status(resp)
+        return self._response(InstanceHealthResponse, resp)
+
     def get_task(self, task_id: "_TaskID") -> TaskInfoResponse:
         if not self.is_api_v2_supported():
             raise ShimAPIVersionError()
@@ -418,14 +438,26 @@ class ShimClient:
         self._negotiated = True
 
 
-def health_response_to_health_status(data: HealthcheckResponse) -> HealthStatus:
-    if data.service == "dstack-shim":
-        return HealthStatus(healthy=True, reason="Service is OK")
-    else:
-        return HealthStatus(
-            healthy=False,
-            reason=f"Service name is {data.service}, service version: {data.version}",
+def healthcheck_response_to_instance_check(
+    response: HealthcheckResponse,
+    instance_health_response: Optional[InstanceHealthResponse] = None,
+) -> InstanceCheck:
+    if response.service == "dstack-shim":
+        message: Optional[str] = None
+        if (
+            instance_health_response is not None
+            and instance_health_response.dcgm is not None
+            and instance_health_response.dcgm.incidents
+        ):
+            message = instance_health_response.dcgm.incidents[0].error_message
+        return InstanceCheck(
+            reachable=True, health_response=instance_health_response, message=message
         )
+    return InstanceCheck(
+        reachable=False,
+        message=f"unexpected service: {response.service} version: {response.version}",
+        health_response=instance_health_response,
+    )
 
 
 def _volume_to_shim_volume_info(volume: Volume, instance_id: str) -> ShimVolumeInfo:

dstack/_internal/server/services/runner/ssh.py

@@ -2,7 +2,7 @@ import functools
 import socket
 import time
 from collections.abc import Iterable
-from typing import Callable, Dict, List, Optional, TypeVar, Union
+from typing import Callable, Dict, List, Literal, Optional, TypeVar, Union
 
 import requests
 from typing_extensions import Concatenate, ParamSpec
@@ -27,7 +27,7 @@ def runner_ssh_tunnel(
     [Callable[Concatenate[Dict[int, int], P], R]],
     Callable[
         Concatenate[PrivateKeyOrPair, JobProvisioningData, Optional[JobRuntimeData], P],
-        Union[bool, R],
+        Union[Literal[False], R],
     ],
 ]:
     """
@@ -42,7 +42,7 @@ def runner_ssh_tunnel(
         func: Callable[Concatenate[Dict[int, int], P], R],
     ) -> Callable[
         Concatenate[PrivateKeyOrPair, JobProvisioningData, Optional[JobRuntimeData], P],
-        Union[bool, R],
+        Union[Literal[False], R],
     ]:
         @functools.wraps(func)
         def wrapper(
@@ -51,7 +51,7 @@ def runner_ssh_tunnel(
             job_runtime_data: Optional[JobRuntimeData],
             *args: P.args,
             **kwargs: P.kwargs,
-        ) -> Union[bool, R]:
+        ) -> Union[Literal[False], R]:
             """
             Returns:
                 is successful