dstack 0.19.15rc1__py3-none-any.whl → 0.19.16__py3-none-any.whl

dstack/_internal/server/background/tasks/process_gateways.py

@@ -40,7 +40,7 @@ async def process_submitted_gateways():
                 .options(lazyload(GatewayModel.gateway_compute))
                 .order_by(GatewayModel.last_processed_at.asc())
                 .limit(1)
-                .with_for_update(skip_locked=True)
+                .with_for_update(skip_locked=True, key_share=True)
             )
             gateway_model = res.scalar()
             if gateway_model is None:

dstack/_internal/server/background/tasks/process_instances.py

@@ -149,7 +149,7 @@ async def _process_next_instance():
                 .options(lazyload(InstanceModel.jobs))
                 .order_by(InstanceModel.last_processed_at.asc())
                 .limit(1)
-                .with_for_update(skip_locked=True)
+                .with_for_update(skip_locked=True, key_share=True)
             )
             instance = res.scalar()
             if instance is None:

dstack/_internal/server/background/tasks/process_placement_groups.py

@@ -30,7 +30,7 @@ async def process_placement_groups():
                     PlacementGroupModel.id.not_in(lockset),
                 )
                 .order_by(PlacementGroupModel.id)  # take locks in order
-                .with_for_update(skip_locked=True)
+                .with_for_update(skip_locked=True, key_share=True)
             )
             placement_group_models = res.scalars().all()
             if len(placement_group_models) == 0:

dstack/_internal/server/background/tasks/process_running_jobs.py

@@ -101,7 +101,7 @@ async def _process_next_running_job():
                 )
                 .order_by(JobModel.last_processed_at.asc())
                 .limit(1)
-                .with_for_update(skip_locked=True)
+                .with_for_update(skip_locked=True, key_share=True)
             )
             job_model = res.unique().scalar()
             if job_model is None:

dstack/_internal/server/background/tasks/process_runs.py

@@ -27,6 +27,7 @@ from dstack._internal.server.services.jobs import (
     group_jobs_by_replica_latest,
 )
 from dstack._internal.server.services.locking import get_locker
+from dstack._internal.server.services.prometheus.client_metrics import run_metrics
 from dstack._internal.server.services.runs import (
     fmt,
     process_terminating_run,
@@ -62,7 +63,7 @@ async def _process_next_run():
                 )
                 .order_by(RunModel.last_processed_at.asc())
                 .limit(1)
-                .with_for_update(skip_locked=True)
+                .with_for_update(skip_locked=True, key_share=True)
             )
             run_model = res.scalar()
             if run_model is None:
@@ -74,7 +75,7 @@ async def _process_next_run():
                     JobModel.id.not_in(job_lockset),
                 )
                 .order_by(JobModel.id)  # take locks in order
-                .with_for_update(skip_locked=True)
+                .with_for_update(skip_locked=True, key_share=True)
             )
             job_models = res.scalars().all()
             if len(run_model.jobs) != len(job_models):
@@ -329,6 +330,24 @@ async def _process_active_run(session: AsyncSession, run_model: RunModel):
             run_model.status.name,
             new_status.name,
         )
+        if run_model.status == RunStatus.SUBMITTED and new_status == RunStatus.PROVISIONING:
+            current_time = common.get_current_datetime()
+            submit_to_provision_duration = (
+                current_time - run_model.submitted_at.replace(tzinfo=datetime.timezone.utc)
+            ).total_seconds()
+            logger.info(
+                "%s: run took %.2f seconds from submision to provisioning.",
+                fmt(run_model),
+                submit_to_provision_duration,
+            )
+            project_name = run_model.project.name
+            run_metrics.log_submit_to_provision_duration(
+                submit_to_provision_duration, project_name, run_spec.configuration.type
+            )
+
+        if new_status == RunStatus.PENDING:
+            run_metrics.increment_pending_runs(run_model.project.name, run_spec.configuration.type)
+
         run_model.status = new_status
         run_model.termination_reason = termination_reason
         # While a run goes to pending without provisioning, resubmission_attempt increases.

dstack/_internal/server/background/tasks/process_submitted_jobs.py

@@ -99,7 +99,7 @@ async def _process_next_submitted_job():
                     JobModel.id.not_in(lockset),
                 )
                 # Jobs are process in FIFO sorted by priority globally,
-                # thus runs from different project can "overtake" each other by using higher priorities.
+                # thus runs from different projects can "overtake" each other by using higher priorities.
                 # That's not a big problem as long as projects do not compete for the same compute resources.
                 # Jobs with lower priorities from other projects will be processed without major lag
                 # as long as new higher priority runs are not constantly submitted.
@@ -108,7 +108,13 @@ async def _process_next_submitted_job():
                 # there can be many projects and we are limited by the max DB connections.
                 .order_by(RunModel.priority.desc(), JobModel.last_processed_at.asc())
                 .limit(1)
-                .with_for_update(skip_locked=True)
+                .with_for_update(
+                    skip_locked=True,
+                    key_share=True,
+                    # Do not lock joined run, only job.
+                    # Locking run here may cause deadlock.
+                    of=JobModel,
+                )
             )
             job_model = res.scalar()
             if job_model is None:
@@ -201,7 +207,7 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
             )
             .options(lazyload(InstanceModel.jobs))
             .order_by(InstanceModel.id)  # take locks in order
-            .with_for_update()
+            .with_for_update(key_share=True)
         )
         pool_instances = list(res.unique().scalars().all())
         instances_ids = sorted([i.id for i in pool_instances])
@@ -326,7 +332,7 @@ async def _process_submitted_job(session: AsyncSession, job_model: JobModel):
         .where(VolumeModel.id.in_(volumes_ids))
         .options(selectinload(VolumeModel.user))
         .order_by(VolumeModel.id)  # take locks in order
-        .with_for_update()
+        .with_for_update(key_share=True)
     )
     async with get_locker().lock_ctx(VolumeModel.__tablename__, volumes_ids):
         if len(volume_models) > 0:

dstack/_internal/server/background/tasks/process_terminating_jobs.py

@@ -45,7 +45,7 @@ async def _process_next_terminating_job():
                 )
                 .order_by(JobModel.last_processed_at.asc())
                 .limit(1)
-                .with_for_update(skip_locked=True)
+                .with_for_update(skip_locked=True, key_share=True)
             )
             job_model = res.scalar()
             if job_model is None:
@@ -58,7 +58,7 @@ async def _process_next_terminating_job():
                         InstanceModel.id.not_in(instance_lockset),
                     )
                     .options(lazyload(InstanceModel.jobs))
-                    .with_for_update(skip_locked=True)
+                    .with_for_update(skip_locked=True, key_share=True)
                 )
                 instance_model = res.scalar()
                 if instance_model is None:

dstack/_internal/server/background/tasks/process_volumes.py

@@ -33,7 +33,7 @@ async def process_submitted_volumes():
                 )
                 .order_by(VolumeModel.last_processed_at.asc())
                 .limit(1)
-                .with_for_update(skip_locked=True)
+                .with_for_update(skip_locked=True, key_share=True)
             )
             volume_model = res.scalar()
             if volume_model is None:

dstack/_internal/server/routers/gateways.py

@@ -9,7 +9,10 @@ import dstack._internal.server.services.gateways as gateways
 from dstack._internal.core.errors import ResourceNotExistsError
 from dstack._internal.server.db import get_session
 from dstack._internal.server.models import ProjectModel, UserModel
-from dstack._internal.server.security.permissions import ProjectAdmin, ProjectMember
+from dstack._internal.server.security.permissions import (
+    ProjectAdmin,
+    ProjectMemberOrPublicAccess,
+)
 from dstack._internal.server.utils.routers import get_base_api_additional_responses
 
 router = APIRouter(
@@ -22,7 +25,7 @@ router = APIRouter(
 @router.post("/list")
 async def list_gateways(
     session: AsyncSession = Depends(get_session),
-    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
+    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMemberOrPublicAccess()),
 ) -> List[models.Gateway]:
     _, project = user_project
     return await gateways.list_project_gateways(session=session, project=project)
@@ -32,7 +35,7 @@ async def list_gateways(
 async def get_gateway(
     body: schemas.GetGatewayRequest,
     session: AsyncSession = Depends(get_session),
-    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMember()),
+    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectMemberOrPublicAccess()),
 ) -> models.Gateway:
     _, project = user_project
     gateway = await gateways.get_gateway_by_name(session=session, project=project, name=body.name)

dstack/_internal/server/routers/projects.py

@@ -7,13 +7,19 @@ from dstack._internal.core.models.projects import Project
 from dstack._internal.server.db import get_session
 from dstack._internal.server.models import ProjectModel, UserModel
 from dstack._internal.server.schemas.projects import (
+    AddProjectMemberRequest,
     CreateProjectRequest,
     DeleteProjectsRequest,
+    RemoveProjectMemberRequest,
     SetProjectMembersRequest,
+    UpdateProjectRequest,
 )
 from dstack._internal.server.security.permissions import (
     Authenticated,
+    ProjectAdmin,
     ProjectManager,
+    ProjectManagerOrPublicProject,
+    ProjectManagerOrSelfLeave,
     ProjectMemberOrPublicAccess,
 )
 from dstack._internal.server.services import projects
@@ -92,3 +98,60 @@ async def set_project_members(
     )
     await session.refresh(project)
     return projects.project_model_to_project(project)
+
+
+@router.post(
+    "/{project_name}/add_members",
+)
+async def add_project_members(
+    body: AddProjectMemberRequest,
+    session: AsyncSession = Depends(get_session),
+    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectManagerOrPublicProject()),
+) -> Project:
+    user, project = user_project
+    await projects.add_project_members(
+        session=session,
+        user=user,
+        project=project,
+        members=body.members,
+    )
+    await session.refresh(project)
+    return projects.project_model_to_project(project)
+
+
+@router.post(
+    "/{project_name}/remove_members",
+)
+async def remove_project_members(
+    body: RemoveProjectMemberRequest,
+    session: AsyncSession = Depends(get_session),
+    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectManagerOrSelfLeave()),
+) -> Project:
+    user, project = user_project
+    await projects.remove_project_members(
+        session=session,
+        user=user,
+        project=project,
+        usernames=body.usernames,
+    )
+    await session.refresh(project)
+    return projects.project_model_to_project(project)
+
+
+@router.post(
+    "/{project_name}/update",
+)
+async def update_project(
+    body: UpdateProjectRequest,
+    session: AsyncSession = Depends(get_session),
+    user_project: Tuple[UserModel, ProjectModel] = Depends(ProjectAdmin()),
+) -> Project:
+    user, project = user_project
+    await projects.update_project(
+        session=session,
+        user=user,
+        project=project,
+        is_public=body.is_public,
+    )
+    await session.refresh(project)
+    return projects.project_model_to_project(project)

dstack/_internal/server/routers/prometheus.py

@@ -1,15 +1,15 @@
 import os
 from typing import Annotated
 
+import prometheus_client
 from fastapi import APIRouter, Depends
 from fastapi.responses import PlainTextResponse
-from prometheus_client import generate_latest
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from dstack._internal.server import settings
 from dstack._internal.server.db import get_session
 from dstack._internal.server.security.permissions import OptionalServiceAccount
-from dstack._internal.server.services import prometheus
+from dstack._internal.server.services.prometheus import custom_metrics
 from dstack._internal.server.utils.routers import error_not_found
 
 _auth = OptionalServiceAccount(os.getenv("DSTACK_PROMETHEUS_AUTH_TOKEN"))
@@ -27,6 +27,6 @@ async def get_prometheus_metrics(
 ) -> str:
     if not settings.ENABLE_PROMETHEUS_METRICS:
         raise error_not_found()
-    custom_metrics = await prometheus.get_metrics(session=session)
-    prometheus_metrics = generate_latest()
-    return custom_metrics + prometheus_metrics.decode()
+    custom_metrics_ = await custom_metrics.get_metrics(session=session)
+    client_metrics = prometheus_client.generate_latest().decode()
+    return custom_metrics_ + client_metrics

dstack/_internal/server/schemas/logs.py

@@ -1,7 +1,7 @@
 from datetime import datetime
 from typing import Optional
 
-from pydantic import UUID4, Field
+from pydantic import UUID4, Field, validator
 
 from dstack._internal.core.models.common import CoreModel
 
@@ -12,5 +12,14 @@ class PollLogsRequest(CoreModel):
     start_time: Optional[datetime]
     end_time: Optional[datetime]
     descending: bool = False
+    next_token: Optional[str] = None
     limit: int = Field(100, ge=0, le=1000)
     diagnose: bool = False
+
+    @validator("descending")
+    @classmethod
+    def validate_descending(cls, v):
+        # Descending is not supported until we migrate from base64-encoded logs to plain text logs.
+        if v is True:
+            raise ValueError("descending: true is not supported")
+        return v

dstack/_internal/server/schemas/projects.py

@@ -11,6 +11,10 @@ class CreateProjectRequest(CoreModel):
     is_public: bool = False
 
 
+class UpdateProjectRequest(CoreModel):
+    is_public: bool
+
+
 class DeleteProjectsRequest(CoreModel):
     projects_names: List[str]
 
@@ -25,3 +29,11 @@ class MemberSetting(CoreModel):
 
 class SetProjectMembersRequest(CoreModel):
     members: List[MemberSetting]
+
+
+class AddProjectMemberRequest(CoreModel):
+    members: List[MemberSetting]
+
+
+class RemoveProjectMemberRequest(CoreModel):
+    usernames: List[str]

dstack/_internal/server/security/permissions.py

@@ -58,7 +58,7 @@ class ProjectAdmin:
             raise error_invalid_token()
         project = await get_project_model_by_name(session=session, project_name=project_name)
         if project is None:
-            raise error_forbidden()
+            raise error_not_found()
         if user.global_role == GlobalRole.ADMIN:
             return user, project
         project_role = get_user_project_role(user=user, project=project)
@@ -68,6 +68,10 @@ class ProjectAdmin:
 
 
 class ProjectManager:
+    """
+    Allows project admins and managers to manage projects.
+    """
+
     async def __call__(
         self,
         project_name: str,
@@ -79,12 +83,15 @@ class ProjectManager:
             raise error_invalid_token()
         project = await get_project_model_by_name(session=session, project_name=project_name)
         if project is None:
-            raise error_forbidden()
+            raise error_not_found()
+
         if user.global_role == GlobalRole.ADMIN:
             return user, project
+
         project_role = get_user_project_role(user=user, project=project)
         if project_role in [ProjectRole.ADMIN, ProjectRole.MANAGER]:
             return user, project
+
         raise error_forbidden()
 
 
@@ -135,6 +142,72 @@ class ProjectMemberOrPublicAccess:
         raise error_forbidden()
 
 
+class ProjectManagerOrPublicProject:
+    """
+    Allows:
+    1. Project managers to perform member management operations
+    2. Access to public projects for any authenticated user
+    """
+
+    def __init__(self):
+        self.project_manager = ProjectManager()
+
+    async def __call__(
+        self,
+        project_name: str,
+        session: AsyncSession = Depends(get_session),
+        token: HTTPAuthorizationCredentials = Security(HTTPBearer()),
+    ) -> Tuple[UserModel, ProjectModel]:
+        user = await log_in_with_token(session=session, token=token.credentials)
+        if user is None:
+            raise error_invalid_token()
+        project = await get_project_model_by_name(session=session, project_name=project_name)
+        if project is None:
+            raise error_not_found()
+
+        if user.global_role == GlobalRole.ADMIN:
+            return user, project
+
+        project_role = get_user_project_role(user=user, project=project)
+        if project_role in [ProjectRole.ADMIN, ProjectRole.MANAGER]:
+            return user, project
+
+        if project.is_public:
+            return user, project
+
+        raise error_forbidden()
+
+
+class ProjectManagerOrSelfLeave:
+    """
+    Allows:
+    1. Project managers to remove any members
+    2. Any project member to leave (remove themselves)
+    """
+
+    async def __call__(
+        self,
+        project_name: str,
+        session: AsyncSession = Depends(get_session),
+        token: HTTPAuthorizationCredentials = Security(HTTPBearer()),
+    ) -> Tuple[UserModel, ProjectModel]:
+        user = await log_in_with_token(session=session, token=token.credentials)
+        if user is None:
+            raise error_invalid_token()
+        project = await get_project_model_by_name(session=session, project_name=project_name)
+        if project is None:
+            raise error_not_found()
+
+        if user.global_role == GlobalRole.ADMIN:
+            return user, project
+
+        project_role = get_user_project_role(user=user, project=project)
+        if project_role is not None:
+            return user, project
+
+        raise error_forbidden()
+
+
 class OptionalServiceAccount:
     def __init__(self, token: Optional[str]) -> None:
         self._token = token

dstack/_internal/server/services/fleets.py

@@ -532,7 +532,7 @@ async def delete_fleets(
             .options(selectinload(FleetModel.runs))
             .execution_options(populate_existing=True)
             .order_by(FleetModel.id)  # take locks in order
-            .with_for_update()
+            .with_for_update(key_share=True)
         )
         fleet_models = res.scalars().unique().all()
         fleets = [fleet_model_to_fleet(m) for m in fleet_models]

dstack/_internal/server/services/gateways/__init__.py

@@ -240,7 +240,7 @@ async def delete_gateways(
             .options(selectinload(GatewayModel.gateway_compute))
             .execution_options(populate_existing=True)
             .order_by(GatewayModel.id)  # take locks in order
-            .with_for_update()
+            .with_for_update(key_share=True)
         )
         gateway_models = res.scalars().all()
         for gateway_model in gateway_models:

dstack/_internal/server/services/jobs/configurators/base.py

@@ -171,6 +171,8 @@ class JobConfigurator(ABC):
         return result
 
     def _dstack_image_commands(self) -> List[str]:
+        if self.run_spec.configuration.docker is True:
+            return ["start-dockerd"]
         if (
             self.run_spec.configuration.image is not None
             or self.run_spec.configuration.entrypoint is not None
@@ -201,7 +203,9 @@ class JobConfigurator(ABC):
         return self.run_spec.configuration.home_dir
 
     def _image_name(self) -> str:
-        if self.run_spec.configuration.image is not None:
+        if self.run_spec.configuration.docker is True:
+            return settings.DSTACK_DIND_IMAGE
+        elif self.run_spec.configuration.image is not None:
             return self.run_spec.configuration.image
         return get_default_image(nvcc=bool(self.run_spec.configuration.nvcc))
 
@@ -215,6 +219,8 @@ class JobConfigurator(ABC):
         return UnixUser.parse(user)
 
     def _privileged(self) -> bool:
+        if self.run_spec.configuration.docker is True:
+            return True
         return self.run_spec.configuration.privileged
 
     def _single_branch(self) -> bool:

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

@@ -78,14 +78,22 @@ class CloudWatchLogStorage(LogStorage):
             project.name, request.run_name, request.job_submission_id, log_producer
         )
         cw_events: List[_CloudWatchLogEvent]
+        next_token: Optional[str] = None
         with self._wrap_boto_errors():
             try:
-                cw_events = self._get_log_events(stream, request)
+                cw_events, next_token = self._get_log_events(stream, request)
             except botocore.exceptions.ClientError as e:
                 if not self._is_resource_not_found_exception(e):
                     raise
-                logger.debug("Stream %s not found, returning dummy response", stream)
-                cw_events = []
+                # Check if the group exists to distinguish between group not found vs stream not found
+                try:
+                    self._check_group_exists(self._group)
+                    # Group exists, so the error must be due to missing stream
+                    logger.debug("Stream %s not found, returning dummy response", stream)
+                    cw_events = []
+                except LogStorageError:
+                    # Group doesn't exist, re-raise the LogStorageError
+                    raise
         logs = [
             LogEvent(
                 timestamp=unix_time_ms_to_datetime(cw_event["timestamp"]),
@@ -94,51 +102,43 @@ class CloudWatchLogStorage(LogStorage):
             )
             for cw_event in cw_events
         ]
-        return JobSubmissionLogs(logs=logs)
+        return JobSubmissionLogs(logs=logs, next_token=next_token if len(logs) > 0 else None)
 
-    def _get_log_events(self, stream: str, request: PollLogsRequest) -> List[_CloudWatchLogEvent]:
-        limit = request.limit
+    def _get_log_events(
+        self, stream: str, request: PollLogsRequest
+    ) -> Tuple[List[_CloudWatchLogEvent], Optional[str]]:
+        start_from_head = not request.descending
         parameters = {
             "logGroupName": self._group,
             "logStreamName": stream,
-            "limit": limit,
+            "limit": request.limit,
+            "startFromHead": start_from_head,
         }
-        start_from_head = not request.descending
-        parameters["startFromHead"] = start_from_head
+
         if request.start_time:
-            # XXX: Since callers use start_time/end_time for pagination, one millisecond is added
-            # to avoid an infinite loop because startTime boundary is inclusive.
             parameters["startTime"] = datetime_to_unix_time_ms(request.start_time) + 1
+
         if request.end_time:
-            # No need to substract one millisecond in this case, though, seems that endTime is
-            # exclusive, that is, time interval boundaries are [startTime, entTime)
             parameters["endTime"] = datetime_to_unix_time_ms(request.end_time)
-        # "Partially full or empty pages don't necessarily mean that pagination is finished.
-        # As long as the nextBackwardToken or nextForwardToken returned is NOT equal to the
-        # nextToken that you passed into the API call, there might be more log events available."
-        events: List[_CloudWatchLogEvent] = []
-        next_token: Optional[str] = None
+        elif start_from_head:
+            # When startFromHead=true and no endTime is provided, set endTime to "now"
+            # to prevent infinite pagination as new logs arrive faster than we can read them
+            parameters["endTime"] = datetime_to_unix_time_ms(datetime.now(timezone.utc))
+
+        if request.next_token:
+            parameters["nextToken"] = request.next_token
+
+        response = self._client.get_log_events(**parameters)
+
+        events = response.get("events", [])
         next_token_key = "nextForwardToken" if start_from_head else "nextBackwardToken"
-        # Limit max tries to avoid a possible infinite loop if the API is misbehaving
-        tries_left = 10
-        while tries_left:
-            if next_token is not None:
-                parameters["nextToken"] = next_token
-            response = self._client.get_log_events(**parameters)
-            if start_from_head:
-                events.extend(response["events"])
-            else:
-                # Regardless of the startFromHead value log events are arranged in
-                # chronological order, from earliest to latest.
-                events.extend(reversed(response["events"]))
-            if len(events) >= limit:
-                return events[:limit]
-            if response[next_token_key] == next_token:
-                return events
-            next_token = response[next_token_key]
-            tries_left -= 1
-        logger.warning("too many requests to stream %s, returning partial response", stream)
-        return events
+        next_token = response.get(next_token_key)
+
+        # TODO: The code below is not going to be used until we migrate from base64-encoded logs to plain text logs.
+        if request.descending:
+            events = list(reversed(events))
+
+        return events, next_token
 
     def write_logs(
         self,