Wake pending retries when capacity is released

src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py

@@ -20,6 +20,7 @@
     JobSpec,
     JobStatus,
     JobTerminationReason,
+    RunStatus,
     RunTerminationReason,
 )
 from dstack._internal.server import settings
@@ -294,6 +295,9 @@ async def process(self, item: JobTerminatingPipelineItem):
             and result.instance_update_map.get("status") == InstanceStatus.TERMINATING
         ):
             self._pipeline_hinter.hint_fetch(InstanceModel.__name__)
+        if result.unassign_event_message is not None:
+            await _wake_pending_runs_on_capacity_release()
+            self._pipeline_hinter.hint_fetch(RunModel.__name__)
         # TODO: Hint RunPipeline to quickly move run to TERMINATED.
         # Currently not implemented since it also requires making run eligible for processing.
         # (This pipeline cannot modify runs so it's not simple).
@@ -343,6 +347,20 @@ class _VolumeDetachResult:
     set_volumes_detached_at: bool = False
 
 
+async def _wake_pending_runs_on_capacity_release() -> None:
+    """Allow waiting retries to compete when a real instance releases capacity."""
+    async with get_session_ctx() as session:
+        await session.execute(
+            update(RunModel)
+            .where(
+                RunModel.status == RunStatus.PENDING,
+                RunModel.resubmission_attempt > 0,
+                RunModel.skip_min_processing_interval == False,
+            )
+            .values(skip_min_processing_interval=True)
+        )
+
+
 async def _refetch_locked_job(
     session: AsyncSession, item: JobTerminatingPipelineItem
 ) -> Optional[JobModel]:

src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py

@@ -239,7 +239,10 @@ async def fetch(self, limit: int) -> list[RunPipelineItem]:
                     run_model.lock_expires_at = lock_expires_at
                     run_model.lock_token = lock_token
                     run_model.lock_owner = RunPipeline.__name__
-                    run_model.skip_min_processing_interval = False
+                    # Pending retries use this flag to consume a capacity-release wake-up in
+                    # the worker. Other run states consume it when they are fetched.
+                    if run_model.status != RunStatus.PENDING:
+                        run_model.skip_min_processing_interval = False
                     items.append(
                         RunPipelineItem(
                             __tablename__=RunModel.__tablename__,
@@ -406,6 +409,7 @@ async def _apply_pending_result(
     context: pending.PendingContext,
     result: pending.PendingResult,
 ) -> None:
+    result.run_update_map["skip_min_processing_interval"] = False
     set_processed_update_map_fields(result.run_update_map)
     set_unlock_update_map_fields(result.run_update_map)
 
@@ -480,6 +484,7 @@ async def _apply_noop_result(
                 lock_expires_at=None,
                 lock_token=None,
                 lock_owner=None,
+                skip_min_processing_interval=False,
                 last_processed_at=now,
             )
         )

src/dstack/_internal/server/background/pipeline_tasks/runs/pending.py

@@ -25,6 +25,7 @@ class PendingRunUpdateMap(ItemUpdateMap, total=False):
     termination_reason: Optional[RunTerminationReason]
     desired_replica_count: int
     desired_replica_counts: Optional[str]
+    skip_min_processing_interval: bool
 
 
 @dataclass
@@ -60,7 +61,11 @@ async def process_pending_run(context: PendingContext) -> Optional[PendingResult
             new_job_models=[],
         )
 
-    if run_model.resubmission_attempt > 0 and not _is_ready_for_resubmission(run_model):
+    if (
+        run_model.resubmission_attempt > 0
+        and not run_model.skip_min_processing_interval
+        and not _is_ready_for_resubmission(run_model)
+    ):
         return None
 
     if run_spec.configuration.type == "service":

src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pending.py

@@ -107,6 +107,39 @@ async def test_skips_retrying_run_when_delay_not_met(
         assert run.lock_expires_at is None
         assert run.lock_owner is None
 
+    async def test_capacity_release_wake_skips_retry_delay(
+        self, test_db, session: AsyncSession, worker: RunWorker
+    ) -> None:
+        project = await create_project(session=session)
+        user = await create_user(session=session)
+        repo = await create_repo(session=session, project_id=project.id)
+        run = await create_run(
+            session=session,
+            project=project,
+            repo=repo,
+            user=user,
+            status=RunStatus.PENDING,
+            resubmission_attempt=6,
+        )
+        await create_job(
+            session=session,
+            run=run,
+            status=JobStatus.FAILED,
+            last_processed_at=get_current_datetime(),
+        )
+        run.skip_min_processing_interval = True
+        lock_run(run)
+        await session.commit()
+
+        await worker.process(run_to_pipeline_item(run))
+
+        await session.refresh(run)
+        assert run.status == RunStatus.SUBMITTED
+        assert run.skip_min_processing_interval is False
+        assert run.lock_token is None
+        assert run.lock_expires_at is None
+        assert run.lock_owner is None
+
     async def test_resubmits_retrying_run_after_delay(
         self, test_db, session: AsyncSession, worker: RunWorker
     ) -> None:

src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pipeline.py

@@ -58,6 +58,7 @@ async def test_fetch_selects_eligible_runs_and_sets_lock_fields(
             submitted_at=stale - dt.timedelta(seconds=3),
             resubmission_attempt=1,
         )
+        pending_retry.skip_min_processing_interval = True
         pending_scheduled_ready = await create_run(
             session=session,
             project=project,
@@ -162,7 +163,12 @@ async def test_fetch_selects_eligible_runs_and_sets_lock_fields(
         assert all(run.lock_owner == RunPipeline.__name__ for run in fetched_runs)
         assert all(run.lock_expires_at is not None for run in fetched_runs)
         assert all(run.lock_token is not None for run in fetched_runs)
-        assert all(not run.skip_min_processing_interval for run in fetched_runs)
+        assert all(
+            not run.skip_min_processing_interval
+            for run in fetched_runs
+            if run.id != pending_retry.id
+        )
+        assert pending_retry.skip_min_processing_interval
         assert len({run.lock_token for run in fetched_runs}) == 1
 
         assert future_scheduled.lock_owner is None

src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py

@@ -1,7 +1,7 @@
 import asyncio
 import uuid
 from datetime import datetime, timedelta, timezone
-from typing import Optional
+from typing import Optional, cast
 from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 import pytest
@@ -12,7 +12,7 @@
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.configurations import TaskConfiguration
 from dstack._internal.core.models.instances import InstanceStatus
-from dstack._internal.core.models.runs import JobStatus, JobTerminationReason
+from dstack._internal.core.models.runs import JobStatus, JobTerminationReason, RunStatus
 from dstack._internal.core.models.volumes import VolumeStatus
 from dstack._internal.core.services.ssh.tunnel import SSHTunnel
 from dstack._internal.server.background.pipeline_tasks.jobs_terminating import (
@@ -22,7 +22,7 @@
     JobTerminatingWorker,
     _get_related_instance_lock_owner,
 )
-from dstack._internal.server.models import InstanceModel, JobModel, VolumeAttachmentModel
+from dstack._internal.server.models import InstanceModel, JobModel, RunModel, VolumeAttachmentModel
 from dstack._internal.server.schemas.runner import LogEvent, PullResponse
 from dstack._internal.server.services.runner.client import (
     PeerConnectionError,
@@ -694,6 +694,50 @@ async def test_terminates_job(
             event.message == "Job status changed TERMINATING -> TERMINATED" for event in events
         )
 
+    async def test_wakes_pending_retries_when_capacity_is_released(
+        self, test_db, session: AsyncSession, worker: JobTerminatingWorker
+    ):
+        project = await create_project(session=session)
+        user = await create_user(session=session)
+        instance = await create_instance(
+            session=session,
+            project=project,
+            status=InstanceStatus.BUSY,
+        )
+        repo = await create_repo(session=session, project_id=project.id)
+        run = await create_run(
+            session=session,
+            project=project,
+            repo=repo,
+            user=user,
+            run_name="active-run",
+        )
+        pending_run = await create_run(
+            session=session,
+            project=project,
+            repo=repo,
+            user=user,
+            run_name="pending-retry",
+            status=RunStatus.PENDING,
+            resubmission_attempt=1,
+        )
+        job = await create_job(
+            session=session,
+            run=run,
+            status=JobStatus.TERMINATING,
+            termination_reason=JobTerminationReason.TERMINATED_BY_USER,
+            job_provisioning_data=get_job_provisioning_data(dockerized=False),
+            instance=instance,
+        )
+        _lock_job(job)
+        await session.commit()
+
+        await worker.process(_job_to_pipeline_item(job))
+
+        await session.refresh(pending_run)
+        assert pending_run.skip_min_processing_interval is True
+        cast(Mock, worker._pipeline_hinter.hint_fetch).assert_any_call(RunModel.__name__)
+
     async def test_detaches_job_volumes(
         self, test_db, session: AsyncSession, worker: JobTerminatingWorker
     ):