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

dstack/_internal/core/models/profiles.py

@@ -8,6 +8,7 @@ from typing_extensions import Annotated, Literal
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import CoreModel, Duration
 from dstack._internal.utils.common import list_enum_values_for_annotation
+from dstack._internal.utils.cron import validate_cron
 from dstack._internal.utils.json_utils import pydantic_orjson_dumps_with_indent
 from dstack._internal.utils.tags import tags_validator
 
@@ -167,6 +168,38 @@ class UtilizationPolicy(CoreModel):
         return v
 
 
+class Schedule(CoreModel):
+    cron: Annotated[
+        Union[List[str], str],
+        Field(
+            description=(
+                "A cron expression or a list of cron expressions specifying the UTC time when the run needs to be started"
+            )
+        ),
+    ]
+
+    @validator("cron")
+    def _validate_cron(cls, v: Union[List[str], str]) -> List[str]:
+        if isinstance(v, str):
+            values = [v]
+        else:
+            values = v
+        if len(values) == 0:
+            raise ValueError("At least one cron expression must be specified")
+        for value in values:
+            validate_cron(value)
+        return values
+
+    @property
+    def crons(self) -> List[str]:
+        """
+        Access `cron` attribute as a list.
+        """
+        if isinstance(self.cron, str):
+            return [self.cron]
+        return self.cron
+
+
 class ProfileParams(CoreModel):
     backends: Annotated[
         Optional[List[BackendType]],
@@ -281,6 +314,10 @@ class ProfileParams(CoreModel):
             )
         ),
     ] = None
+    schedule: Annotated[
+        Optional[Schedule],
+        Field(description=("The schedule for starting the run at specified time")),
+    ] = None
     fleets: Annotated[
         Optional[list[str]], Field(description="The fleets considered for reuse")
     ] = None

dstack/_internal/core/models/runs.py

@@ -1,6 +1,6 @@
 from datetime import datetime, timedelta
 from enum import Enum
-from typing import Any, Dict, List, Optional, Type
+from typing import Any, Dict, List, Literal, Optional, Type
 
 from pydantic import UUID4, Field, root_validator
 from typing_extensions import Annotated
@@ -8,8 +8,11 @@ from typing_extensions import Annotated
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.common import ApplyAction, CoreModel, NetworkMode, RegistryAuth
 from dstack._internal.core.models.configurations import (
+    DEFAULT_PROBE_METHOD,
     DEFAULT_REPO_DIR,
     AnyRunConfiguration,
+    HTTPHeaderSpec,
+    HTTPMethod,
     RunConfiguration,
     ServiceConfiguration,
 )
@@ -223,6 +226,17 @@ class JobSSHKey(CoreModel):
     public: str
 
 
+class ProbeSpec(CoreModel):
+    type: Literal["http"]  # expect other probe types in the future, namely `exec`
+    url: str
+    method: HTTPMethod = DEFAULT_PROBE_METHOD
+    headers: list[HTTPHeaderSpec] = []
+    body: Optional[str] = None
+    timeout: int
+    interval: int
+    ready_after: int
+
+
 class JobSpec(CoreModel):
     replica_num: int = 0  # default value for backward compatibility
     job_num: int
@@ -256,6 +270,7 @@ class JobSpec(CoreModel):
     file_archives: list[FileArchiveMapping] = []
     # None for non-services and pre-0.19.19 services. See `get_service_port`
     service_port: Optional[int] = None
+    probes: list[ProbeSpec] = []
 
 
 class JobProvisioningData(CoreModel):
@@ -325,6 +340,10 @@ class ClusterInfo(CoreModel):
     gpus_per_job: int
 
 
+class Probe(CoreModel):
+    success_streak: int
+
+
 class JobSubmission(CoreModel):
     id: UUID4
     submission_num: int
@@ -341,6 +360,7 @@ class JobSubmission(CoreModel):
     job_provisioning_data: Optional[JobProvisioningData]
     job_runtime_data: Optional[JobRuntimeData]
     error: Optional[str] = None
+    probes: list[Probe] = []
 
     @property
     def age(self) -> timedelta:

dstack/_internal/core/services/ssh/tunnel.py

@@ -236,6 +236,13 @@ class SSHTunnel:
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.close()
 
+    async def __aenter__(self):
+        await self.aopen()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.aclose()
+
     def _get_proxy_command(self) -> Optional[str]:
         proxy_command: Optional[str] = None
         for params, identity_path in self.ssh_proxies:

dstack/_internal/server/app.py

@@ -13,6 +13,7 @@ from fastapi.datastructures import URL
 from fastapi.responses import HTMLResponse, RedirectResponse
 from fastapi.staticfiles import StaticFiles
 from prometheus_client import Counter, Histogram
+from sentry_sdk.types import SamplingContext
 
 from dstack._internal.cli.utils.common import console
 from dstack._internal.core.errors import ForbiddenError, ServerClientError
@@ -21,6 +22,7 @@ from dstack._internal.proxy.lib.deps import get_injector_from_app
 from dstack._internal.proxy.lib.routers import model_proxy
 from dstack._internal.server import settings
 from dstack._internal.server.background import start_background_tasks
+from dstack._internal.server.background.tasks.process_probes import PROBES_SCHEDULER
 from dstack._internal.server.db import get_db, get_session_ctx, migrate
 from dstack._internal.server.routers import (
     backends,
@@ -81,16 +83,6 @@ REQUEST_DURATION = Histogram(
 
 
 def create_app() -> FastAPI:
-    if settings.SENTRY_DSN is not None:
-        sentry_sdk.init(
-            dsn=settings.SENTRY_DSN,
-            release=DSTACK_VERSION,
-            environment=settings.SERVER_ENVIRONMENT,
-            enable_tracing=True,
-            traces_sample_rate=settings.SENTRY_TRACES_SAMPLE_RATE,
-            profiles_sample_rate=settings.SENTRY_PROFILES_SAMPLE_RATE,
-        )
-
     app = FastAPI(
         docs_url="/api/docs",
         lifespan=lifespan,
@@ -102,6 +94,15 @@ def create_app() -> FastAPI:
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     configure_logging()
+    if settings.SENTRY_DSN is not None:
+        sentry_sdk.init(
+            dsn=settings.SENTRY_DSN,
+            release=DSTACK_VERSION,
+            environment=settings.SERVER_ENVIRONMENT,
+            enable_tracing=True,
+            traces_sampler=_sentry_traces_sampler,
+            profiles_sample_rate=settings.SENTRY_PROFILES_SAMPLE_RATE,
+        )
     server_executor = ThreadPoolExecutor(max_workers=settings.SERVER_EXECUTOR_MAX_WORKERS)
     asyncio.get_running_loop().set_default_executor(server_executor)
     await migrate()
@@ -155,6 +156,7 @@ async def lifespan(app: FastAPI):
         scheduler = start_background_tasks()
     else:
         logger.info("Background processing is disabled")
+    PROBES_SCHEDULER.start()
     dstack_version = DSTACK_VERSION if DSTACK_VERSION else "(no version)"
     logger.info(f"The admin token is {admin.token.get_plaintext_or_error()}", {"show_path": False})
     logger.info(
@@ -166,6 +168,7 @@ async def lifespan(app: FastAPI):
     yield
     if settings.SERVER_BACKGROUND_PROCESSING_ENABLED:
         scheduler.shutdown()
+    PROBES_SCHEDULER.shutdown(wait=False)
     await gateway_connections_pool.remove_all()
     service_conn_pool = await get_injector_from_app(app).get_service_connection_pool()
     await service_conn_pool.remove_all()
@@ -197,6 +200,7 @@ def register_routes(app: FastAPI, ui: bool = True):
     app.include_router(fleets.root_router)
     app.include_router(fleets.project_router)
     app.include_router(instances.root_router)
+    app.include_router(instances.project_router)
     app.include_router(repos.router)
     app.include_router(runs.root_router)
     app.include_router(runs.project_router)
@@ -379,3 +383,15 @@ def _print_dstack_logo():
 ╰━━┻━━┻╯╱╰╯╰━━┻╯
 [/]"""
     )
+
+
+def _sentry_traces_sampler(sampling_context: SamplingContext) -> float:
+    parent_sampling_decision = sampling_context["parent_sampled"]
+    if parent_sampling_decision is not None:
+        return float(parent_sampling_decision)
+    transaction_context = sampling_context["transaction_context"]
+    name = transaction_context.get("name")
+    if name is not None:
+        if name.startswith("background."):
+            return settings.SENTRY_TRACES_BACKGROUND_SAMPLE_RATE
+    return settings.SENTRY_TRACES_SAMPLE_RATE

dstack/_internal/server/background/__init__.py

@@ -9,6 +9,7 @@ from dstack._internal.server.background.tasks.process_gateways import (
 )
 from dstack._internal.server.background.tasks.process_idle_volumes import process_idle_volumes
 from dstack._internal.server.background.tasks.process_instances import (
+    delete_instance_health_checks,
     process_instances,
 )
 from dstack._internal.server.background.tasks.process_metrics import (
@@ -18,6 +19,7 @@ from dstack._internal.server.background.tasks.process_metrics import (
 from dstack._internal.server.background.tasks.process_placement_groups import (
     process_placement_groups,
 )
+from dstack._internal.server.background.tasks.process_probes import process_probes
 from dstack._internal.server.background.tasks.process_prometheus_metrics import (
     collect_prometheus_metrics,
     delete_prometheus_metrics,
@@ -63,6 +65,7 @@ def start_background_tasks() -> AsyncIOScheduler:
     # that the first waiting for the lock will acquire it.
     # The jitter is needed to give all tasks a chance to acquire locks.
 
+    _scheduler.add_job(process_probes, IntervalTrigger(seconds=3, jitter=1))
     _scheduler.add_job(collect_metrics, IntervalTrigger(seconds=10), max_instances=1)
     _scheduler.add_job(delete_metrics, IntervalTrigger(minutes=5), max_instances=1)
     if settings.ENABLE_PROMETHEUS_METRICS:
@@ -79,6 +82,12 @@ def start_background_tasks() -> AsyncIOScheduler:
         process_idle_volumes, IntervalTrigger(seconds=60, jitter=10), max_instances=1
     )
     _scheduler.add_job(process_placement_groups, IntervalTrigger(seconds=30, jitter=5))
+    _scheduler.add_job(
+        process_fleets,
+        IntervalTrigger(seconds=10, jitter=2),
+        max_instances=1,
+    )
+    _scheduler.add_job(delete_instance_health_checks, IntervalTrigger(minutes=5), max_instances=1)
     for replica in range(settings.SERVER_BACKGROUND_PROCESSING_FACTOR):
         # Add multiple copies of tasks if requested.
         # max_instances=1 for additional copies to avoid running too many tasks.
@@ -113,11 +122,5 @@ def start_background_tasks() -> AsyncIOScheduler:
             kwargs={"batch_size": 5},
             max_instances=2 if replica == 0 else 1,
         )
-        _scheduler.add_job(
-            process_fleets,
-            IntervalTrigger(seconds=10, jitter=2),
-            kwargs={"batch_size": 5},
-            max_instances=2 if replica == 0 else 1,
-        )
     _scheduler.start()
     return _scheduler

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

@@ -1,36 +1,37 @@
-import asyncio
 from datetime import timedelta
+from typing import List
 
-from sqlalchemy import select
+from sqlalchemy import select, update
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload
+from sqlalchemy.orm import joinedload, load_only
 
 from dstack._internal.core.models.fleets import FleetStatus
 from dstack._internal.server.db import get_db, get_session_ctx
-from dstack._internal.server.models import FleetModel
+from dstack._internal.server.models import (
+    FleetModel,
+    InstanceModel,
+    JobModel,
+    PlacementGroupModel,
+    RunModel,
+)
 from dstack._internal.server.services.fleets import (
     is_fleet_empty,
     is_fleet_in_use,
 )
 from dstack._internal.server.services.locking import get_locker
-from dstack._internal.server.services.placement import schedule_fleet_placement_groups_deletion
+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__)
 
 
+BATCH_SIZE = 10
 MIN_PROCESSING_INTERVAL = timedelta(seconds=30)
 
 
-async def process_fleets(batch_size: int = 1):
-    tasks = []
-    for _ in range(batch_size):
-        tasks.append(_process_next_fleet())
-    await asyncio.gather(*tasks)
-
-
-async def _process_next_fleet():
+@sentry_utils.instrument_background_task
+async def process_fleets():
     lock, lockset = get_locker(get_db().dialect_name).get_lockset(FleetModel.__tablename__)
     async with get_session_ctx() as session:
         async with lock:
@@ -40,51 +41,64 @@ async def _process_next_fleet():
                     FleetModel.deleted == False,
                     FleetModel.id.not_in(lockset),
                     FleetModel.last_processed_at
-                    < get_current_datetime().replace(tzinfo=None) - MIN_PROCESSING_INTERVAL,
+                    < get_current_datetime() - MIN_PROCESSING_INTERVAL,
                 )
+                .options(load_only(FleetModel.id))
                 .order_by(FleetModel.last_processed_at.asc())
-                .limit(1)
+                .limit(BATCH_SIZE)
                 .with_for_update(skip_locked=True, key_share=True)
             )
-            fleet_model = res.scalar()
-            if fleet_model is None:
-                return
-            lockset.add(fleet_model.id)
+            fleet_models = list(res.scalars().all())
+            fleet_ids = [fm.id for fm in fleet_models]
+            for fleet_id in fleet_ids:
+                lockset.add(fleet_id)
         try:
-            fleet_model_id = fleet_model.id
-            await _process_fleet(session=session, fleet_model=fleet_model)
+            await _process_fleets(session=session, fleet_models=fleet_models)
         finally:
-            lockset.difference_update([fleet_model_id])
+            lockset.difference_update(fleet_ids)
 
 
-async def _process_fleet(session: AsyncSession, fleet_model: FleetModel):
-    logger.debug("Processing fleet %s", fleet_model.name)
+async def _process_fleets(session: AsyncSession, fleet_models: List[FleetModel]):
+    fleet_ids = [fm.id for fm in fleet_models]
     # Refetch to load related attributes.
-    # joinedload produces LEFT OUTER JOIN that can't be used with FOR UPDATE.
     res = await session.execute(
         select(FleetModel)
-        .where(FleetModel.id == fleet_model.id)
-        .options(joinedload(FleetModel.project))
-        .options(joinedload(FleetModel.instances))
-        .options(joinedload(FleetModel.runs))
+        .where(FleetModel.id.in_(fleet_ids))
+        .options(joinedload(FleetModel.instances).load_only(InstanceModel.deleted))
+        .options(
+            joinedload(FleetModel.instances).joinedload(InstanceModel.jobs).load_only(JobModel.id)
+        )
+        .options(joinedload(FleetModel.runs).load_only(RunModel.status))
         .execution_options(populate_existing=True)
     )
-    fleet_model = res.unique().scalar_one()
-    await _autodelete_fleet(session=session, fleet_model=fleet_model)
+    fleet_models = list(res.unique().scalars().all())
+
+    deleted_fleets_ids = []
+    now = get_current_datetime()
+    for fleet_model in fleet_models:
+        deleted = _autodelete_fleet(fleet_model)
+        if deleted:
+            deleted_fleets_ids.append(fleet_model.id)
+        fleet_model.last_processed_at = now
+
+    await session.execute(
+        update(PlacementGroupModel)
+        .where(
+            PlacementGroupModel.fleet_id.in_(deleted_fleets_ids),
+        )
+        .values(fleet_deleted=True)
+    )
+    await session.commit()
 
 
-async def _autodelete_fleet(session: AsyncSession, fleet_model: FleetModel):
+def _autodelete_fleet(fleet_model: FleetModel) -> bool:
     # Currently all empty fleets are autodeleted.
     # TODO: If fleets with `nodes: 0..` are supported, their deletion should be skipped.
     if is_fleet_in_use(fleet_model) or not is_fleet_empty(fleet_model):
-        fleet_model.last_processed_at = get_current_datetime()
-        await session.commit()
-        return
+        return False
 
     logger.info("Automatic cleanup of an empty fleet %s", fleet_model.name)
     fleet_model.status = FleetStatus.TERMINATED
     fleet_model.deleted = True
-    fleet_model.last_processed_at = get_current_datetime()
-    await schedule_fleet_placement_groups_deletion(session=session, fleet_id=fleet_model.id)
-    await session.commit()
     logger.info("Fleet %s deleted", fleet_model.name)
+    return True

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

@@ -17,6 +17,7 @@ from dstack._internal.server.services.gateways import (
 )
 from dstack._internal.server.services.locking import advisory_lock_ctx, get_locker
 from dstack._internal.server.services.logging import fmt
+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
 
@@ -28,6 +29,7 @@ async def process_gateways_connections():
     await _process_active_connections()
 
 
+@sentry_utils.instrument_background_task
 async def process_gateways():
     lock, lockset = get_locker(get_db().dialect_name).get_lockset(GatewayModel.__tablename__)
     async with get_session_ctx() as session:
@@ -110,7 +112,6 @@ async def _process_connection(conn: GatewayConnection):
 async def _process_submitted_gateway(session: AsyncSession, gateway_model: GatewayModel):
     logger.info("%s: started gateway provisioning", fmt(gateway_model))
     # Refetch to load related attributes.
-    # joinedload produces LEFT OUTER JOIN that can't be used with FOR UPDATE.
     res = await session.execute(
         select(GatewayModel)
         .where(GatewayModel.id == gateway_model.id)
@@ -157,7 +158,6 @@ async def _process_provisioning_gateway(
     session: AsyncSession, gateway_model: GatewayModel
 ) -> None:
     # Refetch to load related attributes.
-    # joinedload produces LEFT OUTER JOIN that can't be used with FOR UPDATE.
     res = await session.execute(
         select(GatewayModel)
         .where(GatewayModel.id == gateway_model.id)

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

@@ -10,13 +10,14 @@ from dstack._internal.core.errors import BackendNotAvailable
 from dstack._internal.core.models.profiles import parse_duration
 from dstack._internal.core.models.volumes import VolumeStatus
 from dstack._internal.server.db import get_db, get_session_ctx
-from dstack._internal.server.models import ProjectModel, VolumeModel
+from dstack._internal.server.models import ProjectModel, UserModel, VolumeModel
 from dstack._internal.server.services import backends as backends_services
 from dstack._internal.server.services.locking import get_locker
 from dstack._internal.server.services.volumes import (
     get_volume_configuration,
     volume_model_to_volume,
 )
+from dstack._internal.server.utils import sentry_utils
 from dstack._internal.utils import common
 from dstack._internal.utils.common import get_current_datetime
 from dstack._internal.utils.logging import get_logger
@@ -24,6 +25,7 @@ from dstack._internal.utils.logging import get_logger
 logger = get_logger(__name__)
 
 
+@sentry_utils.instrument_background_task
 async def process_idle_volumes():
     lock, lockset = get_locker(get_db().dialect_name).get_lockset(VolumeModel.__tablename__)
     async with get_session_ctx() as session:
@@ -49,7 +51,7 @@ async def process_idle_volumes():
             select(VolumeModel)
             .where(VolumeModel.id.in_(volume_ids))
             .options(joinedload(VolumeModel.project).joinedload(ProjectModel.backends))
-            .options(joinedload(VolumeModel.user))
+            .options(joinedload(VolumeModel.user).load_only(UserModel.name))
             .options(joinedload(VolumeModel.attachments))
             .execution_options(populate_existing=True)
         )
@@ -82,8 +84,7 @@ def _should_delete_volume(volume: VolumeModel) -> bool:
 
 def _get_idle_time(volume: VolumeModel) -> datetime.timedelta:
     last_used = volume.last_job_processed_at or volume.created_at
-    last_used_utc = last_used.replace(tzinfo=datetime.timezone.utc)
-    idle_time = get_current_datetime() - last_used_utc
+    idle_time = get_current_datetime() - last_used
     return max(idle_time, datetime.timedelta(0))