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

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_metrics.py

@@ -42,10 +42,33 @@ async def collect_metrics():
 
 
 async def delete_metrics():
-    cutoff = _get_delete_metrics_cutoff()
+    now_timestamp_micro = int(get_current_datetime().timestamp() * 1_000_000)
+    running_timestamp_micro_cutoff = (
+        now_timestamp_micro - settings.SERVER_METRICS_RUNNING_TTL_SECONDS * 1_000_000
+    )
+    finished_timestamp_micro_cutoff = (
+        now_timestamp_micro - settings.SERVER_METRICS_FINISHED_TTL_SECONDS * 1_000_000
+    )
     async with get_session_ctx() as session:
-        await session.execute(
-            delete(JobMetricsPoint).where(JobMetricsPoint.timestamp_micro < cutoff)
+        await asyncio.gather(
+            session.execute(
+                delete(JobMetricsPoint).where(
+                    JobMetricsPoint.job_id.in_(
+                        select(JobModel.id).where(JobModel.status.in_([JobStatus.RUNNING]))
+                    ),
+                    JobMetricsPoint.timestamp_micro < running_timestamp_micro_cutoff,
+                )
+            ),
+            session.execute(
+                delete(JobMetricsPoint).where(
+                    JobMetricsPoint.job_id.in_(
+                        select(JobModel.id).where(
+                            JobModel.status.in_(JobStatus.finished_statuses())
+                        )
+                    ),
+                    JobMetricsPoint.timestamp_micro < finished_timestamp_micro_cutoff,
+                )
+            ),
         )
         await session.commit()
 
@@ -134,9 +157,3 @@ def _pull_runner_metrics(
 ) -> Optional[MetricsResponse]:
     runner_client = client.RunnerClient(port=ports[DSTACK_RUNNER_HTTP_PORT])
     return runner_client.get_metrics()
-
-
-def _get_delete_metrics_cutoff() -> int:
-    now = int(get_current_datetime().timestamp() * 1_000_000)
-    cutoff = now - (settings.SERVER_METRICS_TTL_SECONDS * 1_000_000)
-    return cutoff

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/background/tasks/process_running_jobs.py

@@ -1,6 +1,6 @@
 import asyncio
 from collections.abc import Iterable
-from datetime import timedelta
+from datetime import timedelta, timezone
 from typing import Dict, List, Optional
 
 from sqlalchemy import select
@@ -71,6 +71,12 @@ from dstack._internal.utils.logging import get_logger
 logger = get_logger(__name__)
 
 
+# Minimum time before terminating active job in case of connectivity issues.
+# Should be sufficient to survive most problems caused by
+# the server network flickering and providers' glitches.
+JOB_DISCONNECTED_RETRY_TIMEOUT = timedelta(minutes=2)
+
+
 async def process_running_jobs(batch_size: int = 1):
     tasks = []
     for _ in range(batch_size):
@@ -202,7 +208,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                 user_ssh_key = run.run_spec.ssh_key_pub.strip()
                 public_keys = [project.ssh_public_key.strip(), user_ssh_key]
                 if job_provisioning_data.backend == BackendType.LOCAL:
-                    # No need to update ~/.ssh/authorized_keys when running shim localy
+                    # No need to update ~/.ssh/authorized_keys when running shim locally
                     user_ssh_key = ""
                 success = await common_utils.run_async(
                     _process_provisioning_with_shim,
@@ -299,19 +305,38 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                 run_model,
                 job_model,
             )
-            if not success:
-                job_model.termination_reason = JobTerminationReason.INTERRUPTED_BY_NO_CAPACITY
 
-        if not success:  # kill the job
-            logger.warning(
-                "%s: failed because runner is not available or return an error,  age=%s",
-                fmt(job_model),
-                job_submission.age,
-            )
-            job_model.status = JobStatus.TERMINATING
-            if not job_model.termination_reason:
-                job_model.termination_reason = JobTerminationReason.INTERRUPTED_BY_NO_CAPACITY
-            # job will be terminated and instance will be emptied by process_terminating_jobs
+        if success:
+            job_model.disconnected_at = None
+        else:
+            if job_model.termination_reason:
+                logger.warning(
+                    "%s: failed because shim/runner returned an error, age=%s",
+                    fmt(job_model),
+                    job_submission.age,
+                )
+                job_model.status = JobStatus.TERMINATING
+                # job will be terminated and instance will be emptied by process_terminating_jobs
+            else:
+                # No job_model.termination_reason set means ssh connection failed
+                if job_model.disconnected_at is None:
+                    job_model.disconnected_at = common_utils.get_current_datetime()
+                if _should_terminate_job_due_to_disconnect(job_model):
+                    logger.warning(
+                        "%s: failed because instance is unreachable, age=%s",
+                        fmt(job_model),
+                        job_submission.age,
+                    )
+                    # TODO: Replace with JobTerminationReason.INSTANCE_UNREACHABLE in 0.20 or
+                    # when CLI <= 0.19.8 is no longer supported
+                    job_model.termination_reason = JobTerminationReason.INTERRUPTED_BY_NO_CAPACITY
+                    job_model.status = JobStatus.TERMINATING
+                else:
+                    logger.warning(
+                        "%s: is unreachable, waiting for the instance to become reachable again, age=%s",
+                        fmt(job_model),
+                        job_submission.age,
+                    )
 
     if (
         initial_status != job_model.status
@@ -543,7 +568,7 @@ def _process_pulling_with_shim(
     if shim_client.is_api_v2_supported():  # raises error if shim is down, causes retry
         task = shim_client.get_task(job_model.id)
 
-        # If task goes to terminated before the job is submitted to runner, then an error occured
+        # If task goes to terminated before the job is submitted to runner, then an error occurred
         if task.status == TaskStatus.TERMINATED:
             logger.warning(
                 "shim failed to execute job %s: %s (%s)",
@@ -572,7 +597,7 @@ def _process_pulling_with_shim(
     else:
         shim_status = shim_client.pull()  # raises error if shim is down, causes retry
 
-        # If shim goes to pending before the job is submitted to runner, then an error occured
+        # If shim goes to pending before the job is submitted to runner, then an error occurred
         if (
             shim_status.state == "pending"
             and shim_status.result is not None
@@ -651,6 +676,10 @@ def _process_running(
                 )
             if latest_state_event.termination_message:
                 job_model.termination_reason_message = latest_state_event.termination_message
+        if (exit_status := latest_state_event.exit_status) is not None:
+            job_model.exit_status = exit_status
+            if exit_status != 0:
+                logger.info("%s: non-zero exit status %s", fmt(job_model), exit_status)
     else:
         _terminate_if_inactivity_duration_exceeded(run_model, job_model, resp.no_connections_secs)
     if job_model.status != previous_status:
@@ -688,6 +717,15 @@ def _terminate_if_inactivity_duration_exceeded(
         )
 
 
+def _should_terminate_job_due_to_disconnect(job_model: JobModel) -> bool:
+    if job_model.disconnected_at is None:
+        return False
+    return (
+        common_utils.get_current_datetime()
+        > job_model.disconnected_at.replace(tzinfo=timezone.utc) + JOB_DISCONNECTED_RETRY_TIMEOUT
+    )
+
+
 async def _check_gpu_utilization(session: AsyncSession, job_model: JobModel, job: Job) -> None:
     policy = job.job_spec.utilization_policy
     if policy is None:
@@ -818,8 +856,8 @@ def _submit_job_to_runner(
         return success_if_not_available
 
     runner_client.submit_job(
-        run_spec=run.run_spec,
-        job_spec=job.job_spec,
+        run=run,
+        job=job,
         cluster_info=cluster_info,
         secrets=secrets,
         repo_credentials=repo_credentials,