dstack 0.0.9__py3-none-any.whl → 0.20.7__py3-none-any.whl

dstack/_internal/server/services/offers.py

@@ -0,0 +1,249 @@
+import itertools
+from collections.abc import Iterable, Iterator
+from typing import List, Literal, Optional, Tuple, Union
+
+import gpuhunt
+
+from dstack._internal.core.backends.base.backend import Backend
+from dstack._internal.core.backends.base.compute import ComputeWithPlacementGroupSupport
+from dstack._internal.core.backends.features import (
+    BACKENDS_WITH_CREATE_INSTANCE_SUPPORT,
+    BACKENDS_WITH_MULTINODE_SUPPORT,
+    BACKENDS_WITH_PRIVILEGED_SUPPORT,
+    BACKENDS_WITH_RESERVATION_SUPPORT,
+)
+from dstack._internal.core.models.backends.base import BackendType
+from dstack._internal.core.models.instances import (
+    InstanceOfferWithAvailability,
+    InstanceType,
+    Resources,
+)
+from dstack._internal.core.models.placement import PlacementGroup
+from dstack._internal.core.models.profiles import Profile
+from dstack._internal.core.models.runs import JobProvisioningData, Requirements
+from dstack._internal.core.models.volumes import Volume
+from dstack._internal.server.models import ProjectModel
+from dstack._internal.server.services import backends as backends_services
+
+
+async def get_offers_by_requirements(
+    project: ProjectModel,
+    profile: Profile,
+    requirements: Requirements,
+    exclude_not_available=False,
+    multinode: bool = False,
+    master_job_provisioning_data: Optional[JobProvisioningData] = None,
+    volumes: Optional[List[List[Volume]]] = None,
+    privileged: bool = False,
+    instance_mounts: bool = False,
+    placement_group: Optional[PlacementGroup] = None,
+    blocks: Union[int, Literal["auto"]] = 1,
+    max_offers: Optional[int] = None,
+) -> List[Tuple[Backend, InstanceOfferWithAvailability]]:
+    backends: List[Backend] = await backends_services.get_project_backends(project=project)
+
+    backend_types = profile.backends
+    regions = profile.regions
+    availability_zones = profile.availability_zones
+    instance_types = profile.instance_types
+
+    if volumes:
+        mount_point_volumes = volumes[0]
+        volumes_backend_types = [v.configuration.backend for v in mount_point_volumes]
+        if backend_types is None:
+            backend_types = volumes_backend_types
+        backend_types = [b for b in backend_types if b in volumes_backend_types]
+        volumes_regions = [v.configuration.region for v in mount_point_volumes]
+        if regions is None:
+            regions = volumes_regions
+        regions = [r for r in regions if r in volumes_regions]
+
+    if multinode:
+        if backend_types is None:
+            backend_types = BACKENDS_WITH_MULTINODE_SUPPORT
+        backend_types = [b for b in backend_types if b in BACKENDS_WITH_MULTINODE_SUPPORT]
+
+    if privileged:
+        if backend_types is None:
+            backend_types = BACKENDS_WITH_PRIVILEGED_SUPPORT
+        backend_types = [b for b in backend_types if b in BACKENDS_WITH_PRIVILEGED_SUPPORT]
+
+    if instance_mounts:
+        if backend_types is None:
+            backend_types = BACKENDS_WITH_CREATE_INSTANCE_SUPPORT
+        backend_types = [b for b in backend_types if b in BACKENDS_WITH_CREATE_INSTANCE_SUPPORT]
+
+    if profile.reservation is not None:
+        if backend_types is None:
+            backend_types = BACKENDS_WITH_RESERVATION_SUPPORT
+        backend_types = [b for b in backend_types if b in BACKENDS_WITH_RESERVATION_SUPPORT]
+
+    # For multi-node, restrict backend and region.
+    # The default behavior is to provision all nodes in the same backend and region.
+    if master_job_provisioning_data is not None:
+        if backend_types is None:
+            backend_types = [master_job_provisioning_data.get_base_backend()]
+        if regions is None:
+            regions = [master_job_provisioning_data.region]
+        backend_types = [
+            b for b in backend_types if b == master_job_provisioning_data.get_base_backend()
+        ]
+        regions = [r for r in regions if r == master_job_provisioning_data.region]
+
+    if backend_types is not None:
+        backends = [b for b in backends if b.TYPE in backend_types or b.TYPE == BackendType.DSTACK]
+
+    offers = await backends_services.get_backend_offers(
+        backends=backends,
+        requirements=requirements,
+        exclude_not_available=exclude_not_available,
+    )
+
+    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:
+        offers = _get_shareable_offers(offers, blocks)
+
+    if max_offers is not None:
+        offers = itertools.islice(offers, max_offers)
+
+    # Put NOT_AVAILABLE and NO_QUOTA offers at the end.
+    # We have to do this after taking max_offers to avoid processing all offers
+    # if all/most offers are unavailable.
+    return sorted(offers, key=lambda i: not i[1].availability.is_available())
+
+
+def is_divisible_into_blocks(
+    cpu_count: int, gpu_count: int, blocks: Union[int, Literal["auto"]]
+) -> tuple[bool, int]:
+    """
+    Returns `True` and number of blocks the instance can be split into or `False` and `0` if
+    is not divisible.
+    Requested number of blocks can be `auto`, which means as many as possible.
+    """
+    if blocks == "auto":
+        if gpu_count == 0:
+            blocks = cpu_count
+        else:
+            blocks = min(cpu_count, gpu_count)
+    if blocks < 1 or cpu_count % blocks or gpu_count % blocks:
+        return False, 0
+    return True, blocks
+
+
+def generate_shared_offer(
+    offer: InstanceOfferWithAvailability, blocks: int, total_blocks: int
+) -> InstanceOfferWithAvailability:
+    full_resources = offer.instance.resources
+    resources = Resources(
+        cpus=full_resources.cpus // total_blocks * blocks,
+        memory_mib=full_resources.memory_mib // total_blocks * blocks,
+        gpus=full_resources.gpus[: len(full_resources.gpus) // total_blocks * blocks],
+        spot=full_resources.spot,
+        disk=full_resources.disk,
+        description=full_resources.description,
+    )
+    return InstanceOfferWithAvailability(
+        backend=offer.backend,
+        instance=InstanceType(
+            name=offer.instance.name,
+            resources=resources,
+        ),
+        region=offer.region,
+        price=offer.price,
+        backend_data=offer.backend_data,
+        availability=offer.availability,
+        blocks=blocks,
+        total_blocks=total_blocks,
+    )
+
+
+def get_instance_offer_with_restricted_az(
+    instance_offer: InstanceOfferWithAvailability,
+    master_job_provisioning_data: Optional[JobProvisioningData],
+) -> InstanceOfferWithAvailability:
+    instance_offer = instance_offer.copy()
+    if (
+        master_job_provisioning_data is not None
+        and master_job_provisioning_data.availability_zone is not None
+    ):
+        if instance_offer.availability_zones is None:
+            instance_offer.availability_zones = [master_job_provisioning_data.availability_zone]
+        instance_offer.availability_zones = [
+            z
+            for z in instance_offer.availability_zones
+            if z == master_job_provisioning_data.availability_zone
+        ]
+    return instance_offer
+
+
+def _filter_offers(
+    offers: Iterable[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,
+) -> Iterator[Tuple[Backend, InstanceOfferWithAvailability]]:
+    """
+    Yields filtered offers. May return modified offers to match the filters.
+    """
+    if regions is not None:
+        regions = [r.lower() for r in regions]
+    if instance_types is not None:
+        instance_types = [i.lower() for i in instance_types]
+
+    for b, offer in offers:
+        if backend_types is not None and offer.backend not in backend_types:
+            continue
+        if regions is not None and offer.region.lower() not in regions:
+            continue
+        if instance_types is not None and offer.instance.name.lower() not in instance_types:
+            continue
+        if placement_group is not None:
+            compute = b.compute()
+            if not isinstance(
+                compute, ComputeWithPlacementGroupSupport
+            ) or not compute.is_suitable_placement_group(placement_group, offer):
+                continue
+        if availability_zones is not None:
+            if offer.availability_zones is None:
+                continue
+            new_offer = offer.copy()
+            new_offer.availability_zones = [
+                z for z in offer.availability_zones if z in availability_zones
+            ]
+            if not new_offer.availability_zones:
+                continue
+            offer = new_offer
+        yield (b, offer)
+
+
+def _get_shareable_offers(
+    offers: Iterable[Tuple[Backend, InstanceOfferWithAvailability]],
+    blocks: Union[int, Literal["auto"]],
+) -> Iterator[Tuple[Backend, InstanceOfferWithAvailability]]:
+    """
+    Yields offers that can be shared with `total_blocks` set.
+    """
+    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, total_blocks = is_divisible_into_blocks(cpu_count, gpu_count, blocks)
+        if not divisible:
+            continue
+        new_offer = offer.copy()
+        new_offer.total_blocks = total_blocks
+        yield (backend, new_offer)

dstack/_internal/server/services/permissions.py

@@ -0,0 +1,37 @@
+from pydantic import Field
+from typing_extensions import Annotated
+
+from dstack._internal.core.models.common import CoreModel
+
+
+class DefaultPermissions(CoreModel):
+    allow_non_admins_create_projects: Annotated[
+        bool,
+        Field(
+            description=(
+                "This flag controls whether regular users (non-global admins)"
+                " can create and manage their own projects"
+            )
+        ),
+    ] = True
+    allow_non_admins_manage_ssh_fleets: Annotated[
+        bool,
+        Field(
+            description=(
+                "This flag controls whether regular project members (i.e. Users)"
+                " can add and delete SSH fleets"
+            )
+        ),
+    ] = True
+
+
+_default_permissions = DefaultPermissions()
+
+
+def set_default_permissions(default_permissions: DefaultPermissions):
+    global _default_permissions
+    _default_permissions = default_permissions
+
+
+def get_default_permissions() -> DefaultPermissions:
+    return _default_permissions

dstack/_internal/server/services/placement.py

@@ -0,0 +1,234 @@
+from collections.abc import Iterable
+from typing import Optional
+from uuid import UUID
+
+from git import List
+from sqlalchemy import and_, select, update
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from dstack._internal.core.backends.base.compute import (
+    ComputeWithPlacementGroupSupport,
+    generate_unique_placement_group_name,
+)
+from dstack._internal.core.errors import BackendError, PlacementGroupNotSupportedError
+from dstack._internal.core.models.instances import InstanceOffer
+from dstack._internal.core.models.placement import (
+    PlacementGroup,
+    PlacementGroupConfiguration,
+    PlacementGroupProvisioningData,
+    PlacementStrategy,
+)
+from dstack._internal.server.models import FleetModel, InstanceModel, PlacementGroupModel
+from dstack._internal.utils.common import run_async
+from dstack._internal.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+def placement_group_model_to_placement_group(
+    placement_group_model: PlacementGroupModel,
+) -> PlacementGroup:
+    configuration = get_placement_group_configuration(placement_group_model)
+    provisioning_data = get_placement_group_provisioning_data(placement_group_model)
+    return PlacementGroup(
+        name=placement_group_model.name,
+        project_name=placement_group_model.project.name,
+        configuration=configuration,
+        provisioning_data=provisioning_data,
+    )
+
+
+def placement_group_model_to_placement_group_optional(
+    placement_group_model: Optional[PlacementGroupModel],
+) -> Optional[PlacementGroup]:
+    if placement_group_model is None:
+        return None
+    return placement_group_model_to_placement_group(placement_group_model)
+
+
+def get_placement_group_configuration(
+    placement_group_model: PlacementGroupModel,
+) -> PlacementGroupConfiguration:
+    return PlacementGroupConfiguration.__response__.parse_raw(placement_group_model.configuration)
+
+
+def get_placement_group_provisioning_data(
+    placement_group_model: PlacementGroupModel,
+) -> Optional[PlacementGroupProvisioningData]:
+    if placement_group_model.provisioning_data is None:
+        return None
+    return PlacementGroupProvisioningData.__response__.parse_raw(
+        placement_group_model.provisioning_data
+    )
+
+
+async def get_fleet_placement_group_models(
+    session: AsyncSession,
+    fleet_id: Optional[UUID],
+) -> List[PlacementGroupModel]:
+    if fleet_id is None:
+        return []
+    res = await session.execute(
+        select(PlacementGroupModel).where(
+            and_(
+                PlacementGroupModel.fleet_id == fleet_id,
+                PlacementGroupModel.deleted == False,
+                PlacementGroupModel.fleet_deleted == False,
+            )
+        )
+    )
+    return list(res.scalars().all())
+
+
+async def schedule_fleet_placement_groups_deletion(
+    session: AsyncSession, fleet_id: UUID, except_placement_group_ids: Iterable[UUID] = ()
+):
+    await session.execute(
+        update(PlacementGroupModel)
+        .where(
+            and_(
+                PlacementGroupModel.fleet_id == fleet_id,
+                PlacementGroupModel.id.not_in(except_placement_group_ids),
+            )
+        )
+        .values(fleet_deleted=True)  # TODO: rename `fleet_deleted` -> `to_be_deleted`
+    )
+
+
+def get_placement_group_model_for_instance(
+    placement_group_models: list[PlacementGroupModel],
+    instance_model: InstanceModel,
+    master_instance_model: InstanceModel,
+) -> Optional[PlacementGroupModel]:
+    placement_group_model = None
+    if instance_model.id != master_instance_model.id:
+        if placement_group_models:
+            placement_group_model = placement_group_models[0]
+        if len(placement_group_models) > 1:
+            logger.error(
+                (
+                    "Expected 0 or 1 placement groups associated with fleet %s, found %s."
+                    " An incorrect placement group might have been selected for instance %s"
+                ),
+                instance_model.fleet_id,
+                len(placement_group_models),
+                instance_model.name,
+            )
+    return placement_group_model
+
+
+def get_placement_group_model_for_job(
+    placement_group_models: list[PlacementGroupModel],
+    fleet_model: Optional[FleetModel],
+) -> Optional[PlacementGroupModel]:
+    """
+    Returns any fleet placement group for jobs that provision
+    in non-empty fleets and `None` for empty fleets.
+    This is so that only the first job creates placement groups.
+    """
+    placement_group_model = None
+    active_instances = []
+    if fleet_model is not None:
+        active_instances = [i for i in fleet_model.instances if not i.deleted]
+    if len(active_instances) > 0 and len(placement_group_models) > 0:
+        placement_group_model = placement_group_models[0]
+    return placement_group_model
+
+
+async def find_or_create_suitable_placement_group(
+    fleet_model: FleetModel,
+    placement_groups: List[PlacementGroupModel],
+    instance_offer: InstanceOffer,
+    compute: ComputeWithPlacementGroupSupport,
+) -> Optional[PlacementGroupModel]:
+    placement_group_model = find_suitable_placement_group(
+        placement_groups=placement_groups,
+        instance_offer=instance_offer,
+        compute=compute,
+    )
+    if placement_group_model is None:
+        placement_group_model = await create_placement_group(
+            fleet_model=fleet_model,
+            master_instance_offer=instance_offer,
+            compute=compute,
+        )
+    return placement_group_model
+
+
+def find_suitable_placement_group(
+    placement_groups: List[PlacementGroupModel],
+    instance_offer: InstanceOffer,
+    compute: ComputeWithPlacementGroupSupport,
+) -> Optional[PlacementGroupModel]:
+    for pg in placement_groups:
+        if compute.is_suitable_placement_group(
+            placement_group_model_to_placement_group(pg), instance_offer
+        ):
+            return pg
+    return None
+
+
+async def create_placement_group(
+    fleet_model: FleetModel,
+    master_instance_offer: InstanceOffer,
+    compute: ComputeWithPlacementGroupSupport,
+) -> Optional[PlacementGroupModel]:
+    placement_group_model = PlacementGroupModel(
+        # TODO: generate the name in Compute.create_placement_group to allow
+        # backend-specific name length limits
+        name=generate_unique_placement_group_name(
+            project_name=fleet_model.project.name,
+            fleet_name=fleet_model.name,
+        ),
+        project=fleet_model.project,
+        fleet=fleet_model,
+        configuration=PlacementGroupConfiguration(
+            backend=master_instance_offer.backend,
+            region=master_instance_offer.region,
+            placement_strategy=PlacementStrategy.CLUSTER,
+        ).json(),
+    )
+    placement_group = placement_group_model_to_placement_group(placement_group_model)
+    logger.debug(
+        "Creating placement group %s in %s/%s",
+        placement_group.name,
+        placement_group.configuration.backend.value,
+        placement_group.configuration.region,
+    )
+    try:
+        pgpd = await run_async(
+            compute.create_placement_group,
+            placement_group_model_to_placement_group(placement_group_model),
+            master_instance_offer,
+        )
+    except PlacementGroupNotSupportedError:
+        logger.debug(
+            "Skipping offer %s because placement group not supported",
+            master_instance_offer.instance.name,
+        )
+        return None
+    except BackendError as e:
+        logger.warning(
+            "Failed to create placement group %s in %s/%s: %r",
+            placement_group.name,
+            placement_group.configuration.backend.value,
+            placement_group.configuration.region,
+            e,
+        )
+        return None
+    except Exception:
+        logger.exception(
+            "Got exception when creating placement group %s in %s/%s",
+            placement_group.name,
+            placement_group.configuration.backend.value,
+            placement_group.configuration.region,
+        )
+        return None
+    logger.info(
+        "Created placement group %s in %s/%s",
+        placement_group.name,
+        placement_group.configuration.backend.value,
+        placement_group.configuration.region,
+    )
+    placement_group_model.provisioning_data = pgpd.json()
+    return placement_group_model

dstack/_internal/server/services/plugins.py

@@ -0,0 +1,109 @@
+import itertools
+from importlib import import_module
+from typing import Dict
+
+from backports.entry_points_selectable import entry_points  # backport for Python 3.9
+
+from dstack._internal.core.errors import ServerClientError
+from dstack._internal.utils.common import run_async
+from dstack._internal.utils.logging import get_logger
+from dstack.plugins import ApplyPolicy, ApplySpec, Plugin
+
+logger = get_logger(__name__)
+
+
+_PLUGINS: list[Plugin] = []
+
+_BUILTIN_PLUGINS: Dict[str, str] = {"rest_plugin": "dstack.plugins.builtin.rest_plugin:RESTPlugin"}
+
+
+class PluginEntrypoint:
+    def __init__(self, name: str, import_path: str, is_builtin: bool = False):
+        self.name = name
+        self.import_path = import_path
+        self.is_builtin = is_builtin
+
+    def load(self):
+        module_path, _, class_name = self.import_path.partition(":")
+        try:
+            module = import_module(module_path)
+            plugin_class = getattr(module, class_name, None)
+            if plugin_class is None:
+                logger.warning(
+                    ("Failed to load plugin %s: plugin class %s not found in module %s."),
+                    self.name,
+                    class_name,
+                    module_path,
+                )
+                return None
+            if not issubclass(plugin_class, Plugin):
+                logger.warning(
+                    ("Failed to load plugin %s: plugin class %s is not a subclass of Plugin."),
+                    self.name,
+                    class_name,
+                )
+                return None
+            return plugin_class()
+        except ImportError:
+            logger.warning(
+                (
+                    "Failed to load plugin %s when importing %s."
+                    " Ensure the module is on the import path."
+                ),
+                self.name,
+                self.import_path,
+            )
+            return None
+
+
+def load_plugins(enabled_plugins: list[str]):
+    _PLUGINS.clear()
+    entrypoints: dict[str, PluginEntrypoint] = {}
+    plugins_to_load = enabled_plugins.copy()
+    for entrypoint in entry_points(group="dstack.plugins"):  # type: ignore[call-arg]
+        if entrypoint.name not in enabled_plugins:
+            logger.info(
+                ("Found not enabled plugin %s. Plugin will not be loaded."),
+                entrypoint.name,
+            )
+            continue
+        else:
+            entrypoints[entrypoint.name] = PluginEntrypoint(
+                entrypoint.name, entrypoint.value, is_builtin=False
+            )
+
+    for name, import_path in _BUILTIN_PLUGINS.items():
+        if name not in enabled_plugins:
+            logger.debug(
+                ("Found not enabled builtin plugin %s. Plugin will not be loaded."),
+                name,
+            )
+        else:
+            entrypoints[name] = PluginEntrypoint(name, import_path, is_builtin=True)
+
+    for plugin_name, plugin_entrypoint in entrypoints.items():
+        plugin_instance = plugin_entrypoint.load()
+        if plugin_instance is not None:
+            _PLUGINS.append(plugin_instance)
+            plugins_to_load.remove(plugin_name)
+            logger.info("Loaded plugin %s", plugin_name)
+
+    if plugins_to_load:
+        logger.warning("Enabled plugins not found: %s", plugins_to_load)
+
+
+async def apply_plugin_policies(user: str, project: str, spec: ApplySpec) -> ApplySpec:
+    policies = _get_apply_policies()
+    for policy in policies:
+        try:
+            spec = await run_async(policy.on_apply, user=user, project=project, spec=spec)
+        except ValueError as e:
+            msg = None
+            if len(e.args) > 0:
+                msg = e.args[0]
+            raise ServerClientError(msg)
+    return spec
+
+
+def _get_apply_policies() -> list[ApplyPolicy]:
+    return list(itertools.chain(*[p.get_apply_policies() for p in _PLUGINS]))

dstack/_internal/server/services/probes.py

@@ -0,0 +1,10 @@
+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