diff --git a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py index 511823fff..bc551a539 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py @@ -43,11 +43,13 @@ from dstack._internal.core.models.volumes import InstanceMountPoint, Volume, VolumeMountPoint from dstack._internal.server import settings as server_settings from dstack._internal.server.background.pipeline_tasks.base import ( + NOW_PLACEHOLDER, Fetcher, Heartbeater, ItemUpdateMap, Pipeline, PipelineItem, + UpdateMapDateTime, Worker, log_lock_token_changed_after_processing, log_lock_token_mismatch, @@ -146,6 +148,12 @@ JOB_DISCONNECTED_RETRY_TIMEOUT = timedelta(minutes=2) """`The minimum time before terminating active job in case of connectivity issues.""" +MAX_DURATION_ENFORCEMENT_GRACE = timedelta(minutes=2) +"""How long the server waits past `max_duration` before terminating the job itself. +The runner enforces `max_duration` too and does it gracefully, so it normally stops the job +well within the grace period. The server only steps in when the runner failed to. +""" + @dataclass class JobRunningPipelineItem(PipelineItem): @@ -380,6 +388,7 @@ class _JobUpdateMap(ItemUpdateMap, total=False): job_provisioning_data: Optional[str] job_runtime_data: Optional[str] runner_timestamp: Optional[int] + running_at: UpdateMapDateTime disconnected_at: Optional[datetime] inactivity_secs: Optional[int] exit_status: Optional[int] @@ -1045,6 +1054,10 @@ async def _process_running_status( fmt(context.job_model), context.job_submission.age, ) + # Checked before pulling the runner: the runner may be stuck or unreachable, and that is + # exactly when server-side enforcement is needed. + if _terminate_if_max_duration_exceeded(context, result): + return try: process_running_result = await run_async( _process_running, @@ -1808,6 +1821,48 @@ def _terminate_if_inactivity_duration_exceeded( ) +def _terminate_if_max_duration_exceeded( + context: _ProcessContext, + result: _ProcessResult, +) -> bool: + """ + Terminates the job if it has been running longer than `max_duration` plus a grace period. + + A backstop for the runner, which enforces `max_duration` itself and does it gracefully. + The server steps in only when the runner failed to stop the job -- e.g. the workload + survived the termination signals and the runner never reported the timeout. + + Returns `True` if the job was terminated. + """ + job_model = context.job_model + max_duration = context.job.job_spec.max_duration + if max_duration is None: + return False + if job_model.running_at is None: + # Jobs that started running before the server was upgraded have no reference point. + # They are still enforced by the runner. + return False + deadline = ( + job_model.running_at + timedelta(seconds=max_duration) + MAX_DURATION_ENFORCEMENT_GRACE + ) + if get_current_datetime() < deadline: + return False + logger.warning( + "%s: max duration exceeded and the runner did not stop the job, terminating", + fmt(job_model), + ) + _terminate_job( + job_model=job_model, + job_update_map=result.job_update_map, + termination_reason=JobTerminationReason.MAX_DURATION_EXCEEDED, + termination_reason_message=( + f"The job exceeded the max_duration of {max_duration} seconds" + " and did not stop on its own" + ), + ) + return True + + def _should_terminate_job_due_to_disconnect(disconnected_at: Optional[datetime]) -> bool: if disconnected_at is None: return False @@ -2088,6 +2143,10 @@ def _set_job_update_status( ) -> None: if job_update_map.get("status", job_model.status) != new_status: job_update_map["status"] = new_status + if new_status == JobStatus.RUNNING: + # Stamped here rather than at the call site so that `running_at` cannot drift from + # `status`: it is the reference point for server-side `max_duration` enforcement. + job_update_map["running_at"] = NOW_PLACEHOLDER def _set_job_status(job_model: JobModel, result: _ProcessResult, new_status: JobStatus) -> None: diff --git a/src/dstack/_internal/server/migrations/versions/2026/09_08_1002_652af7a3c9c9_add_jobmodel_running_at.py b/src/dstack/_internal/server/migrations/versions/2026/09_08_1002_652af7a3c9c9_add_jobmodel_running_at.py new file mode 100644 index 000000000..e0236ea3e --- /dev/null +++ b/src/dstack/_internal/server/migrations/versions/2026/09_08_1002_652af7a3c9c9_add_jobmodel_running_at.py @@ -0,0 +1,30 @@ +"""Add JobModel.running_at + +Revision ID: 652af7a3c9c9 +Revises: 620892d149b5 +Create Date: 2026-09-08 10:02:00.000000 + +""" + +import sqlalchemy as sa +from alembic import op + +import dstack._internal.server.models + +# revision identifiers, used by Alembic. +revision = "652af7a3c9c9" +down_revision = "620892d149b5" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + with op.batch_alter_table("jobs", schema=None) as batch_op: + batch_op.add_column( + sa.Column("running_at", dstack._internal.server.models.NaiveDateTime(), nullable=True) + ) + + +def downgrade() -> None: + with op.batch_alter_table("jobs", schema=None) as batch_op: + batch_op.drop_column("running_at") diff --git a/src/dstack/_internal/server/models.py b/src/dstack/_internal/server/models.py index 5b5e95102..e72153cae 100644 --- a/src/dstack/_internal/server/models.py +++ b/src/dstack/_internal/server/models.py @@ -535,6 +535,12 @@ class JobModel(PipelineModelMixin, BaseModel): job_name: Mapped[str] = mapped_column(String(100)) submission_num: Mapped[int] = mapped_column(Integer) submitted_at: Mapped[datetime] = mapped_column(NaiveDateTime) + running_at: Mapped[Optional[datetime]] = mapped_column(NaiveDateTime) + """`running_at` stores when the job entered the `RUNNING` status, that is, when the workload + started, excluding provisioning and pulling. It is the reference point for server-side + `max_duration` enforcement. `None` for jobs that never started running and for jobs that were + already running before the server was upgraded. + """ last_processed_at: Mapped[datetime] = mapped_column(NaiveDateTime) skip_min_processing_interval: Mapped[bool] = mapped_column( Boolean, default=False, server_default=false() diff --git a/src/dstack/_internal/server/testing/common.py b/src/dstack/_internal/server/testing/common.py index 18aef0b0a..f6dc63181 100644 --- a/src/dstack/_internal/server/testing/common.py +++ b/src/dstack/_internal/server/testing/common.py @@ -432,6 +432,7 @@ async def create_job( submission_num: int = 0, status: JobStatus = JobStatus.SUBMITTED, submitted_at: datetime = datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), + running_at: Optional[datetime] = None, last_processed_at: datetime = datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc), termination_reason: Optional[JobTerminationReason] = None, job_provisioning_data: Optional[JobProvisioningData] = None, @@ -475,6 +476,7 @@ async def create_job( deployment_num=deployment_num, submission_num=submission_num, submitted_at=submitted_at, + running_at=running_at, last_processed_at=last_processed_at, status=status, termination_reason=termination_reason, diff --git a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py index 88a2245a1..1e13956ca 100644 --- a/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py +++ b/src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py @@ -26,7 +26,7 @@ from dstack._internal.core.models.duration import Duration from dstack._internal.core.models.gateways import GatewayReplicaStatus, GatewayStatus from dstack._internal.core.models.instances import InstanceStatus -from dstack._internal.core.models.profiles import StartupOrder, UtilizationPolicy +from dstack._internal.core.models.profiles import Profile, StartupOrder, UtilizationPolicy from dstack._internal.core.models.runs import ( ClusterInfo, ImagePullProgress, @@ -44,6 +44,7 @@ from dstack._internal.server import settings as server_settings from dstack._internal.server.background.pipeline_tasks.jobs_running import ( JOB_DISCONNECTED_RETRY_TIMEOUT, + MAX_DURATION_ENFORCEMENT_GRACE, ROUTER_PROVISIONING_WAIT_TIMEOUT_SECONDS, JobRunningFetcher, JobRunningPipeline, @@ -1752,6 +1753,128 @@ async def test_inactivity_duration( assert job.termination_reason == expected_termination_reason assert job.inactivity_secs == expected_inactivity_secs + @pytest.mark.parametrize( + ( + "max_duration", + "running_for", + "stamp_running_at", + "expected_status", + "expected_termination_reason", + "expect_pull", + ), + [ + pytest.param( + 600, + timedelta(seconds=600), + True, + JobStatus.RUNNING, + None, + True, + id="deadline-reached-but-runner-still-within-grace", + ), + pytest.param( + 600, + timedelta(seconds=600) + MAX_DURATION_ENFORCEMENT_GRACE - timedelta(seconds=1), + True, + JobStatus.RUNNING, + None, + True, + id="grace-not-elapsed", + ), + pytest.param( + 600, + timedelta(seconds=600) + MAX_DURATION_ENFORCEMENT_GRACE, + True, + JobStatus.TERMINATING, + JobTerminationReason.MAX_DURATION_EXCEEDED, + False, + id="grace-elapsed", + ), + pytest.param( + 600, + timedelta(days=1), + False, + JobStatus.RUNNING, + None, + True, + id="job-started-before-upgrade", + ), + pytest.param( + "off", + timedelta(days=1), + True, + JobStatus.RUNNING, + None, + True, + id="max-duration-off", + ), + ], + ) + async def test_max_duration_enforced_by_server( + self, + test_db, + session: AsyncSession, + worker: JobRunningWorker, + max_duration, + running_for: timedelta, + stamp_running_at: bool, + expected_status: JobStatus, + expected_termination_reason: Optional[JobTerminationReason], + expect_pull: bool, + ) -> None: + now = datetime(2023, 1, 2, 5, 0, tzinfo=timezone.utc) + 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.RUNNING, + run_name="test-run", + run_spec=get_run_spec( + run_name="test-run", + repo_id=repo.name, + profile=Profile(name="default", max_duration=max_duration), + ), + ) + instance = await create_instance( + session=session, + project=project, + status=InstanceStatus.BUSY, + ) + job = await create_job( + session=session, + run=run, + status=JobStatus.RUNNING, + running_at=(now - running_for) if stamp_running_at else None, + job_provisioning_data=get_job_provisioning_data(), + instance=instance, + instance_assigned=True, + ) + with ( + patch("dstack._internal.server.services.runner.pool.SSHTunnel"), + patch( + "dstack._internal.server.services.runner.client.RunnerClient.from_address" + ) as runner_client_cls, + freeze_time(now), + ): + runner_client_mock = runner_client_cls.return_value + runner_client_mock.pull.return_value = PullResponse( + job_states=[], + job_logs=[], + runner_logs=[], + last_updated=0, + no_connections_secs=0, + ) + await _process_job(session, worker, job) + assert runner_client_mock.pull.called == expect_pull + + await session.refresh(job) + assert job.status == expected_status + assert job.termination_reason == expected_termination_reason + @pytest.mark.parametrize( ["samples", "expected_status"], [