dstack 0.19.7__py3-none-any.whl → 0.19.8__py3-none-any.whl

dstack/_internal/core/models/resources.py

@@ -1,8 +1,9 @@
 import math
+from collections.abc import Mapping
 from typing import Any, Dict, Generic, List, Optional, Tuple, TypeVar, Union
 
 import gpuhunt
-from pydantic import Field, root_validator, validator
+from pydantic import Field, parse_obj_as, root_validator, validator
 from pydantic.generics import GenericModel
 from typing_extensions import Annotated
 
@@ -128,6 +129,67 @@ DEFAULT_MEMORY_SIZE = Range[Memory](min=Memory.parse("8GB"))
 DEFAULT_GPU_COUNT = Range[int](min=1, max=1)
 
 
+class CPUSpec(CoreModel):
+    class Config:
+        @staticmethod
+        def schema_extra(schema: Dict[str, Any]):
+            add_extra_schema_types(
+                schema["properties"]["count"],
+                extra_types=[{"type": "integer"}, {"type": "string"}],
+            )
+
+    arch: Annotated[
+        Optional[gpuhunt.CPUArchitecture],
+        Field(description="The CPU architecture, one of: `x86`, `arm`"),
+    ] = None
+    count: Annotated[Range[int], Field(description="The number of CPU cores")] = DEFAULT_CPU_COUNT
+
+    @classmethod
+    def __get_validators__(cls):
+        yield cls.parse
+        yield cls.validate
+
+    @classmethod
+    def parse(cls, v: Any) -> Any:
+        if isinstance(v, int):
+            v = str(v)
+        if isinstance(v, str):
+            tokens = v.replace(" ", "").split(":")
+            spec = {}
+            for token in tokens:
+                if not token:
+                    raise ValueError(f"CPU spec contains empty token: {v}")
+                if ".." in token or token.isdigit():
+                    if "count" in spec:
+                        raise ValueError(f"CPU spec count conflict: {v}")
+                    spec["count"] = token
+                else:
+                    try:
+                        arch = gpuhunt.CPUArchitecture.cast(token)
+                    except ValueError:
+                        raise ValueError(f"Invalid CPU architecture: {v}")
+                    if "arch" in spec:
+                        raise ValueError(f"CPU spec arch conflict: {v}")
+                    spec["arch"] = arch
+            return spec
+        # Range and min/max dict - for backward compatibility
+        if isinstance(v, Range):
+            return {"arch": None, "count": v}
+        if isinstance(v, Mapping) and v.keys() == {"min", "max"}:
+            return {"arch": None, "count": v}
+        return v
+
+    @validator("arch", pre=True)
+    def _validate_arch(cls, v: Any) -> Any:
+        if v is None:
+            return None
+        if isinstance(v, gpuhunt.CPUArchitecture):
+            return v
+        if isinstance(v, str):
+            return gpuhunt.CPUArchitecture.cast(v)
+        return v
+
+
 class GPUSpec(CoreModel):
     class Config:
         @staticmethod
@@ -302,7 +364,10 @@ class ResourcesSpec(CoreModel):
                 extra_types=[{"type": "integer"}, {"type": "string"}],
             )
 
-    cpu: Annotated[Range[int], Field(description="The number of CPU cores")] = DEFAULT_CPU_COUNT
+    # TODO: Remove Range[int] in 0.20. Range[int] for backward compatibility only.
+    cpu: Annotated[Union[CPUSpec, Range[int]], Field(description="The CPU requirements")] = (
+        CPUSpec()
+    )
     memory: Annotated[Range[Memory], Field(description="The RAM size (e.g., `8GB`)")] = (
         DEFAULT_MEMORY_SIZE
     )
@@ -317,8 +382,18 @@ class ResourcesSpec(CoreModel):
     gpu: Annotated[Optional[GPUSpec], Field(description="The GPU requirements")] = None
     disk: Annotated[Optional[DiskSpec], Field(description="The disk resources")] = DEFAULT_DISK
 
+    # TODO: Remove in 0.20. Added for backward compatibility.
+    @root_validator
+    def _post_validate(cls, values):
+        cpu = values.get("cpu")
+        if isinstance(cpu, CPUSpec) and cpu.arch in [None, gpuhunt.CPUArchitecture.X86]:
+            values["cpu"] = cpu.count
+        return values
+
     def pretty_format(self) -> str:
-        resources: Dict[str, Any] = dict(cpus=self.cpu, memory=self.memory)
+        # TODO: Remove in 0.20. Use self.cpu directly
+        cpu = parse_obj_as(CPUSpec, self.cpu)
+        resources: Dict[str, Any] = dict(cpu_arch=cpu.arch, cpus=cpu.count, memory=self.memory)
         if self.gpu:
             gpu = self.gpu
             resources.update(

dstack/_internal/core/models/runs.py

@@ -439,9 +439,14 @@ class Run(CoreModel):
 
     @root_validator
     def _error(cls, values) -> Dict:
+        try:
+            termination_reason = values["termination_reason"]
+            jobs = values["jobs"]
+        except KeyError:
+            return values
         values["error"] = _get_run_error(
-            run_termination_reason=values["termination_reason"],
-            run_jobs=values["jobs"],
+            run_termination_reason=termination_reason,
+            run_jobs=jobs,
         )
         return values
 

dstack/_internal/core/models/volumes.py

@@ -159,7 +159,7 @@ class VolumeMountPoint(CoreModel):
             description=(
                 "The network volume name or the list of network volume names to mount."
                 " If a list is specified, one of the volumes in the list will be mounted."
-                " Specify volumes from different backends/regions to increase availability."
+                " Specify volumes from different backends/regions to increase availability"
             )
         ),
     ]

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

@@ -1,15 +1,16 @@
-from sqlalchemy import select, update
+from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import joinedload
 
 from dstack._internal.core.models.fleets import FleetStatus
 from dstack._internal.server.db import get_session_ctx
-from dstack._internal.server.models import FleetModel, PlacementGroupModel
+from dstack._internal.server.models import FleetModel
 from dstack._internal.server.services.fleets import (
     is_fleet_empty,
     is_fleet_in_use,
 )
 from dstack._internal.server.services.locking import get_locker
+from dstack._internal.server.services.placement import schedule_fleet_placement_groups_deletion
 from dstack._internal.utils.common import get_current_datetime
 from dstack._internal.utils.logging import get_logger
 
@@ -68,16 +69,6 @@ async def _autodelete_fleet(session: AsyncSession, fleet_model: FleetModel):
     fleet_model.status = FleetStatus.TERMINATED
     fleet_model.deleted = True
     fleet_model.last_processed_at = get_current_datetime()
-    await _mark_placement_groups_as_ready_for_deletion(session=session, fleet_model=fleet_model)
+    await schedule_fleet_placement_groups_deletion(session=session, fleet_id=fleet_model.id)
     await session.commit()
     logger.info("Fleet %s deleted", fleet_model.name)
-
-
-async def _mark_placement_groups_as_ready_for_deletion(
-    session: AsyncSession, fleet_model: FleetModel
-):
-    await session.execute(
-        update(PlacementGroupModel)
-        .where(PlacementGroupModel.fleet_id == fleet_model.id)
-        .values(fleet_deleted=True)
-    )

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

@@ -19,6 +19,8 @@ from dstack._internal.core.backends import (
 from dstack._internal.core.backends.base.compute import (
     ComputeWithCreateInstanceSupport,
     ComputeWithPlacementGroupSupport,
+    GoArchType,
+    generate_unique_placement_group_name,
     get_dstack_runner_binary_path,
     get_dstack_shim_binary_path,
     get_dstack_working_dir,
@@ -26,6 +28,7 @@ from dstack._internal.core.backends.base.compute import (
     get_shim_pre_start_commands,
 )
 from dstack._internal.core.backends.remote.provisioning import (
+    detect_cpu_arch,
     get_host_info,
     get_paramiko_connection,
     get_shim_healthcheck,
@@ -39,11 +42,16 @@ from dstack._internal.core.backends.remote.provisioning import (
 from dstack._internal.core.consts import DSTACK_SHIM_HTTP_PORT
 
 # FIXME: ProvisioningError is a subclass of ComputeError and should not be used outside of Compute
-from dstack._internal.core.errors import BackendError, NotYetTerminated, ProvisioningError
+from dstack._internal.core.errors import (
+    BackendError,
+    NotYetTerminated,
+    ProvisioningError,
+)
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.fleets import InstanceGroupPlacement
 from dstack._internal.core.models.instances import (
     InstanceAvailability,
+    InstanceOffer,
     InstanceOfferWithAvailability,
     InstanceRuntime,
     InstanceStatus,
@@ -51,7 +59,6 @@ from dstack._internal.core.models.instances import (
     SSHKey,
 )
 from dstack._internal.core.models.placement import (
-    PlacementGroup,
     PlacementGroupConfiguration,
     PlacementStrategy,
 )
@@ -89,8 +96,9 @@ from dstack._internal.server.services.instances import (
 from dstack._internal.server.services.locking import get_locker
 from dstack._internal.server.services.offers import is_divisible_into_blocks
 from dstack._internal.server.services.placement import (
-    get_fleet_placement_groups,
+    get_fleet_placement_group_models,
     placement_group_model_to_placement_group,
+    schedule_fleet_placement_groups_deletion,
 )
 from dstack._internal.server.services.runner import client as runner_client
 from dstack._internal.server.services.runner.client import HealthStatus
@@ -264,7 +272,7 @@ async def _add_remote(instance: InstanceModel) -> None:
             )
             deploy_timeout = 20 * 60  # 20 minutes
             result = await asyncio.wait_for(future, timeout=deploy_timeout)
-            health, host_info = result
+            health, host_info, cpu_arch = result
         except (asyncio.TimeoutError, TimeoutError) as e:
             raise ProvisioningError(f"Deploy timeout: {e}") from e
         except Exception as e:
@@ -285,7 +293,7 @@ async def _add_remote(instance: InstanceModel) -> None:
         instance.last_retry_at = get_current_datetime()
         return
 
-    instance_type = host_info_to_instance_type(host_info)
+    instance_type = host_info_to_instance_type(host_info, cpu_arch)
     instance_network = None
     internal_ip = None
     try:
@@ -388,7 +396,7 @@ def _deploy_instance(
     pkeys: List[PKey],
     ssh_proxy_pkeys: Optional[list[PKey]],
     authorized_keys: List[str],
-) -> Tuple[HealthStatus, Dict[str, Any]]:
+) -> Tuple[HealthStatus, Dict[str, Any], GoArchType]:
     with get_paramiko_connection(
         remote_details.ssh_user,
         remote_details.host,
@@ -399,13 +407,16 @@ def _deploy_instance(
     ) as client:
         logger.info(f"Connected to {remote_details.ssh_user} {remote_details.host}")
 
+        arch = detect_cpu_arch(client)
+        logger.info("%s: CPU arch is %s", remote_details.host, arch)
+
         # Execute pre start commands
-        shim_pre_start_commands = get_shim_pre_start_commands()
+        shim_pre_start_commands = get_shim_pre_start_commands(arch=arch)
         run_pre_start_commands(client, shim_pre_start_commands, authorized_keys)
         logger.debug("The script for installing dstack has been executed")
 
         # Upload envs
-        shim_envs = get_shim_env(authorized_keys)
+        shim_envs = get_shim_env(authorized_keys, arch=arch)
         try:
             fleet_configuration_envs = remote_details.env.as_dict()
         except ValueError as e:
@@ -440,7 +451,7 @@ def _deploy_instance(
             raise ProvisioningError("Cannot read HealthcheckResponse") from e
         health = runner_client.health_response_to_health_status(health_response)
 
-        return health, host_info
+        return health, host_info, arch
 
 
 async def _create_instance(session: AsyncSession, instance: InstanceModel) -> None:
@@ -509,11 +520,39 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
             )
             return
 
+    placement_group_models = []
+    placement_group_model = None
+    if instance.fleet_id:
+        placement_group_models = await get_fleet_placement_group_models(
+            session=session,
+            fleet_id=instance.fleet_id,
+        )
+        # The placement group is determined when provisioning the master instance
+        # and used for all other instances in the fleet.
+        if not _is_fleet_master_instance(instance):
+            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.fleet_id,
+                    len(placement_group_models),
+                    instance.name,
+                )
+
     offers = await get_create_instance_offers(
         project=instance.project,
         profile=profile,
         requirements=requirements,
         fleet_model=instance.fleet,
+        placement_group=(
+            placement_group_model_to_placement_group(placement_group_model)
+            if placement_group_model
+            else None
+        ),
         blocks="auto" if instance.total_blocks is None else instance.total_blocks,
         exclude_not_available=True,
     )
@@ -527,12 +566,6 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
         )
         return
 
-    placement_groups = []
-    if instance.fleet_id:
-        placement_groups = await get_fleet_placement_groups(
-            session=session, fleet_id=instance.fleet_id
-        )
-
     # Limit number of offers tried to prevent long-running processing
     # in case all offers fail.
     for backend, instance_offer in offers[: server_settings.MAX_OFFERS_TRIED]:
@@ -542,25 +575,28 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
         assert isinstance(compute, ComputeWithCreateInstanceSupport)
         instance_offer = _get_instance_offer_for_instance(instance_offer, instance)
         if (
-            instance_offer.backend in BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT
+            _is_fleet_master_instance(instance)
+            and instance_offer.backend in BACKENDS_WITH_PLACEMENT_GROUPS_SUPPORT
             and instance.fleet
-            and instance_configuration.placement_group_name
+            and _is_cloud_cluster(instance.fleet)
         ):
             assert isinstance(compute, ComputeWithPlacementGroupSupport)
-            placement_group_model = _create_placement_group_if_does_not_exist(
-                session=session,
-                fleet_model=instance.fleet,
-                placement_groups=placement_groups,
-                name=instance_configuration.placement_group_name,
-                backend=instance_offer.backend,
-                region=instance_offer.region,
+            placement_group_model = _find_suitable_placement_group(
+                placement_groups=placement_group_models,
+                instance_offer=instance_offer,
+                compute=compute,
             )
-            if placement_group_model is not None:
-                placement_group = placement_group_model_to_placement_group(placement_group_model)
-                pgpd = await run_async(compute.create_placement_group, placement_group)
-                placement_group_model.provisioning_data = pgpd.json()
+            if placement_group_model is None:
+                placement_group_model = await _create_placement_group(
+                    fleet_model=instance.fleet,
+                    master_instance_offer=instance_offer,
+                    compute=compute,
+                )
+                if placement_group_model is None:  # error occurred
+                    continue
                 session.add(placement_group_model)
-                placement_groups.append(placement_group)
+                await session.flush()
+                placement_group_models.append(placement_group_model)
         logger.debug(
             "Trying %s in %s/%s for $%0.4f per hour",
             instance_offer.instance.name,
@@ -573,6 +609,11 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
                 compute.create_instance,
                 instance_offer,
                 instance_configuration,
+                (
+                    placement_group_model_to_placement_group(placement_group_model)
+                    if placement_group_model
+                    else None
+                ),
             )
         except BackendError as e:
             logger.warning(
@@ -612,22 +653,46 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
                 "instance_status": InstanceStatus.PROVISIONING.value,
             },
         )
+        if instance.fleet_id and _is_fleet_master_instance(instance):
+            # Clean up placement groups that did not end up being used
+            await schedule_fleet_placement_groups_deletion(
+                session=session,
+                fleet_id=instance.fleet_id,
+                except_placement_group_ids=(
+                    [placement_group_model.id] if placement_group_model is not None else []
+                ),
+            )
         return
 
     instance.last_retry_at = get_current_datetime()
 
     if not should_retry:
-        instance.status = InstanceStatus.TERMINATED
-        instance.termination_reason = "All offers failed" if offers else "No offers found"
-        logger.info(
-            "Terminated instance %s: %s",
-            instance.name,
-            instance.termination_reason,
-            extra={
-                "instance_name": instance.name,
-                "instance_status": InstanceStatus.TERMINATED.value,
-            },
-        )
+        _mark_terminated(instance, "All offers failed" if offers else "No offers found")
+        if (
+            instance.fleet
+            and _is_fleet_master_instance(instance)
+            and _is_cloud_cluster(instance.fleet)
+        ):
+            # Do not attempt to deploy other instances, as they won't determine the correct cluster
+            # backend, region, and placement group without a successfully deployed master instance
+            for sibling_instance in instance.fleet.instances:
+                if sibling_instance.id == instance.id:
+                    continue
+                _mark_terminated(sibling_instance, "Master instance failed to start")
+
+
+def _mark_terminated(instance: InstanceModel, termination_reason: str) -> None:
+    instance.status = InstanceStatus.TERMINATED
+    instance.termination_reason = termination_reason
+    logger.info(
+        "Terminated instance %s: %s",
+        instance.name,
+        instance.termination_reason,
+        extra={
+            "instance_name": instance.name,
+            "instance_status": InstanceStatus.TERMINATED.value,
+        },
+    )
 
 
 async def _check_instance(instance: InstanceModel) -> None:
@@ -906,12 +971,20 @@ def _need_to_wait_fleet_provisioning(instance: InstanceModel) -> bool:
     if instance.fleet is None:
         return False
     if (
-        instance.id == instance.fleet.instances[0].id
+        _is_fleet_master_instance(instance)
         or instance.fleet.instances[0].job_provisioning_data is not None
         or instance.fleet.instances[0].status == InstanceStatus.TERMINATED
     ):
         return False
-    fleet = fleet_model_to_fleet(instance.fleet)
+    return _is_cloud_cluster(instance.fleet)
+
+
+def _is_fleet_master_instance(instance: InstanceModel) -> bool:
+    return instance.fleet is not None and instance.id == instance.fleet.instances[0].id
+
+
+def _is_cloud_cluster(fleet_model: FleetModel) -> bool:
+    fleet = fleet_model_to_fleet(fleet_model)
     return (
         fleet.spec.configuration.placement == InstanceGroupPlacement.CLUSTER
         and fleet.spec.configuration.ssh_config is None
@@ -944,28 +1017,76 @@ def _get_instance_offer_for_instance(
     return instance_offer
 
 
-def _create_placement_group_if_does_not_exist(
-    session: AsyncSession,
-    fleet_model: FleetModel,
-    placement_groups: List[PlacementGroup],
-    name: str,
-    backend: BackendType,
-    region: str,
+def _find_suitable_placement_group(
+    placement_groups: List[PlacementGroupModel],
+    instance_offer: InstanceOffer,
+    compute: ComputeWithPlacementGroupSupport,
 ) -> Optional[PlacementGroupModel]:
     for pg in placement_groups:
-        if pg.configuration.backend == backend and pg.configuration.region == region:
-            return None
+        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(
-        name=name,
+        # 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=backend,
-            region=region,
+            backend=master_instance_offer.backend,
+            region=master_instance_offer.region,
             placement_strategy=PlacementStrategy.CLUSTER,
         ).json(),
     )
-    session.add(placement_group_model)
+    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 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/background/tasks/process_placement_groups.py

@@ -66,7 +66,7 @@ async def _delete_placement_groups(
 
 
 async def _delete_placement_group(placement_group_model: PlacementGroupModel):
-    logger.info("Deleting placement group %s", placement_group_model.name)
+    logger.debug("Deleting placement group %s", placement_group_model.name)
     placement_group = placement_group_model_to_placement_group(placement_group_model)
     if placement_group.provisioning_data is None:
         logger.error(

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

@@ -99,11 +99,14 @@ async def _collect_jobs_metrics(job_models: list[JobModel], collected_at: dateti
 
 
 async def _collect_job_metrics(job_model: JobModel) -> Optional[str]:
-    ssh_private_keys = get_instance_ssh_private_keys(get_or_error(job_model.instance))
     jpd = get_job_provisioning_data(job_model)
-    jrd = get_job_runtime_data(job_model)
     if jpd is None:
         return None
+    if not jpd.dockerized:
+        # Container-based backend, no shim
+        return None
+    ssh_private_keys = get_instance_ssh_private_keys(get_or_error(job_model.instance))
+    jrd = get_job_runtime_data(job_model)
     try:
         res = await run_async(
             _pull_job_metrics,

dstack/_internal/server/models.py

@@ -659,6 +659,7 @@ class PlacementGroupModel(BaseModel):
 
     fleet_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("fleets.id"))
     fleet: Mapped["FleetModel"] = relationship(foreign_keys=[fleet_id])
+    # TODO: rename `fleet_deleted` -> `to_be_deleted`
     fleet_deleted: Mapped[bool] = mapped_column(Boolean, default=False)
 
     created_at: Mapped[datetime] = mapped_column(NaiveDateTime, default=get_current_datetime)