dstack 0.19.34__py3-none-any.whl → 0.19.35__py3-none-any.whl

dstack/_internal/core/models/compute_groups.py

@@ -0,0 +1,39 @@
+import enum
+import uuid
+from datetime import datetime
+from typing import List, Optional
+
+from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.common import CoreModel
+from dstack._internal.core.models.runs import JobProvisioningData
+
+
+class ComputeGroupStatus(str, enum.Enum):
+    RUNNING = "running"
+    TERMINATED = "terminated"
+
+
+class ComputeGroupProvisioningData(CoreModel):
+    compute_group_id: str
+    compute_group_name: str
+    backend: BackendType
+    # In case backend provisions instance in another backend,
+    # it may set that backend as base_backend.
+    base_backend: Optional[BackendType] = None
+    region: str
+    job_provisioning_datas: List[JobProvisioningData]
+    backend_data: Optional[str] = None  # backend-specific data in json
+
+
+class ComputeGroup(CoreModel):
+    """
+    Compute group is a group of instances managed as a single unit via the provider API,
+    i.e. instances are not created/deleted one-by-one but all at once.
+    """
+
+    id: uuid.UUID
+    name: str
+    project_name: str
+    created_at: datetime
+    status: ComputeGroupStatus
+    provisioning_data: ComputeGroupProvisioningData

dstack/_internal/core/models/fleets.py

@@ -309,7 +309,12 @@ class InstanceGroupParams(CoreModel):
     idle_duration: Annotated[
         Optional[int],
         Field(
-            description="Time to wait before terminating idle instances. Defaults to `5m` for runs and `3d` for fleets. Use `off` for unlimited duration"
+            description=(
+                "Time to wait before terminating idle instances."
+                " Instances are not terminated if the fleet is already at `nodes.min`."
+                " Defaults to `5m` for runs and `3d` for fleets."
+                " Use `off` for unlimited duration"
+            )
         ),
     ] = None
 

dstack/_internal/core/models/profiles.py

@@ -341,7 +341,9 @@ class ProfileParams(CoreModel):
         Field(
             description=(
                 "Time to wait before terminating idle instances."
-                " Defaults to `5m` for runs and `3d` for fleets. Use `off` for unlimited duration"
+                " Instances are not terminated if the fleet is already at `nodes.min`."
+                " Defaults to `5m` for runs and `3d` for fleets."
+                " Use `off` for unlimited duration"
             )
         ),
     ] = None

dstack/_internal/core/models/runs.py

@@ -207,6 +207,9 @@ class Requirements(CoreModel):
     max_price: Optional[float] = None
     spot: Optional[bool] = None
     reservation: Optional[str] = None
+    # Backends can use `multinode` to filter out offers if
+    # some offers support multinode and some do not.
+    multinode: Optional[bool] = None
 
     def pretty_format(self, resources_only: bool = False):
         res = self.resources.pretty_format()

dstack/_internal/server/app.py

@@ -5,7 +5,7 @@ import time
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import asynccontextmanager
 from pathlib import Path
-from typing import Awaitable, Callable, List
+from typing import Awaitable, Callable, List, Optional
 
 import sentry_sdk
 from fastapi import FastAPI, Request, Response, status
@@ -62,6 +62,7 @@ from dstack._internal.server.utils.routers import (
     CustomORJSONResponse,
     check_client_server_compatibility,
     error_detail,
+    get_client_version,
     get_server_client_error_details,
 )
 from dstack._internal.settings import DSTACK_VERSION
@@ -319,8 +320,19 @@ def register_routes(app: FastAPI, ui: bool = True):
             or request.url.path in _NO_API_VERSION_CHECK_ROUTES
         ):
             return await call_next(request)
+        try:
+            client_version = get_client_version(request)
+        except ValueError as e:
+            return CustomORJSONResponse(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                content={"detail": [error_detail(str(e))]},
+            )
+        client_release: Optional[tuple[int, ...]] = None
+        if client_version is not None:
+            client_release = client_version.release
+        request.state.client_release = client_release
         response = check_client_server_compatibility(
-            client_version=request.headers.get("x-api-version"),
+            client_version=client_version,
             server_version=DSTACK_VERSION,
         )
         if response is not None:

dstack/_internal/server/background/__init__.py

@@ -2,6 +2,7 @@ from apscheduler.schedulers.asyncio import AsyncIOScheduler
 from apscheduler.triggers.interval import IntervalTrigger
 
 from dstack._internal.server import settings
+from dstack._internal.server.background.tasks.process_compute_groups import process_compute_groups
 from dstack._internal.server.background.tasks.process_fleets import process_fleets
 from dstack._internal.server.background.tasks.process_gateways import (
     process_gateways,
@@ -122,5 +123,11 @@ def start_background_tasks() -> AsyncIOScheduler:
             kwargs={"batch_size": 5},
             max_instances=2 if replica == 0 else 1,
         )
+        _scheduler.add_job(
+            process_compute_groups,
+            IntervalTrigger(seconds=15, jitter=2),
+            kwargs={"batch_size": 1},
+            max_instances=2 if replica == 0 else 1,
+        )
     _scheduler.start()
     return _scheduler

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

@@ -0,0 +1,164 @@
+import asyncio
+import datetime
+from datetime import timedelta
+
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import joinedload, load_only
+
+from dstack._internal.core.backends.base.compute import ComputeWithGroupProvisioningSupport
+from dstack._internal.core.errors import BackendError
+from dstack._internal.core.models.compute_groups import ComputeGroupStatus
+from dstack._internal.core.models.instances import InstanceStatus
+from dstack._internal.server.db import get_db, get_session_ctx
+from dstack._internal.server.models import (
+    ComputeGroupModel,
+    ProjectModel,
+)
+from dstack._internal.server.services import backends as backends_services
+from dstack._internal.server.services.compute_groups import compute_group_model_to_compute_group
+from dstack._internal.server.services.locking import get_locker
+from dstack._internal.server.utils import sentry_utils
+from dstack._internal.utils.common import get_current_datetime, run_async
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+MIN_PROCESSING_INTERVAL = timedelta(seconds=30)
+
+TERMINATION_RETRY_TIMEOUT = timedelta(seconds=60)
+TERMINATION_RETRY_MAX_DURATION = timedelta(minutes=15)
+
+
+async def process_compute_groups(batch_size: int = 1):
+    tasks = []
+    for _ in range(batch_size):
+        tasks.append(_process_next_compute_group())
+    await asyncio.gather(*tasks)
+
+
+@sentry_utils.instrument_background_task
+async def _process_next_compute_group():
+    lock, lockset = get_locker(get_db().dialect_name).get_lockset(ComputeGroupModel.__tablename__)
+    async with get_session_ctx() as session:
+        async with lock:
+            res = await session.execute(
+                select(ComputeGroupModel)
+                .where(
+                    ComputeGroupModel.deleted == False,
+                    ComputeGroupModel.id.not_in(lockset),
+                    ComputeGroupModel.last_processed_at
+                    < get_current_datetime() - MIN_PROCESSING_INTERVAL,
+                )
+                .options(load_only(ComputeGroupModel.id))
+                .order_by(ComputeGroupModel.last_processed_at.asc())
+                .limit(1)
+                .with_for_update(skip_locked=True, key_share=True)
+            )
+            compute_group_model = res.scalar()
+            if compute_group_model is None:
+                return
+            compute_group_model_id = compute_group_model.id
+            lockset.add(compute_group_model_id)
+        try:
+            await _process_compute_group(
+                session=session,
+                compute_group_model=compute_group_model,
+            )
+        finally:
+            lockset.difference_update([compute_group_model_id])
+
+
+async def _process_compute_group(session: AsyncSession, compute_group_model: ComputeGroupModel):
+    # Refetch to load related attributes.
+    res = await session.execute(
+        select(ComputeGroupModel)
+        .where(ComputeGroupModel.id == compute_group_model.id)
+        .options(
+            joinedload(ComputeGroupModel.instances),
+            joinedload(ComputeGroupModel.project).joinedload(ProjectModel.backends),
+        )
+        .execution_options(populate_existing=True)
+    )
+    compute_group_model = res.unique().scalar_one()
+    if all(i.status == InstanceStatus.TERMINATING for i in compute_group_model.instances):
+        await _terminate_compute_group(compute_group_model)
+    compute_group_model.last_processed_at = get_current_datetime()
+    await session.commit()
+
+
+async def _terminate_compute_group(compute_group_model: ComputeGroupModel) -> None:
+    if (
+        compute_group_model.last_termination_retry_at is not None
+        and _next_termination_retry_at(compute_group_model) > get_current_datetime()
+    ):
+        return
+    compute_group = compute_group_model_to_compute_group(compute_group_model)
+    cgpd = compute_group.provisioning_data
+    backend = await backends_services.get_project_backend_by_type(
+        project=compute_group_model.project,
+        backend_type=cgpd.backend,
+    )
+    if backend is None:
+        logger.error(
+            "Failed to terminate compute group %s. Backend %s not available.",
+            compute_group.name,
+            cgpd.backend,
+        )
+    else:
+        logger.debug("Terminating compute group %s", compute_group.name)
+        compute = backend.compute()
+        assert isinstance(compute, ComputeWithGroupProvisioningSupport)
+        try:
+            await run_async(
+                compute.terminate_compute_group,
+                compute_group,
+            )
+        except Exception as e:
+            if compute_group_model.first_termination_retry_at is None:
+                compute_group_model.first_termination_retry_at = get_current_datetime()
+            compute_group_model.last_termination_retry_at = get_current_datetime()
+            if _next_termination_retry_at(compute_group_model) < _get_termination_deadline(
+                compute_group_model
+            ):
+                logger.warning(
+                    "Failed to terminate compute group %s. Will retry. Error: %r",
+                    compute_group.name,
+                    e,
+                    exc_info=not isinstance(e, BackendError),
+                )
+                return
+            logger.error(
+                "Failed all attempts to terminate compute group %s."
+                " Please terminate it manually to avoid unexpected charges."
+                " Error: %r",
+                compute_group.name,
+                e,
+                exc_info=not isinstance(e, BackendError),
+            )
+
+    compute_group_model.deleted = True
+    compute_group_model.deleted_at = get_current_datetime()
+    compute_group_model.status = ComputeGroupStatus.TERMINATED
+    # Terminating instances belonging to a compute group are locked implicitly
+    # by locking the compute group.
+    for instance_model in compute_group_model.instances:
+        instance_model.deleted = True
+        instance_model.deleted_at = get_current_datetime()
+        instance_model.finished_at = get_current_datetime()
+        instance_model.status = InstanceStatus.TERMINATED
+    logger.info(
+        "Terminated compute group %s",
+        compute_group.name,
+    )
+
+
+def _next_termination_retry_at(compute_group_model: ComputeGroupModel) -> datetime.datetime:
+    assert compute_group_model.last_termination_retry_at is not None
+    return compute_group_model.last_termination_retry_at + TERMINATION_RETRY_TIMEOUT
+
+
+def _get_termination_deadline(compute_group_model: ComputeGroupModel) -> datetime.datetime:
+    assert compute_group_model.first_termination_retry_at is not None
+    return compute_group_model.first_termination_retry_at + TERMINATION_RETRY_MAX_DURATION

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

@@ -8,7 +8,7 @@ import requests
 from paramiko.pkey import PKey
 from paramiko.ssh_exception import PasswordRequiredException
 from pydantic import ValidationError
-from sqlalchemy import delete, func, select
+from sqlalchemy import and_, delete, func, not_, select
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import joinedload
 
@@ -57,7 +57,6 @@ from dstack._internal.core.models.profiles import (
 )
 from dstack._internal.core.models.runs import (
     JobProvisioningData,
-    Retry,
 )
 from dstack._internal.server import settings as server_settings
 from dstack._internal.server.background.tasks.common import get_provisioning_timeout
@@ -167,6 +166,14 @@ async def _process_next_instance():
                             InstanceStatus.TERMINATING,
                         ]
                     ),
+                    # Terminating instances belonging to a compute group
+                    # are handled by process_compute_groups.
+                    not_(
+                        and_(
+                            InstanceModel.status == InstanceStatus.TERMINATING,
+                            InstanceModel.compute_group_id.is_not(None),
+                        )
+                    ),
                     InstanceModel.id.not_in(lockset),
                     InstanceModel.last_processed_at
                     < get_current_datetime() - MIN_PROCESSING_INTERVAL,
@@ -189,12 +196,12 @@ async def _process_next_instance():
 
 
 async def _process_instance(session: AsyncSession, instance: InstanceModel):
+    # Refetch to load related attributes.
+    # Load related attributes only for statuses that always need them.
     if instance.status in (
         InstanceStatus.PENDING,
         InstanceStatus.TERMINATING,
     ):
-        # Refetch to load related attributes.
-        # Load related attributes only for statuses that always need them.
         res = await session.execute(
             select(InstanceModel)
             .where(InstanceModel.id == instance.id)
@@ -204,6 +211,16 @@ async def _process_instance(session: AsyncSession, instance: InstanceModel):
             .execution_options(populate_existing=True)
         )
         instance = res.unique().scalar_one()
+    elif instance.status == InstanceStatus.IDLE:
+        res = await session.execute(
+            select(InstanceModel)
+            .where(InstanceModel.id == instance.id)
+            .options(joinedload(InstanceModel.project))
+            .options(joinedload(InstanceModel.jobs).load_only(JobModel.id, JobModel.status))
+            .options(joinedload(InstanceModel.fleet).joinedload(FleetModel.instances))
+            .execution_options(populate_existing=True)
+        )
+        instance = res.unique().scalar_one()
 
     if instance.status == InstanceStatus.PENDING:
         if instance.remote_connection_info is not None:
@@ -235,6 +252,14 @@ def _check_and_mark_terminating_if_idle_duration_expired(instance: InstanceModel
         and not instance.jobs
     ):
         return False
+    if instance.fleet is not None and not _can_terminate_fleet_instances_on_idle_duration(
+        instance.fleet
+    ):
+        logger.debug(
+            "Skipping instance %s termination on idle duration. Fleet is already at `nodes.min`.",
+            instance.name,
+        )
+        return False
     idle_duration = _get_instance_idle_duration(instance)
     idle_seconds = instance.termination_idle_time
     delta = datetime.timedelta(seconds=idle_seconds)
@@ -254,6 +279,20 @@ def _check_and_mark_terminating_if_idle_duration_expired(instance: InstanceModel
     return False
 
 
+def _can_terminate_fleet_instances_on_idle_duration(fleet_model: FleetModel) -> bool:
+    # Do not terminate instances on idle duration if fleet is already at `nodes.min`.
+    # This is an optimization to avoid terminate-create loop.
+    # There may be race conditions since we don't take the fleet lock.
+    # That's ok: in the worst case we go below `nodes.min`, but
+    # the fleet consolidation logic will provision new nodes.
+    fleet = fleet_model_to_fleet(fleet_model)
+    if fleet.spec.configuration.nodes is None or fleet.spec.autocreated:
+        return True
+    active_instances = [i for i in fleet_model.instances if i.status.is_active()]
+    active_instances_num = len(active_instances)
+    return active_instances_num > fleet.spec.configuration.nodes.min
+
+
 async def _add_remote(instance: InstanceModel) -> None:
     logger.info("Adding ssh instance %s...", instance.name)
     if instance.status == InstanceStatus.PENDING:
@@ -918,51 +957,48 @@ async def _terminate(instance: InstanceModel) -> None:
     ):
         return
     jpd = get_instance_provisioning_data(instance)
-    if jpd is not None:
-        if jpd.backend != BackendType.REMOTE:
-            backend = await backends_services.get_project_backend_by_type(
-                project=instance.project, backend_type=jpd.backend
+    if jpd is not None and jpd.backend != BackendType.REMOTE:
+        backend = await backends_services.get_project_backend_by_type(
+            project=instance.project, backend_type=jpd.backend
+        )
+        if backend is None:
+            logger.error(
+                "Failed to terminate instance %s. Backend %s not available.",
+                instance.name,
+                jpd.backend,
             )
-            if backend is None:
+        else:
+            logger.debug("Terminating runner instance %s", jpd.hostname)
+            try:
+                await run_async(
+                    backend.compute().terminate_instance,
+                    jpd.instance_id,
+                    jpd.region,
+                    jpd.backend_data,
+                )
+            except Exception as e:
+                if instance.first_termination_retry_at is None:
+                    instance.first_termination_retry_at = get_current_datetime()
+                instance.last_termination_retry_at = get_current_datetime()
+                if _next_termination_retry_at(instance) < _get_termination_deadline(instance):
+                    if isinstance(e, NotYetTerminated):
+                        logger.debug("Instance %s termination in progress: %s", instance.name, e)
+                    else:
+                        logger.warning(
+                            "Failed to terminate instance %s. Will retry. Error: %r",
+                            instance.name,
+                            e,
+                            exc_info=not isinstance(e, BackendError),
+                        )
+                    return
                 logger.error(
-                    "Failed to terminate instance %s. Backend %s not available.",
+                    "Failed all attempts to terminate instance %s."
+                    " Please terminate the instance manually to avoid unexpected charges."
+                    " Error: %r",
                     instance.name,
-                    jpd.backend,
+                    e,
+                    exc_info=not isinstance(e, BackendError),
                 )
-            else:
-                logger.debug("Terminating runner instance %s", jpd.hostname)
-                try:
-                    await run_async(
-                        backend.compute().terminate_instance,
-                        jpd.instance_id,
-                        jpd.region,
-                        jpd.backend_data,
-                    )
-                except Exception as e:
-                    if instance.first_termination_retry_at is None:
-                        instance.first_termination_retry_at = get_current_datetime()
-                    instance.last_termination_retry_at = get_current_datetime()
-                    if _next_termination_retry_at(instance) < _get_termination_deadline(instance):
-                        if isinstance(e, NotYetTerminated):
-                            logger.debug(
-                                "Instance %s termination in progress: %s", instance.name, e
-                            )
-                        else:
-                            logger.warning(
-                                "Failed to terminate instance %s. Will retry. Error: %r",
-                                instance.name,
-                                e,
-                                exc_info=not isinstance(e, BackendError),
-                            )
-                        return
-                    logger.error(
-                        "Failed all attempts to terminate instance %s."
-                        " Please terminate the instance manually to avoid unexpected charges."
-                        " Error: %r",
-                        instance.name,
-                        e,
-                        exc_info=not isinstance(e, BackendError),
-                    )
 
     instance.deleted = True
     instance.deleted_at = get_current_datetime()
@@ -1126,10 +1162,6 @@ def _get_instance_idle_duration(instance: InstanceModel) -> datetime.timedelta:
     return get_current_datetime() - last_time
 
 
-def _get_retry_duration_deadline(instance: InstanceModel, retry: Retry) -> datetime.datetime:
-    return instance.created_at + timedelta(seconds=retry.duration)
-
-
 def _get_provisioning_deadline(
     instance: InstanceModel,
     job_provisioning_data: JobProvisioningData,