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

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

@@ -1,11 +1,13 @@
 from datetime import timedelta
 from typing import List
+from uuid import UUID
 
 from sqlalchemy import select, update
 from sqlalchemy.ext.asyncio import AsyncSession
 from sqlalchemy.orm import joinedload, load_only
 
-from dstack._internal.core.models.fleets import FleetStatus
+from dstack._internal.core.models.fleets import FleetSpec, FleetStatus
+from dstack._internal.core.models.instances import InstanceStatus
 from dstack._internal.server.db import get_db, get_session_ctx
 from dstack._internal.server.models import (
     FleetModel,
@@ -15,7 +17,9 @@ from dstack._internal.server.models import (
     RunModel,
 )
 from dstack._internal.server.services.fleets import (
+    create_fleet_instance_model,
     get_fleet_spec,
+    get_next_instance_num,
     is_fleet_empty,
     is_fleet_in_use,
 )
@@ -65,31 +69,111 @@ async def _process_fleets(session: AsyncSession, fleet_models: List[FleetModel])
     res = await session.execute(
         select(FleetModel)
         .where(FleetModel.id.in_(fleet_ids))
-        .options(joinedload(FleetModel.instances).load_only(InstanceModel.deleted))
         .options(
-            joinedload(FleetModel.instances).joinedload(InstanceModel.jobs).load_only(JobModel.id)
+            joinedload(FleetModel.instances).joinedload(InstanceModel.jobs).load_only(JobModel.id),
+            joinedload(FleetModel.project),
         )
         .options(joinedload(FleetModel.runs).load_only(RunModel.status))
         .execution_options(populate_existing=True)
     )
     fleet_models = list(res.unique().scalars().all())
 
+    # TODO: Drop fleets auto-deletion after dropping fleets auto-creation.
     deleted_fleets_ids = []
-    now = get_current_datetime()
     for fleet_model in fleet_models:
+        _consolidate_fleet_state_with_spec(session, fleet_model)
         deleted = _autodelete_fleet(fleet_model)
         if deleted:
             deleted_fleets_ids.append(fleet_model.id)
-        fleet_model.last_processed_at = now
+        fleet_model.last_processed_at = get_current_datetime()
+    await _update_deleted_fleets_placement_groups(session, deleted_fleets_ids)
+    await session.commit()
 
-    await session.execute(
-        update(PlacementGroupModel)
-        .where(
-            PlacementGroupModel.fleet_id.in_(deleted_fleets_ids),
+
+def _consolidate_fleet_state_with_spec(session: AsyncSession, fleet_model: FleetModel):
+    if fleet_model.status == FleetStatus.TERMINATING:
+        return
+    fleet_spec = get_fleet_spec(fleet_model)
+    if fleet_spec.configuration.nodes is None or fleet_spec.autocreated:
+        # Only explicitly created cloud fleets are consolidated.
+        return
+    if not _is_fleet_ready_for_consolidation(fleet_model):
+        return
+    added_instances = _maintain_fleet_nodes_min(session, fleet_model, fleet_spec)
+    if added_instances:
+        fleet_model.consolidation_attempt += 1
+    else:
+        # The fleet is already consolidated or consolidation is in progress.
+        # We reset consolidation_attempt in both cases for simplicity.
+        # The second case does not need reset but is ok to do since
+        # it means consolidation is longer than delay, so it won't happen too often.
+        # TODO: Reset consolidation_attempt on fleet in-place update.
+        fleet_model.consolidation_attempt = 0
+    fleet_model.last_consolidated_at = get_current_datetime()
+
+
+def _is_fleet_ready_for_consolidation(fleet_model: FleetModel) -> bool:
+    consolidation_retry_delay = _get_consolidation_retry_delay(fleet_model.consolidation_attempt)
+    last_consolidated_at = fleet_model.last_consolidated_at or fleet_model.last_processed_at
+    duration_since_last_consolidation = get_current_datetime() - last_consolidated_at
+    return duration_since_last_consolidation >= consolidation_retry_delay
+
+
+# We use exponentially increasing consolidation retry delays so that
+# consolidation does not happen too often. In particular, this prevents
+# retrying instance provisioning constantly in case of no offers.
+# TODO: Adjust delays.
+_CONSOLIDATION_RETRY_DELAYS = [
+    timedelta(seconds=30),
+    timedelta(minutes=1),
+    timedelta(minutes=2),
+    timedelta(minutes=5),
+    timedelta(minutes=10),
+]
+
+
+def _get_consolidation_retry_delay(consolidation_attempt: int) -> timedelta:
+    if consolidation_attempt < len(_CONSOLIDATION_RETRY_DELAYS):
+        return _CONSOLIDATION_RETRY_DELAYS[consolidation_attempt]
+    return _CONSOLIDATION_RETRY_DELAYS[-1]
+
+
+def _maintain_fleet_nodes_min(
+    session: AsyncSession,
+    fleet_model: FleetModel,
+    fleet_spec: FleetSpec,
+) -> bool:
+    """
+    Ensures the fleet has at least `nodes.min` instances.
+    Returns `True` if retried or added new instances and `False` otherwise.
+    """
+    assert fleet_spec.configuration.nodes is not None
+    for instance in fleet_model.instances:
+        # Delete terminated but not deleted instances since
+        # they are going to be replaced with new pending instances.
+        if instance.status == InstanceStatus.TERMINATED and not instance.deleted:
+            # It's safe to modify instances without instance lock since
+            # no other task modifies already terminated instances.
+            instance.deleted = True
+            instance.deleted_at = get_current_datetime()
+    active_instances = [i for i in fleet_model.instances if not i.deleted]
+    active_instances_num = len(active_instances)
+    if active_instances_num >= fleet_spec.configuration.nodes.min:
+        return False
+    nodes_missing = fleet_spec.configuration.nodes.min - active_instances_num
+    for i in range(nodes_missing):
+        instance_model = create_fleet_instance_model(
+            session=session,
+            project=fleet_model.project,
+            # TODO: Store fleet.user and pass it instead of the project owner.
+            username=fleet_model.project.owner.name,
+            spec=fleet_spec,
+            instance_num=get_next_instance_num({i.instance_num for i in active_instances}),
         )
-        .values(fleet_deleted=True)
-    )
-    await session.commit()
+        active_instances.append(instance_model)
+        fleet_model.instances.append(instance_model)
+    logger.info("Added %s instances to fleet %s", nodes_missing, fleet_model.name)
+    return True
 
 
 def _autodelete_fleet(fleet_model: FleetModel) -> bool:
@@ -100,7 +184,7 @@ def _autodelete_fleet(fleet_model: FleetModel) -> bool:
     if (
         fleet_model.status != FleetStatus.TERMINATING
         and fleet_spec.configuration.nodes is not None
-        and (fleet_spec.configuration.nodes.min is None or fleet_spec.configuration.nodes.min == 0)
+        and fleet_spec.configuration.nodes.min == 0
     ):
         # Empty fleets that allow 0 nodes should not be auto-deleted
         return False
@@ -110,3 +194,15 @@ def _autodelete_fleet(fleet_model: FleetModel) -> bool:
     fleet_model.deleted = True
     logger.info("Fleet %s deleted", fleet_model.name)
     return True
+
+
+async def _update_deleted_fleets_placement_groups(session: AsyncSession, fleets_ids: list[UUID]):
+    if len(fleets_ids) == 0:
+        return
+    await session.execute(
+        update(PlacementGroupModel)
+        .where(
+            PlacementGroupModel.fleet_id.in_(fleets_ids),
+        )
+        .values(fleet_deleted=True)
+    )

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

@@ -49,8 +49,8 @@ async def process_gateways():
             if gateway_model is None:
                 return
             lockset.add(gateway_model.id)
+        gateway_model_id = gateway_model.id
         try:
-            gateway_model_id = gateway_model.id
             initial_status = gateway_model.status
             if initial_status == GatewayStatus.SUBMITTED:
                 await _process_submitted_gateway(session=session, gateway_model=gateway_model)
@@ -165,6 +165,9 @@ async def _process_provisioning_gateway(
     )
     gateway_model = res.unique().scalar_one()
 
+    # Provisioning gateways must have compute.
+    assert gateway_model.gateway_compute is not None
+
     # FIXME: problems caused by blocking on connect_to_gateway_with_retry and configure_gateway:
     # - cannot delete the gateway before it is provisioned because the DB model is locked
     # - connection retry counter is reset on server restart

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

@@ -53,14 +53,12 @@ from dstack._internal.core.models.placement import (
     PlacementStrategy,
 )
 from dstack._internal.core.models.profiles import (
-    RetryEvent,
     TerminationPolicy,
 )
 from dstack._internal.core.models.runs import (
     JobProvisioningData,
     Retry,
 )
-from dstack._internal.core.services.profiles import get_retry
 from dstack._internal.server import settings as server_settings
 from dstack._internal.server.background.tasks.common import get_provisioning_timeout
 from dstack._internal.server.db import get_db, get_session_ctx
@@ -85,8 +83,10 @@ from dstack._internal.server.services.instances import (
     get_instance_provisioning_data,
     get_instance_requirements,
     get_instance_ssh_private_keys,
+    remove_dangling_tasks_from_instance,
 )
 from dstack._internal.server.services.locking import get_locker
+from dstack._internal.server.services.logging import fmt
 from dstack._internal.server.services.offers import is_divisible_into_blocks
 from dstack._internal.server.services.placement import (
     get_fleet_placement_group_models,
@@ -181,8 +181,8 @@ async def _process_next_instance():
             if instance is None:
                 return
             lockset.add(instance.id)
+        instance_model_id = instance.id
         try:
-            instance_model_id = instance.id
             await _process_instance(session=session, instance=instance)
         finally:
             lockset.difference_update([instance_model_id])
@@ -325,7 +325,6 @@ async def _add_remote(instance: InstanceModel) -> None:
             e,
         )
         instance.status = InstanceStatus.PENDING
-        instance.last_retry_at = get_current_datetime()
         return
 
     instance_type = host_info_to_instance_type(host_info, cpu_arch)
@@ -393,6 +392,7 @@ async def _add_remote(instance: InstanceModel) -> None:
         return
 
     region = instance.region
+    assert region is not None  # always set for ssh instances
     jpd = JobProvisioningData(
         backend=BackendType.REMOTE,
         instance_type=instance_type,
@@ -423,7 +423,6 @@ async def _add_remote(instance: InstanceModel) -> None:
     instance.offer = instance_offer.json()
     instance.job_provisioning_data = jpd.json()
     instance.started_at = get_current_datetime()
-    instance.last_retry_at = get_current_datetime()
 
 
 def _deploy_instance(
@@ -490,29 +489,6 @@ def _deploy_instance(
 
 
 async def _create_instance(session: AsyncSession, instance: InstanceModel) -> None:
-    if instance.last_retry_at is not None:
-        last_retry = instance.last_retry_at
-        if get_current_datetime() < last_retry + timedelta(minutes=1):
-            return
-
-    if (
-        instance.profile is None
-        or instance.requirements is None
-        or instance.instance_configuration is None
-    ):
-        instance.status = InstanceStatus.TERMINATED
-        instance.termination_reason = "Empty profile, requirements or instance_configuration"
-        instance.last_retry_at = get_current_datetime()
-        logger.warning(
-            "Empty profile, requirements or instance_configuration. Terminate instance: %s",
-            instance.name,
-            extra={
-                "instance_name": instance.name,
-                "instance_status": InstanceStatus.TERMINATED.value,
-            },
-        )
-        return
-
     if _need_to_wait_fleet_provisioning(instance):
         logger.debug("Waiting for the first instance in the fleet to be provisioned")
         return
@@ -526,7 +502,6 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
         instance.termination_reason = (
             f"Error to parse profile, requirements or instance_configuration: {e}"
         )
-        instance.last_retry_at = get_current_datetime()
         logger.warning(
             "Error to parse profile, requirements or instance_configuration. Terminate instance: %s",
             instance.name,
@@ -537,24 +512,6 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
         )
         return
 
-    retry = get_retry(profile)
-    should_retry = retry is not None and RetryEvent.NO_CAPACITY in retry.on_events
-
-    if retry is not None:
-        retry_duration_deadline = _get_retry_duration_deadline(instance, retry)
-        if get_current_datetime() > retry_duration_deadline:
-            instance.status = InstanceStatus.TERMINATED
-            instance.termination_reason = "Retry duration expired"
-            logger.warning(
-                "Retry duration expired. Terminating instance %s",
-                instance.name,
-                extra={
-                    "instance_name": instance.name,
-                    "instance_status": InstanceStatus.TERMINATED.value,
-                },
-            )
-            return
-
     placement_group_models = []
     placement_group_model = None
     if instance.fleet_id:
@@ -592,15 +549,6 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
         exclude_not_available=True,
     )
 
-    if not offers and should_retry:
-        instance.last_retry_at = get_current_datetime()
-        logger.debug(
-            "No offers for instance %s. Next retry",
-            instance.name,
-            extra={"instance_name": instance.name},
-        )
-        return
-
     # 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]:
@@ -678,7 +626,6 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
         instance.offer = instance_offer.json()
         instance.total_blocks = instance_offer.total_blocks
         instance.started_at = get_current_datetime()
-        instance.last_retry_at = get_current_datetime()
 
         logger.info(
             "Created instance %s",
@@ -699,21 +646,18 @@ async def _create_instance(session: AsyncSession, instance: InstanceModel) -> No
             )
         return
 
-    instance.last_retry_at = get_current_datetime()
-
-    if not should_retry:
-        _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")
+    _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:
@@ -788,6 +732,7 @@ async def _check_instance(session: AsyncSession, instance: InstanceModel) -> Non
         ssh_private_keys,
         job_provisioning_data,
         None,
+        instance=instance,
         check_instance_health=check_instance_health,
     )
     if instance_check is False:
@@ -934,7 +879,7 @@ async def _wait_for_instance_provisioning_data(
 
 @runner_ssh_tunnel(ports=[DSTACK_SHIM_HTTP_PORT], retries=1)
 def _check_instance_inner(
-    ports: Dict[int, int], *, check_instance_health: bool = False
+    ports: Dict[int, int], *, instance: InstanceModel, check_instance_health: bool = False
 ) -> InstanceCheck:
     instance_health_response: Optional[InstanceHealthResponse] = None
     shim_client = runner_client.ShimClient(port=ports[DSTACK_SHIM_HTTP_PORT])
@@ -954,6 +899,10 @@ def _check_instance_inner(
         args = (method.__func__.__name__, e.__class__.__name__, e)
         logger.exception(template, *args)
         return InstanceCheck(reachable=False, message=template % args)
+    try:
+        remove_dangling_tasks_from_instance(shim_client, instance)
+    except Exception as e:
+        logger.exception("%s: error removing dangling tasks: %s", fmt(instance), e)
     return runner_client.healthcheck_response_to_instance_check(
         healthcheck_response, instance_health_response
     )

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

@@ -120,7 +120,7 @@ async def _execute_probe(probe: ProbeModel, probe_spec: ProbeSpec) -> bool:
                 method=probe_spec.method,
                 url="http://dstack" + probe_spec.url,
                 headers=[(h.name, h.value) for h in probe_spec.headers],
-                data=probe_spec.body,
+                content=probe_spec.body,
                 timeout=probe_spec.timeout,
                 follow_redirects=False,
             )

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

@@ -41,6 +41,7 @@ from dstack._internal.core.models.volumes import InstanceMountPoint, Volume, Vol
 from dstack._internal.server.background.tasks.common import get_provisioning_timeout
 from dstack._internal.server.db import get_db, get_session_ctx
 from dstack._internal.server.models import (
+    FleetModel,
     InstanceModel,
     JobModel,
     ProbeModel,
@@ -128,9 +129,8 @@ async def _process_next_running_job():
             if job_model is None:
                 return
             lockset.add(job_model.id)
-
+        job_model_id = job_model.id
         try:
-            job_model_id = job_model.id
             await _process_running_job(session=session, job_model=job_model)
         finally:
             lockset.difference_update([job_model_id])
@@ -152,6 +152,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
         .options(joinedload(RunModel.project))
         .options(joinedload(RunModel.user))
         .options(joinedload(RunModel.repo))
+        .options(joinedload(RunModel.fleet).load_only(FleetModel.id, FleetModel.name))
         .options(joinedload(RunModel.jobs))
     )
     run_model = res.unique().scalar_one()
@@ -170,6 +171,11 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
 
     job = find_job(run.jobs, job_model.replica_num, job_model.job_num)
 
+    volumes = []
+    secrets = {}
+    cluster_info = None
+    repo_creds = None
+
     initial_status = job_model.status
     if initial_status in [JobStatus.PROVISIONING, JobStatus.PULLING]:
         # Wait until all other jobs in the replica are provisioned
@@ -257,6 +263,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                 user_ssh_key,
             )
         else:
+            assert cluster_info is not None
             logger.debug(
                 "%s: process provisioning job without shim, age=%s",
                 fmt(job_model),
@@ -275,7 +282,6 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                 repo=repo_model,
                 code_hash=_get_repo_code_hash(run, job),
             )
-
             success = await common_utils.run_async(
                 _submit_job_to_runner,
                 server_ssh_private_keys,
@@ -309,6 +315,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
 
     else:  # fails are not acceptable
         if initial_status == JobStatus.PULLING:
+            assert cluster_info is not None
             logger.debug(
                 "%s: process pulling job with shim, age=%s", fmt(job_model), job_submission.age
             )
@@ -341,7 +348,7 @@ async def _process_running_job(session: AsyncSession, job_model: JobModel):
                 server_ssh_private_keys,
                 job_provisioning_data,
             )
-        elif initial_status == JobStatus.RUNNING:
+        else:
             logger.debug("%s: process running job, age=%s", fmt(job_model), job_submission.age)
             success = await common_utils.run_async(
                 _process_running,
@@ -632,6 +639,7 @@ def _process_pulling_with_shim(
         is successful
     """
     shim_client = client.ShimClient(port=ports[DSTACK_SHIM_HTTP_PORT])
+    job_runtime_data = None
     if shim_client.is_api_v2_supported():  # raises error if shim is down, causes retry
         task = shim_client.get_task(job_model.id)
 

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

@@ -21,6 +21,7 @@ from dstack._internal.core.models.runs import (
 )
 from dstack._internal.server.db import get_db, get_session_ctx
 from dstack._internal.server.models import (
+    FleetModel,
     InstanceModel,
     JobModel,
     ProjectModel,
@@ -129,8 +130,8 @@ async def _process_next_run():
             job_ids = [j.id for j in run_model.jobs]
             run_lockset.add(run_model.id)
             job_lockset.update(job_ids)
+        run_model_id = run_model.id
         try:
-            run_model_id = run_model.id
             await _process_run(session=session, run_model=run_model)
         finally:
             run_lockset.difference_update([run_model_id])
@@ -145,6 +146,7 @@ async def _process_run(session: AsyncSession, run_model: RunModel):
         .execution_options(populate_existing=True)
         .options(joinedload(RunModel.project).load_only(ProjectModel.id, ProjectModel.name))
         .options(joinedload(RunModel.user).load_only(UserModel.name))
+        .options(joinedload(RunModel.fleet).load_only(FleetModel.id, FleetModel.name))
         .options(
             selectinload(RunModel.jobs)
             .joinedload(JobModel.instance)