dstack 0.19.24__py3-none-any.whl → 0.19.25__py3-none-any.whl

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/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/runs.py

@@ -41,6 +41,7 @@ from dstack._internal.core.models.runs import (
     JobStatus,
     JobSubmission,
     JobTerminationReason,
+    ProbeSpec,
     Run,
     RunPlan,
     RunSpec,
@@ -58,6 +59,7 @@ from dstack._internal.server import settings
 from dstack._internal.server.db import get_db
 from dstack._internal.server.models import (
     JobModel,
+    ProbeModel,
     ProjectModel,
     RepoModel,
     RunModel,
@@ -86,6 +88,7 @@ from dstack._internal.server.services.locking import get_locker, string_to_lock_
 from dstack._internal.server.services.logging import fmt
 from dstack._internal.server.services.offers import get_offers_by_requirements
 from dstack._internal.server.services.plugins import apply_plugin_policies
+from dstack._internal.server.services.probes import is_probe_ready
 from dstack._internal.server.services.projects import list_user_project_models
 from dstack._internal.server.services.resources import set_resources_defaults
 from dstack._internal.server.services.secrets import get_project_secrets_mapping
@@ -1185,8 +1188,8 @@ async def scale_run_replicas(session: AsyncSession, run_model: RunModel, replica
         elif {JobStatus.PROVISIONING, JobStatus.PULLING} & statuses:
             # if there are any provisioning or pulling jobs, the replica is active and has the importance of 1
             active_replicas.append((1, is_out_of_date, replica_num, replica_jobs))
-        elif not is_replica_ready(replica_jobs):
-            # all jobs are running, but probes are failing, the replica is active and has the importance of 2
+        elif not is_replica_registered(replica_jobs):
+            # all jobs are running, but not receiving traffic, the replica is active and has the importance of 2
             active_replicas.append((2, is_out_of_date, replica_num, replica_jobs))
         else:
             # all jobs are running and ready, the replica is active and has the importance of 3
@@ -1273,15 +1276,13 @@ async def retry_run_replica_jobs(
         session.add(new_job_model)
 
 
-def is_replica_ready(jobs: Iterable[JobModel]) -> bool:
-    if not all(job.status == JobStatus.RUNNING for job in jobs):
-        return False
-    for job in jobs:
-        job_spec: JobSpec = JobSpec.__response__.parse_raw(job.job_spec_data)
-        for probe_spec, probe in zip(job_spec.probes, job.probes):
-            if probe.success_streak < probe_spec.ready_after:
-                return False
-    return True
+def is_job_ready(probes: Iterable[ProbeModel], probe_specs: Iterable[ProbeSpec]) -> bool:
+    return all(is_probe_ready(probe, probe_spec) for probe, probe_spec in zip(probes, probe_specs))
+
+
+def is_replica_registered(jobs: list[JobModel]) -> bool:
+    # Only job_num=0 is supposed to receive service requests
+    return jobs[0].registered
 
 
 def _remove_job_spec_sensitive_info(spec: JobSpec):