diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py index 4b2ddf6e2..52ddf09d1 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py +++ b/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]: diff --git a/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py b/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py index 713828d27..b9b08c125 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/runs/__init__.py +++ b/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, ) ) diff --git a/src/dstack/_internal/server/background/pipeline_tasks/runs/pending.py b/src/dstack/_internal/server/background/pipeline_tasks/runs/pending.py index 80aade47f..35a81d9bb 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/runs/pending.py +++ b/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": diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pending.py b/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pending.py index 6cb006821..e643df405 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pending.py +++ b/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: diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pipeline.py b/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pipeline.py index feda67cf3..fdb5d8ce7 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_runs/test_pipeline.py +++ b/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 diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py index 80e062246..0802d594b 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_terminating_jobs.py +++ b/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 ):