dstack 0.19.25rc1__py3-none-any.whl → 0.19.26__py3-none-any.whl

dstack/_internal/server/services/gpus.py

@@ -1,8 +1,8 @@
 from typing import Dict, List, Literal, Optional, Tuple
 
-from sqlalchemy.ext.asyncio import AsyncSession
-
 from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.errors import ServerClientError
+from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.instances import InstanceOfferWithAvailability
 from dstack._internal.core.models.profiles import SpotPolicy
 from dstack._internal.core.models.resources import Range
@@ -15,10 +15,43 @@ from dstack._internal.server.schemas.gpus import (
     ListGpusResponse,
 )
 from dstack._internal.server.services.offers import get_offers_by_requirements
+from dstack._internal.utils.common import get_or_error
+
+
+async def list_gpus_grouped(
+    project: ProjectModel,
+    run_spec: RunSpec,
+    group_by: Optional[List[Literal["backend", "region", "count"]]] = None,
+) -> ListGpusResponse:
+    """Retrieves available GPU specifications based on a run spec, with optional grouping."""
+    offers = await _get_gpu_offers(project=project, run_spec=run_spec)
+    backend_gpus = _process_offers_into_backend_gpus(offers)
+    group_by_set = set(group_by) if group_by else set()
+    if "region" in group_by_set and "backend" not in group_by_set:
+        raise ServerClientError("Cannot group by 'region' without also grouping by 'backend'")
+
+    # Determine grouping strategy based on combination
+    has_backend = "backend" in group_by_set
+    has_region = "region" in group_by_set
+    has_count = "count" in group_by_set
+    if has_backend and has_region and has_count:
+        gpus = _get_gpus_grouped_by_backend_region_and_count(backend_gpus)
+    elif has_backend and has_count:
+        gpus = _get_gpus_grouped_by_backend_and_count(backend_gpus)
+    elif has_backend and has_region:
+        gpus = _get_gpus_grouped_by_backend_and_region(backend_gpus)
+    elif has_backend:
+        gpus = _get_gpus_grouped_by_backend(backend_gpus)
+    elif has_count:
+        gpus = _get_gpus_grouped_by_count(backend_gpus)
+    else:
+        gpus = _get_gpus_with_no_grouping(backend_gpus)
+
+    return ListGpusResponse(gpus=gpus)
 
 
 async def _get_gpu_offers(
-    session: AsyncSession, project: ProjectModel, run_spec: RunSpec
+    project: ProjectModel, run_spec: RunSpec
 ) -> List[Tuple[Backend, InstanceOfferWithAvailability]]:
     """Fetches all available instance offers that match the run spec's GPU requirements."""
     profile = run_spec.merged_profile
@@ -28,7 +61,6 @@ async def _get_gpu_offers(
         spot=get_policy_map(profile.spot_policy, default=SpotPolicy.AUTO),
         reservation=profile.reservation,
     )
-
     return await get_offers_by_requirements(
         project=project,
         profile=profile,
@@ -45,10 +77,10 @@ def _process_offers_into_backend_gpus(
     offers: List[Tuple[Backend, InstanceOfferWithAvailability]],
 ) -> List[BackendGpus]:
     """Transforms raw offers into a structured list of BackendGpus, aggregating GPU info."""
-    backend_data: Dict[str, Dict] = {}
+    backend_data: Dict[BackendType, Dict] = {}
 
-    for backend, offer in offers:
-        backend_type = backend.TYPE
+    for _, offer in offers:
+        backend_type = offer.backend
         if backend_type not in backend_data:
             backend_data[backend_type] = {"gpus": {}, "regions": set()}
 
@@ -111,7 +143,7 @@ def _process_offers_into_backend_gpus(
     return backend_gpus_list
 
 
-def _update_gpu_group(row: GpuGroup, gpu: BackendGpu, backend_type: str):
+def _update_gpu_group(row: GpuGroup, gpu: BackendGpu, backend_type: BackendType):
     """Updates an existing GpuGroup with new data from another GPU offer."""
     spot_type: Literal["spot", "on-demand"] = "spot" if gpu.spot else "on-demand"
 
@@ -122,6 +154,12 @@ def _update_gpu_group(row: GpuGroup, gpu: BackendGpu, backend_type: str):
     if row.backends and backend_type not in row.backends:
         row.backends.append(backend_type)
 
+    # FIXME: Consider using non-optional range
+    assert row.count.min is not None
+    assert row.count.max is not None
+    assert row.price.min is not None
+    assert row.price.max is not None
+
     row.count.min = min(row.count.min, gpu.count)
     row.count.max = max(row.count.max, gpu.count)
     per_gpu_price = gpu.price / gpu.count
@@ -194,7 +232,7 @@ def _get_gpus_grouped_by_backend(backend_gpus: List[BackendGpus]) -> List[GpuGro
             not any(av.is_available() for av in g.availability),
             g.price.min,
             g.price.max,
-            g.backend.value,
+            get_or_error(g.backend).value,
             g.name,
             g.memory_mib,
         ),
@@ -229,7 +267,7 @@ def _get_gpus_grouped_by_backend_and_region(backend_gpus: List[BackendGpus]) ->
             not any(av.is_available() for av in g.availability),
             g.price.min,
             g.price.max,
-            g.backend.value,
+            get_or_error(g.backend).value,
             g.region,
             g.name,
             g.memory_mib,
@@ -299,7 +337,7 @@ def _get_gpus_grouped_by_backend_and_count(backend_gpus: List[BackendGpus]) -> L
             not any(av.is_available() for av in g.availability),
             g.price.min,
             g.price.max,
-            g.backend.value,
+            get_or_error(g.backend).value,
             g.count.min,
             g.name,
             g.memory_mib,
@@ -344,47 +382,10 @@ def _get_gpus_grouped_by_backend_region_and_count(
             not any(av.is_available() for av in g.availability),
             g.price.min,
             g.price.max,
-            g.backend.value,
+            get_or_error(g.backend).value,
             g.region,
             g.count.min,
             g.name,
             g.memory_mib,
         ),
     )
-
-
-async def list_gpus_grouped(
-    session: AsyncSession,
-    project: ProjectModel,
-    run_spec: RunSpec,
-    group_by: Optional[List[Literal["backend", "region", "count"]]] = None,
-) -> ListGpusResponse:
-    """Retrieves available GPU specifications based on a run spec, with optional grouping."""
-    offers = await _get_gpu_offers(session, project, run_spec)
-    backend_gpus = _process_offers_into_backend_gpus(offers)
-
-    group_by_set = set(group_by) if group_by else set()
-
-    if "region" in group_by_set and "backend" not in group_by_set:
-        from dstack._internal.core.errors import ServerClientError
-
-        raise ServerClientError("Cannot group by 'region' without also grouping by 'backend'")
-
-    # Determine grouping strategy based on combination
-    has_backend = "backend" in group_by_set
-    has_region = "region" in group_by_set
-    has_count = "count" in group_by_set
-    if has_backend and has_region and has_count:
-        gpus = _get_gpus_grouped_by_backend_region_and_count(backend_gpus)
-    elif has_backend and has_count:
-        gpus = _get_gpus_grouped_by_backend_and_count(backend_gpus)
-    elif has_backend and has_region:
-        gpus = _get_gpus_grouped_by_backend_and_region(backend_gpus)
-    elif has_backend:
-        gpus = _get_gpus_grouped_by_backend(backend_gpus)
-    elif has_count:
-        gpus = _get_gpus_grouped_by_count(backend_gpus)
-    else:
-        gpus = _get_gpus_with_no_grouping(backend_gpus)
-
-    return ListGpusResponse(gpus=gpus)

dstack/_internal/server/services/instances.py

@@ -39,6 +39,7 @@ from dstack._internal.core.models.profiles import (
 from dstack._internal.core.models.runs import JobProvisioningData, Requirements
 from dstack._internal.core.models.volumes import Volume
 from dstack._internal.core.services.profiles import get_termination
+from dstack._internal.server import settings as server_settings
 from dstack._internal.server.models import (
     FleetModel,
     InstanceHealthCheckModel,
@@ -47,9 +48,11 @@ from dstack._internal.server.models import (
     UserModel,
 )
 from dstack._internal.server.schemas.health.dcgm import DCGMHealthResponse
-from dstack._internal.server.schemas.runner import InstanceHealthResponse
+from dstack._internal.server.schemas.runner import InstanceHealthResponse, TaskStatus
+from dstack._internal.server.services.logging import fmt
 from dstack._internal.server.services.offers import generate_shared_offer
 from dstack._internal.server.services.projects import list_user_project_models
+from dstack._internal.server.services.runner.client import ShimClient
 from dstack._internal.utils import common as common_utils
 from dstack._internal.utils.logging import get_logger
 
@@ -633,3 +636,40 @@ async def create_ssh_instance_model(
         busy_blocks=0,
     )
     return im
+
+
+def remove_dangling_tasks_from_instance(shim_client: ShimClient, instance: InstanceModel) -> None:
+    if not shim_client.is_api_v2_supported():
+        return
+    assigned_to_instance_job_ids = {str(j.id) for j in instance.jobs}
+    task_list_response = shim_client.list_tasks()
+    tasks: list[tuple[str, Optional[TaskStatus]]]
+    if task_list_response.tasks is not None:
+        tasks = [(t.id, t.status) for t in task_list_response.tasks]
+    elif task_list_response.ids is not None:
+        # compatibility with pre-0.19.26 shim
+        tasks = [(t_id, None) for t_id in task_list_response.ids]
+    else:
+        raise ValueError("Unexpected task list response, neither `tasks` nor `ids` is set")
+    for task_id, task_status in tasks:
+        if task_id in assigned_to_instance_job_ids:
+            continue
+        should_terminate = task_status != TaskStatus.TERMINATED
+        should_remove = not server_settings.SERVER_KEEP_SHIM_TASKS
+        if not (should_terminate or should_remove):
+            continue
+        logger.warning(
+            "%s: dangling task found, id=%s, status=%s. Terminating and/or removing",
+            fmt(instance),
+            task_id,
+            task_status or "<unknown>",
+        )
+        if should_terminate:
+            shim_client.terminate_task(
+                task_id=task_id,
+                reason=None,
+                message=None,
+                timeout=0,
+            )
+        if should_remove:
+            shim_client.remove_task(task_id=task_id)

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

@@ -256,7 +256,16 @@ async def process_terminating_job(
     if jpd is not None:
         logger.debug("%s: stopping container", fmt(job_model))
         ssh_private_keys = get_instance_ssh_private_keys(instance_model)
-        await stop_container(job_model, jpd, ssh_private_keys)
+        if not await stop_container(job_model, jpd, ssh_private_keys):
+            # The dangling container can be removed later during instance processing
+            logger.warning(
+                (
+                    "%s: could not stop container, possibly due to a communication error."
+                    " See debug logs for details."
+                    " Ignoring, can attempt to remove the container later"
+                ),
+                fmt(job_model),
+            )
         if jrd is not None and jrd.volume_names is not None:
             volume_names = jrd.volume_names
         else:
@@ -378,21 +387,22 @@ async def stop_container(
     job_model: JobModel,
     job_provisioning_data: JobProvisioningData,
     ssh_private_keys: tuple[str, Optional[str]],
-):
+) -> bool:
     if job_provisioning_data.dockerized:
         # send a request to the shim to terminate the docker container
         # SSHError and RequestException are caught in the `runner_ssh_tunner` decorator
-        await run_async(
+        return await run_async(
             _shim_submit_stop,
             ssh_private_keys,
             job_provisioning_data,
             None,
             job_model,
         )
+    return True
 
 
 @runner_ssh_tunnel(ports=[DSTACK_SHIM_HTTP_PORT])
-def _shim_submit_stop(ports: Dict[int, int], job_model: JobModel):
+def _shim_submit_stop(ports: Dict[int, int], job_model: JobModel) -> bool:
     shim_client = client.ShimClient(port=ports[DSTACK_SHIM_HTTP_PORT])
 
     resp = shim_client.healthcheck()
@@ -418,6 +428,7 @@ def _shim_submit_stop(ports: Dict[int, int], job_model: JobModel):
             shim_client.remove_task(task_id=job_model.id)
     else:
         shim_client.stop(force=True)
+    return True
 
 
 def group_jobs_by_replica_latest(jobs: List[JobModel]) -> Iterable[Tuple[int, List[JobModel]]]:

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

@@ -3,7 +3,7 @@ import sys
 import threading
 from abc import ABC, abstractmethod
 from pathlib import PurePosixPath
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional
 
 from cachetools import TTLCache, cached
 
@@ -179,6 +179,7 @@ class JobConfigurator(ABC):
 
     async def _commands(self) -> List[str]:
         if self.run_spec.configuration.entrypoint is not None:  # docker-like format
+            assert self.run_spec.configuration.type != "dev-environment"
             entrypoint = shlex.split(self.run_spec.configuration.entrypoint)
             commands = self.run_spec.configuration.commands
         elif shell_commands := self._shell_commands():
@@ -258,19 +259,17 @@ class JobConfigurator(ABC):
         return self.run_spec.configuration.single_branch
 
     def _max_duration(self) -> Optional[int]:
-        if self.run_spec.merged_profile.max_duration in [None, True]:
+        if self.run_spec.merged_profile.max_duration is None:
             return self._default_max_duration()
-        if self.run_spec.merged_profile.max_duration in ["off", False]:
+        if self.run_spec.merged_profile.max_duration == "off":
             return None
-        # pydantic validator ensures this is int
         return self.run_spec.merged_profile.max_duration
 
     def _stop_duration(self) -> Optional[int]:
-        if self.run_spec.merged_profile.stop_duration in [None, True]:
+        if self.run_spec.merged_profile.stop_duration is None:
             return DEFAULT_STOP_DURATION
-        if self.run_spec.merged_profile.stop_duration in ["off", False]:
+        if self.run_spec.merged_profile.stop_duration == "off":
             return None
-        # pydantic validator ensures this is int
         return self.run_spec.merged_profile.stop_duration
 
     def _utilization_policy(self) -> Optional[UtilizationPolicy]:
@@ -328,7 +327,7 @@ class JobConfigurator(ABC):
 
 
 def interpolate_job_volumes(
-    run_volumes: List[Union[MountPoint, str]],
+    run_volumes: List[MountPoint],
     job_num: int,
 ) -> List[MountPoint]:
     if len(run_volumes) == 0:
@@ -343,9 +342,6 @@ def interpolate_job_volumes(
     )
     job_volumes = []
     for mount_point in run_volumes:
-        if isinstance(mount_point, str):
-            # pydantic validator ensures strings are converted to MountPoint
-            continue
         if not isinstance(mount_point, VolumeMountPoint):
             job_volumes.append(mount_point.copy())
             continue

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

@@ -18,6 +18,8 @@ class DevEnvironmentJobConfigurator(JobConfigurator):
     TYPE: RunConfigurationType = RunConfigurationType.DEV_ENVIRONMENT
 
     def __init__(self, run_spec: RunSpec, secrets: Dict[str, str]):
+        assert run_spec.configuration.type == "dev-environment"
+
         if run_spec.configuration.ide == "vscode":
             __class = VSCodeDesktop
         elif run_spec.configuration.ide == "cursor":
@@ -32,6 +34,8 @@ class DevEnvironmentJobConfigurator(JobConfigurator):
         super().__init__(run_spec=run_spec, secrets=secrets)
 
     def _shell_commands(self) -> List[str]:
+        assert self.run_spec.configuration.type == "dev-environment"
+
         commands = self.ide.get_install_commands()
         commands.append(INSTALL_IPYKERNEL)
         commands += self.run_spec.configuration.setup
@@ -56,4 +60,5 @@ class DevEnvironmentJobConfigurator(JobConfigurator):
         return self.run_spec.merged_profile.spot_policy or SpotPolicy.ONDEMAND
 
     def _ports(self) -> List[PortMapping]:
+        assert self.run_spec.configuration.type == "dev-environment"
         return self.run_spec.configuration.ports

dstack/_internal/server/services/jobs/configurators/extensions/cursor.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional
 
 from dstack._internal.core.models.configurations import DEFAULT_REPO_DIR
 
@@ -6,8 +6,8 @@ from dstack._internal.core.models.configurations import DEFAULT_REPO_DIR
 class CursorDesktop:
     def __init__(
         self,
-        run_name: str,
-        version: str,
+        run_name: Optional[str],
+        version: Optional[str],
         extensions: List[str],
     ):
         self.run_name = run_name

dstack/_internal/server/services/jobs/configurators/extensions/vscode.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional
 
 from dstack._internal.core.models.configurations import DEFAULT_REPO_DIR
 
@@ -6,8 +6,8 @@ from dstack._internal.core.models.configurations import DEFAULT_REPO_DIR
 class VSCodeDesktop:
     def __init__(
         self,
-        run_name: str,
-        version: str,
+        run_name: Optional[str],
+        version: Optional[str],
         extensions: List[str],
     ):
         self.run_name = run_name

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

@@ -9,6 +9,7 @@ class ServiceJobConfigurator(JobConfigurator):
     TYPE: RunConfigurationType = RunConfigurationType.SERVICE
 
     def _shell_commands(self) -> List[str]:
+        assert self.run_spec.configuration.type == "service"
         return self.run_spec.configuration.commands
 
     def _default_single_branch(self) -> bool:

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

@@ -10,6 +10,7 @@ class TaskJobConfigurator(JobConfigurator):
     TYPE: RunConfigurationType = RunConfigurationType.TASK
 
     async def get_job_specs(self, replica_num: int) -> List[JobSpec]:
+        assert self.run_spec.configuration.type == "task"
         job_specs = []
         for job_num in range(self.run_spec.configuration.nodes):
             job_spec = await self._get_job_spec(
@@ -21,6 +22,7 @@ class TaskJobConfigurator(JobConfigurator):
         return job_specs
 
     def _shell_commands(self) -> List[str]:
+        assert self.run_spec.configuration.type == "task"
         return self.run_spec.configuration.commands
 
     def _default_single_branch(self) -> bool:
@@ -33,6 +35,7 @@ class TaskJobConfigurator(JobConfigurator):
         return self.run_spec.merged_profile.spot_policy or SpotPolicy.ONDEMAND
 
     def _ports(self) -> List[PortMapping]:
+        assert self.run_spec.configuration.type == "task"
         return self.run_spec.configuration.ports
 
     def _working_dir(self) -> Optional[str]:

dstack/_internal/server/services/locking.py

@@ -23,13 +23,13 @@ T = TypeVar("T")
 
 
 class Lockset(Protocol[T]):
-    def __contains__(self, item: T) -> bool: ...
+    def __contains__(self, item: T, /) -> bool: ...
     def __iter__(self) -> Iterator[T]: ...
     def __len__(self) -> int: ...
-    def add(self, item: T) -> None: ...
-    def discard(self, item: T) -> None: ...
-    def update(self, other: Iterable[T]) -> None: ...
-    def difference_update(self, other: Iterable[T]) -> None: ...
+    def add(self, item: T, /) -> None: ...
+    def discard(self, item: T, /) -> None: ...
+    def update(self, other: Iterable[T], /) -> None: ...
+    def difference_update(self, other: Iterable[T], /) -> None: ...
 
 
 class ResourceLocker:

dstack/_internal/server/services/logging.py

@@ -1,14 +1,22 @@
 from typing import Union
 
-from dstack._internal.server.models import GatewayModel, JobModel, ProbeModel, RunModel
+from dstack._internal.server.models import (
+    GatewayModel,
+    InstanceModel,
+    JobModel,
+    ProbeModel,
+    RunModel,
+)
 
 
-def fmt(model: Union[RunModel, JobModel, GatewayModel, ProbeModel]) -> str:
+def fmt(model: Union[RunModel, JobModel, InstanceModel, GatewayModel, ProbeModel]) -> str:
     """Consistent string representation of a model for logging."""
     if isinstance(model, RunModel):
         return f"run({model.id.hex[:6]}){model.run_name}"
     if isinstance(model, JobModel):
         return f"job({model.id.hex[:6]}){model.job_name}"
+    if isinstance(model, InstanceModel):
+        return f"instance({model.id.hex[:6]}){model.name}"
     if isinstance(model, GatewayModel):
         return f"gateway({model.id.hex[:6]}){model.name}"
     if isinstance(model, ProbeModel):

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

@@ -7,14 +7,14 @@ from dstack._internal.server import settings
 from dstack._internal.server.models import ProjectModel
 from dstack._internal.server.schemas.logs import PollLogsRequest
 from dstack._internal.server.schemas.runner import LogEvent as RunnerLogEvent
-from dstack._internal.server.services.logs.aws import BOTO_AVAILABLE, CloudWatchLogStorage
+from dstack._internal.server.services.logs import aws as aws_logs
+from dstack._internal.server.services.logs import gcp as gcp_logs
 from dstack._internal.server.services.logs.base import (
     LogStorage,
     LogStorageError,
     b64encode_raw_message,
 )
 from dstack._internal.server.services.logs.filelog import FileLogStorage
-from dstack._internal.server.services.logs.gcp import GCP_LOGGING_AVAILABLE, GCPLogStorage
 from dstack._internal.utils.common import run_async
 from dstack._internal.utils.logging import get_logger
 
@@ -29,9 +29,9 @@ def get_log_storage() -> LogStorage:
     if _log_storage is not None:
         return _log_storage
     if settings.SERVER_CLOUDWATCH_LOG_GROUP:
-        if BOTO_AVAILABLE:
+        if aws_logs.BOTO_AVAILABLE:
             try:
-                _log_storage = CloudWatchLogStorage(
+                _log_storage = aws_logs.CloudWatchLogStorage(
                     group=settings.SERVER_CLOUDWATCH_LOG_GROUP,
                     region=settings.SERVER_CLOUDWATCH_LOG_REGION,
                 )
@@ -44,9 +44,11 @@ def get_log_storage() -> LogStorage:
         else:
             logger.error("Cannot use CloudWatch Logs storage: boto3 is not installed")
     elif settings.SERVER_GCP_LOGGING_PROJECT:
-        if GCP_LOGGING_AVAILABLE:
+        if gcp_logs.GCP_LOGGING_AVAILABLE:
             try:
-                _log_storage = GCPLogStorage(project_id=settings.SERVER_GCP_LOGGING_PROJECT)
+                _log_storage = gcp_logs.GCPLogStorage(
+                    project_id=settings.SERVER_GCP_LOGGING_PROJECT
+                )
             except LogStorageError as e:
                 logger.error("Failed to initialize GCP Logs storage: %s", e)
             except Exception: