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

dstack/_internal/server/schemas/gateways.py

@@ -14,7 +14,7 @@ class CreateGatewayRequest(CoreModel):
     backend_type: Annotated[Optional[BackendType], Field(exclude=True)] = None
     region: Annotated[Optional[str], Field(exclude=True)] = None
 
-    class Config:
+    class Config(CoreModel.Config):
         @staticmethod
         def schema_extra(schema: Dict[str, Any]) -> None:
             del schema["properties"]["name"]

dstack/_internal/server/schemas/gpus.py

@@ -0,0 +1,66 @@
+from typing import List, Literal, Optional
+
+import gpuhunt
+from pydantic import Field
+
+from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.common import CoreModel
+from dstack._internal.core.models.instances import InstanceAvailability
+from dstack._internal.core.models.resources import Range
+from dstack._internal.core.models.runs import RunSpec
+
+
+class BackendGpu(CoreModel):
+    """GPU specification from a backend offer."""
+
+    name: str
+    memory_mib: int
+    vendor: gpuhunt.AcceleratorVendor
+    availability: InstanceAvailability
+    spot: bool
+    count: int
+    price: float
+    region: str
+
+
+class BackendGpus(CoreModel):
+    """Backend GPU specifications."""
+
+    backend_type: BackendType
+    gpus: List[BackendGpu]
+    regions: List[str]
+
+
+class ListGpusRequest(CoreModel):
+    """Request for listing GPUs with optional grouping."""
+
+    run_spec: RunSpec
+    group_by: Optional[List[Literal["backend", "region", "count"]]] = Field(
+        default=None,
+        description="List of fields to group by. Valid values: 'backend', 'region', 'count'. "
+        "Note: 'region' can only be used together with 'backend'.",
+    )
+
+
+class GpuGroup(CoreModel):
+    """GPU group that can handle all grouping scenarios."""
+
+    name: str
+    memory_mib: int
+    vendor: gpuhunt.AcceleratorVendor
+    availability: List[InstanceAvailability]
+    spot: List[Literal["spot", "on-demand"]]
+    count: Range[int]
+    price: Range[float]
+    backends: Optional[List[BackendType]] = None
+    backend: Optional[BackendType] = None
+    regions: Optional[List[str]] = None
+    region: Optional[str] = None
+
+
+class ListGpusResponse(CoreModel):
+    """Response containing GPU specifications."""
+
+    gpus: List[GpuGroup] = Field(
+        description="List of GPU specifications, grouped according to the group_by parameter"
+    )

dstack/_internal/server/services/docker.py

@@ -32,7 +32,7 @@ class DXFAuthAdapter:
 
 
 class DockerImage(CoreModel):
-    class Config:
+    class Config(CoreModel.Config):
         frozen = True
 
     image: str

dstack/_internal/server/services/gpus.py

@@ -0,0 +1,390 @@
+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.models.instances import InstanceOfferWithAvailability
+from dstack._internal.core.models.profiles import SpotPolicy
+from dstack._internal.core.models.resources import Range
+from dstack._internal.core.models.runs import Requirements, RunSpec, get_policy_map
+from dstack._internal.server.models import ProjectModel
+from dstack._internal.server.schemas.gpus import (
+    BackendGpu,
+    BackendGpus,
+    GpuGroup,
+    ListGpusResponse,
+)
+from dstack._internal.server.services.offers import get_offers_by_requirements
+
+
+async def _get_gpu_offers(
+    session: AsyncSession, 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
+    requirements = Requirements(
+        resources=run_spec.configuration.resources,
+        max_price=profile.max_price,
+        spot=get_policy_map(profile.spot_policy, default=SpotPolicy.AUTO),
+        reservation=profile.reservation,
+    )
+
+    return await get_offers_by_requirements(
+        project=project,
+        profile=profile,
+        requirements=requirements,
+        exclude_not_available=False,
+        multinode=False,
+        volumes=None,
+        privileged=False,
+        instance_mounts=False,
+    )
+
+
+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] = {}
+
+    for backend, offer in offers:
+        backend_type = backend.TYPE
+        if backend_type not in backend_data:
+            backend_data[backend_type] = {"gpus": {}, "regions": set()}
+
+        backend_data[backend_type]["regions"].add(offer.region)
+
+        if not offer.instance.resources.gpus:
+            continue
+
+        gpu_types_in_offer = {}
+        for gpu in offer.instance.resources.gpus:
+            gpu_type_key = (gpu.name, gpu.memory_mib, gpu.vendor)
+            if gpu_type_key not in gpu_types_in_offer:
+                gpu_types_in_offer[gpu_type_key] = 0
+            gpu_types_in_offer[gpu_type_key] += 1
+
+        for (
+            gpu_name,
+            gpu_memory_mib,
+            gpu_vendor,
+        ), gpu_count_in_offer in gpu_types_in_offer.items():
+            instance_config_key = (
+                gpu_name,
+                gpu_memory_mib,
+                gpu_vendor,
+                gpu_count_in_offer,
+                offer.instance.resources.spot,
+                offer.region,
+            )
+
+            if instance_config_key not in backend_data[backend_type]["gpus"]:
+                backend_data[backend_type]["gpus"][instance_config_key] = BackendGpu(
+                    name=gpu_name,
+                    memory_mib=gpu_memory_mib,
+                    vendor=gpu_vendor,
+                    availability=offer.availability,
+                    spot=offer.instance.resources.spot,
+                    count=gpu_count_in_offer,
+                    price=offer.price,
+                    region=offer.region,
+                )
+
+    backend_gpus_list = []
+    for backend_type, data in backend_data.items():
+        gpus_list = sorted(
+            list(data["gpus"].values()),
+            key=lambda g: (
+                not g.availability.is_available(),
+                g.vendor.value,
+                g.name,
+                g.memory_mib,
+            ),
+        )
+        backend_gpus_list.append(
+            BackendGpus(
+                backend_type=backend_type,
+                gpus=gpus_list,
+                regions=sorted(list(data["regions"])),
+            )
+        )
+    return backend_gpus_list
+
+
+def _update_gpu_group(row: GpuGroup, gpu: BackendGpu, backend_type: str):
+    """Updates an existing GpuGroup with new data from another GPU offer."""
+    spot_type: Literal["spot", "on-demand"] = "spot" if gpu.spot else "on-demand"
+
+    if gpu.availability not in row.availability:
+        row.availability.append(gpu.availability)
+    if spot_type not in row.spot:
+        row.spot.append(spot_type)
+    if row.backends and backend_type not in row.backends:
+        row.backends.append(backend_type)
+
+    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
+    row.price.min = min(row.price.min, per_gpu_price)
+    row.price.max = max(row.price.max, per_gpu_price)
+
+
+def _get_gpus_with_no_grouping(backend_gpus: List[BackendGpus]) -> List[GpuGroup]:
+    """Aggregates GPU specs into a flat list, without any grouping."""
+    gpu_rows: Dict[Tuple, GpuGroup] = {}
+    for backend in backend_gpus:
+        for gpu in backend.gpus:
+            key = (gpu.name, gpu.memory_mib, gpu.vendor)
+            if key not in gpu_rows:
+                per_gpu_price = gpu.price / gpu.count
+                price_range = Range[float](min=per_gpu_price, max=per_gpu_price)
+
+                gpu_rows[key] = GpuGroup(
+                    name=gpu.name,
+                    memory_mib=gpu.memory_mib,
+                    vendor=gpu.vendor,
+                    availability=[gpu.availability],
+                    spot=["spot" if gpu.spot else "on-demand"],
+                    count=Range[int](min=gpu.count, max=gpu.count),
+                    price=price_range,
+                    backends=[backend.backend_type],
+                )
+            else:
+                _update_gpu_group(gpu_rows[key], gpu, backend.backend_type)
+
+    result = sorted(
+        list(gpu_rows.values()),
+        key=lambda g: (
+            not any(av.is_available() for av in g.availability),
+            g.price.min,
+            g.price.max,
+            g.name,
+            g.memory_mib,
+        ),
+    )
+
+    return result
+
+
+def _get_gpus_grouped_by_backend(backend_gpus: List[BackendGpus]) -> List[GpuGroup]:
+    """Aggregates GPU specs, grouping them by backend."""
+    gpu_rows: Dict[Tuple, GpuGroup] = {}
+    for backend in backend_gpus:
+        for gpu in backend.gpus:
+            key = (gpu.name, gpu.memory_mib, gpu.vendor, backend.backend_type)
+            if key not in gpu_rows:
+                per_gpu_price = gpu.price / gpu.count
+                gpu_rows[key] = GpuGroup(
+                    name=gpu.name,
+                    memory_mib=gpu.memory_mib,
+                    vendor=gpu.vendor,
+                    availability=[gpu.availability],
+                    spot=["spot" if gpu.spot else "on-demand"],
+                    count=Range[int](min=gpu.count, max=gpu.count),
+                    price=Range[float](min=per_gpu_price, max=per_gpu_price),
+                    backend=backend.backend_type,
+                    regions=backend.regions.copy(),
+                )
+            else:
+                _update_gpu_group(gpu_rows[key], gpu, backend.backend_type)
+
+    return sorted(
+        list(gpu_rows.values()),
+        key=lambda g: (
+            not any(av.is_available() for av in g.availability),
+            g.price.min,
+            g.price.max,
+            g.backend.value,
+            g.name,
+            g.memory_mib,
+        ),
+    )
+
+
+def _get_gpus_grouped_by_backend_and_region(backend_gpus: List[BackendGpus]) -> List[GpuGroup]:
+    """Aggregates GPU specs, grouping them by both backend and region."""
+    gpu_rows: Dict[Tuple, GpuGroup] = {}
+    for backend in backend_gpus:
+        for gpu in backend.gpus:
+            key = (gpu.name, gpu.memory_mib, gpu.vendor, backend.backend_type, gpu.region)
+            if key not in gpu_rows:
+                per_gpu_price = gpu.price / gpu.count
+                gpu_rows[key] = GpuGroup(
+                    name=gpu.name,
+                    memory_mib=gpu.memory_mib,
+                    vendor=gpu.vendor,
+                    availability=[gpu.availability],
+                    spot=["spot" if gpu.spot else "on-demand"],
+                    count=Range[int](min=gpu.count, max=gpu.count),
+                    price=Range[float](min=per_gpu_price, max=per_gpu_price),
+                    backend=backend.backend_type,
+                    region=gpu.region,
+                )
+            else:
+                _update_gpu_group(gpu_rows[key], gpu, backend.backend_type)
+
+    return sorted(
+        list(gpu_rows.values()),
+        key=lambda g: (
+            not any(av.is_available() for av in g.availability),
+            g.price.min,
+            g.price.max,
+            g.backend.value,
+            g.region,
+            g.name,
+            g.memory_mib,
+        ),
+    )
+
+
+def _get_gpus_grouped_by_count(backend_gpus: List[BackendGpus]) -> List[GpuGroup]:
+    """Aggregates GPU specs, grouping them by GPU count."""
+    gpu_rows: Dict[Tuple, GpuGroup] = {}
+    for backend in backend_gpus:
+        for gpu in backend.gpus:
+            key = (gpu.name, gpu.memory_mib, gpu.vendor, gpu.count)
+            if key not in gpu_rows:
+                per_gpu_price = gpu.price / gpu.count
+                gpu_rows[key] = GpuGroup(
+                    name=gpu.name,
+                    memory_mib=gpu.memory_mib,
+                    vendor=gpu.vendor,
+                    availability=[gpu.availability],
+                    spot=["spot" if gpu.spot else "on-demand"],
+                    count=Range[int](min=gpu.count, max=gpu.count),
+                    price=Range[float](min=per_gpu_price, max=per_gpu_price),
+                    backends=[backend.backend_type],
+                )
+            else:
+                _update_gpu_group(gpu_rows[key], gpu, backend.backend_type)
+
+    return sorted(
+        list(gpu_rows.values()),
+        key=lambda g: (
+            not any(av.is_available() for av in g.availability),
+            g.price.min,
+            g.price.max,
+            g.count.min,
+            g.name,
+            g.memory_mib,
+        ),
+    )
+
+
+def _get_gpus_grouped_by_backend_and_count(backend_gpus: List[BackendGpus]) -> List[GpuGroup]:
+    """Aggregates GPU specs, grouping them by backend and GPU count."""
+    gpu_rows: Dict[Tuple, GpuGroup] = {}
+    for backend in backend_gpus:
+        for gpu in backend.gpus:
+            key = (gpu.name, gpu.memory_mib, gpu.vendor, backend.backend_type, gpu.count)
+            if key not in gpu_rows:
+                per_gpu_price = gpu.price / gpu.count
+                gpu_rows[key] = GpuGroup(
+                    name=gpu.name,
+                    memory_mib=gpu.memory_mib,
+                    vendor=gpu.vendor,
+                    availability=[gpu.availability],
+                    spot=["spot" if gpu.spot else "on-demand"],
+                    count=Range[int](min=gpu.count, max=gpu.count),
+                    price=Range[float](min=per_gpu_price, max=per_gpu_price),
+                    backend=backend.backend_type,
+                    regions=backend.regions.copy(),
+                )
+            else:
+                _update_gpu_group(gpu_rows[key], gpu, backend.backend_type)
+
+    return sorted(
+        list(gpu_rows.values()),
+        key=lambda g: (
+            not any(av.is_available() for av in g.availability),
+            g.price.min,
+            g.price.max,
+            g.backend.value,
+            g.count.min,
+            g.name,
+            g.memory_mib,
+        ),
+    )
+
+
+def _get_gpus_grouped_by_backend_region_and_count(
+    backend_gpus: List[BackendGpus],
+) -> List[GpuGroup]:
+    """Aggregates GPU specs, grouping them by backend, region, and GPU count."""
+    gpu_rows: Dict[Tuple, GpuGroup] = {}
+    for backend in backend_gpus:
+        for gpu in backend.gpus:
+            key = (
+                gpu.name,
+                gpu.memory_mib,
+                gpu.vendor,
+                backend.backend_type,
+                gpu.region,
+                gpu.count,
+            )
+            if key not in gpu_rows:
+                per_gpu_price = gpu.price / gpu.count
+                gpu_rows[key] = GpuGroup(
+                    name=gpu.name,
+                    memory_mib=gpu.memory_mib,
+                    vendor=gpu.vendor,
+                    availability=[gpu.availability],
+                    spot=["spot" if gpu.spot else "on-demand"],
+                    count=Range[int](min=gpu.count, max=gpu.count),
+                    price=Range[float](min=per_gpu_price, max=per_gpu_price),
+                    backend=backend.backend_type,
+                    region=gpu.region,
+                )
+            else:
+                _update_gpu_group(gpu_rows[key], gpu, backend.backend_type)
+
+    return sorted(
+        list(gpu_rows.values()),
+        key=lambda g: (
+            not any(av.is_available() for av in g.availability),
+            g.price.min,
+            g.price.max,
+            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/jobs/__init__.py

@@ -152,7 +152,9 @@ def job_model_to_job_submission(
         inactivity_secs=job_model.inactivity_secs,
         status=job_model.status,
         status_message=status_message,
-        termination_reason=job_model.termination_reason,
+        termination_reason=job_model.termination_reason.value
+        if job_model.termination_reason
+        else None,
         termination_reason_message=job_model.termination_reason_message,
         exit_status=job_model.exit_status,
         job_provisioning_data=job_provisioning_data,

dstack/_internal/server/services/offers.py

@@ -49,6 +49,7 @@ async def get_offers_by_requirements(
     backend_types = profile.backends
     regions = profile.regions
     availability_zones = profile.availability_zones
+    instance_types = profile.instance_types
 
     if volumes:
         mount_point_volumes = volumes[0]
@@ -97,9 +98,43 @@ async def get_offers_by_requirements(
         exclude_not_available=exclude_not_available,
     )
 
-    # Filter offers again for backends since a backend
-    # can return offers of different backend types (e.g. BackendType.DSTACK).
-    # The first filter should remain as an optimization.
+    offers = filter_offers(
+        offers=offers,
+        # Double filtering by backends if backend returns offers for other backend.
+        backend_types=backend_types,
+        regions=regions,
+        availability_zones=availability_zones,
+        instance_types=instance_types,
+        placement_group=placement_group,
+    )
+
+    if blocks == 1:
+        return offers
+
+    shareable_offers = []
+    for backend, offer in offers:
+        resources = offer.instance.resources
+        cpu_count = resources.cpus
+        gpu_count = len(resources.gpus)
+        if gpu_count > 0 and resources.gpus[0].vendor == gpuhunt.AcceleratorVendor.GOOGLE:
+            # TPUs cannot be shared
+            gpu_count = 1
+        divisible, _blocks = is_divisible_into_blocks(cpu_count, gpu_count, blocks)
+        if not divisible:
+            continue
+        offer.total_blocks = _blocks
+        shareable_offers.append((backend, offer))
+    return shareable_offers
+
+
+def filter_offers(
+    offers: List[Tuple[Backend, InstanceOfferWithAvailability]],
+    backend_types: Optional[List[BackendType]] = None,
+    regions: Optional[List[str]] = None,
+    availability_zones: Optional[List[str]] = None,
+    instance_types: Optional[List[str]] = None,
+    placement_group: Optional[PlacementGroup] = None,
+) -> List[Tuple[Backend, InstanceOfferWithAvailability]]:
     if backend_types is not None:
         offers = [(b, o) for b, o in offers if o.backend in backend_types]
 
@@ -119,39 +154,21 @@ async def get_offers_by_requirements(
                     new_offers.append((b, new_offer))
         offers = new_offers
 
+    if instance_types is not None:
+        instance_types = [i.lower() for i in instance_types]
+        offers = [(b, o) for b, o in offers if o.instance.name.lower() in instance_types]
+
     if placement_group is not None:
         new_offers = []
         for b, o in offers:
-            for backend in backends:
-                compute = backend.compute()
-                if isinstance(
-                    compute, ComputeWithPlacementGroupSupport
-                ) and compute.is_suitable_placement_group(placement_group, o):
-                    new_offers.append((b, o))
-                    break
+            compute = b.compute()
+            if isinstance(
+                compute, ComputeWithPlacementGroupSupport
+            ) and compute.is_suitable_placement_group(placement_group, o):
+                new_offers.append((b, o))
         offers = new_offers
 
-    if profile.instance_types is not None:
-        instance_types = [i.lower() for i in profile.instance_types]
-        offers = [(b, o) for b, o in offers if o.instance.name.lower() in instance_types]
-
-    if blocks == 1:
-        return offers
-
-    shareable_offers = []
-    for backend, offer in offers:
-        resources = offer.instance.resources
-        cpu_count = resources.cpus
-        gpu_count = len(resources.gpus)
-        if gpu_count > 0 and resources.gpus[0].vendor == gpuhunt.AcceleratorVendor.GOOGLE:
-            # TPUs cannot be shared
-            gpu_count = 1
-        divisible, _blocks = is_divisible_into_blocks(cpu_count, gpu_count, blocks)
-        if not divisible:
-            continue
-        offer.total_blocks = _blocks
-        shareable_offers.append((backend, offer))
-    return shareable_offers
+    return offers
 
 
 def is_divisible_into_blocks(

dstack/_internal/server/services/probes.py

@@ -1,6 +1,10 @@
-from dstack._internal.core.models.runs import Probe
+from dstack._internal.core.models.runs import Probe, ProbeSpec
 from dstack._internal.server.models import ProbeModel
 
 
 def probe_model_to_probe(probe_model: ProbeModel) -> Probe:
     return Probe(success_streak=probe_model.success_streak)
+
+
+def is_probe_ready(probe: ProbeModel, spec: ProbeSpec) -> bool:
+    return probe.success_streak >= spec.ready_after

dstack/_internal/server/services/proxy/repo.py

@@ -54,6 +54,7 @@ class ServerProxyRepo(BaseProxyRepo):
                 RunModel.gateway_id.is_(None),
                 JobModel.run_name == run_name,
                 JobModel.status == JobStatus.RUNNING,
+                JobModel.registered == True,
                 JobModel.job_num == 0,
             )
             .options(

dstack/_internal/server/services/repos.py

@@ -129,7 +129,7 @@ async def create_repo(
     repo = RepoModel(
         project_id=project.id,
         name=repo_id,
-        type=repo_info.repo_type,
+        type=RepoType(repo_info.repo_type),
         info=repo_info.json(),
     )
     try: