Skip to content

Commit b60d98c

Browse files
authored
Enforce max_duration on the server (#4273)
Previously, `max_duration` was enforced only by the runner, which cancels the job context once the deadline passes. When that failed to stop the workload, nothing else did: a RUNNING job has no server-side liveness bound, and the runner keeps answering `/api/pull` with the job still running. Now the server enforces the deadline as well, without depending on the runner being healthy. `JobModel.running_at` records when the job entered RUNNING, and a job still running that long plus `max_duration` and `MAX_DURATION_ENFORCEMENT_GRACE` is terminated with `MAX_DURATION_EXCEEDED`. The grace period keeps the runner primary, since it stops the job gracefully and reports a more accurate state, and lets the server act only when the runner did not. Jobs whose `running_at` is unset -- those already running when the server is upgraded -- are left to the runner. Fixes: #4259
1 parent 0140d63 commit b60d98c

5 files changed

Lines changed: 221 additions & 1 deletion

File tree

‎src/dstack/_internal/server/background/pipeline_tasks/jobs_running.py‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,11 +43,13 @@
4343
from dstack._internal.core.models.volumes import InstanceMountPoint, Volume, VolumeMountPoint
4444
from dstack._internal.server import settings as server_settings
4545
from dstack._internal.server.background.pipeline_tasks.base import (
46+
NOW_PLACEHOLDER,
4647
Fetcher,
4748
Heartbeater,
4849
ItemUpdateMap,
4950
Pipeline,
5051
PipelineItem,
52+
UpdateMapDateTime,
5153
Worker,
5254
log_lock_token_changed_after_processing,
5355
log_lock_token_mismatch,
@@ -146,6 +148,12 @@
146148
JOB_DISCONNECTED_RETRY_TIMEOUT = timedelta(minutes=2)
147149
"""`The minimum time before terminating active job in case of connectivity issues."""
148150

151+
MAX_DURATION_ENFORCEMENT_GRACE = timedelta(minutes=2)
152+
"""How long the server waits past `max_duration` before terminating the job itself.
153+
The runner enforces `max_duration` too and does it gracefully, so it normally stops the job
154+
well within the grace period. The server only steps in when the runner failed to.
155+
"""
156+
149157

150158
@dataclass
151159
class JobRunningPipelineItem(PipelineItem):
@@ -380,6 +388,7 @@ class _JobUpdateMap(ItemUpdateMap, total=False):
380388
job_provisioning_data: Optional[str]
381389
job_runtime_data: Optional[str]
382390
runner_timestamp: Optional[int]
391+
running_at: UpdateMapDateTime
383392
disconnected_at: Optional[datetime]
384393
inactivity_secs: Optional[int]
385394
exit_status: Optional[int]
@@ -1045,6 +1054,10 @@ async def _process_running_status(
10451054
fmt(context.job_model),
10461055
context.job_submission.age,
10471056
)
1057+
# Checked before pulling the runner: the runner may be stuck or unreachable, and that is
1058+
# exactly when server-side enforcement is needed.
1059+
if _terminate_if_max_duration_exceeded(context, result):
1060+
return
10481061
try:
10491062
process_running_result = await run_async(
10501063
_process_running,
@@ -1808,6 +1821,48 @@ def _terminate_if_inactivity_duration_exceeded(
18081821
)
18091822

18101823

1824+
def _terminate_if_max_duration_exceeded(
1825+
context: _ProcessContext,
1826+
result: _ProcessResult,
1827+
) -> bool:
1828+
"""
1829+
Terminates the job if it has been running longer than `max_duration` plus a grace period.
1830+
1831+
A backstop for the runner, which enforces `max_duration` itself and does it gracefully.
1832+
The server steps in only when the runner failed to stop the job -- e.g. the workload
1833+
survived the termination signals and the runner never reported the timeout.
1834+
1835+
Returns `True` if the job was terminated.
1836+
"""
1837+
job_model = context.job_model
1838+
max_duration = context.job.job_spec.max_duration
1839+
if max_duration is None:
1840+
return False
1841+
if job_model.running_at is None:
1842+
# Jobs that started running before the server was upgraded have no reference point.
1843+
# They are still enforced by the runner.
1844+
return False
1845+
deadline = (
1846+
job_model.running_at + timedelta(seconds=max_duration) + MAX_DURATION_ENFORCEMENT_GRACE
1847+
)
1848+
if get_current_datetime() < deadline:
1849+
return False
1850+
logger.warning(
1851+
"%s: max duration exceeded and the runner did not stop the job, terminating",
1852+
fmt(job_model),
1853+
)
1854+
_terminate_job(
1855+
job_model=job_model,
1856+
job_update_map=result.job_update_map,
1857+
termination_reason=JobTerminationReason.MAX_DURATION_EXCEEDED,
1858+
termination_reason_message=(
1859+
f"The job exceeded the max_duration of {max_duration} seconds"
1860+
" and did not stop on its own"
1861+
),
1862+
)
1863+
return True
1864+
1865+
18111866
def _should_terminate_job_due_to_disconnect(disconnected_at: Optional[datetime]) -> bool:
18121867
if disconnected_at is None:
18131868
return False
@@ -2088,6 +2143,10 @@ def _set_job_update_status(
20882143
) -> None:
20892144
if job_update_map.get("status", job_model.status) != new_status:
20902145
job_update_map["status"] = new_status
2146+
if new_status == JobStatus.RUNNING:
2147+
# Stamped here rather than at the call site so that `running_at` cannot drift from
2148+
# `status`: it is the reference point for server-side `max_duration` enforcement.
2149+
job_update_map["running_at"] = NOW_PLACEHOLDER
20912150

20922151

20932152
def _set_job_status(job_model: JobModel, result: _ProcessResult, new_status: JobStatus) -> None:
Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
"""Add JobModel.running_at
2+
3+
Revision ID: 652af7a3c9c9
4+
Revises: 620892d149b5
5+
Create Date: 2026-09-08 10:02:00.000000
6+
7+
"""
8+
9+
import sqlalchemy as sa
10+
from alembic import op
11+
12+
import dstack._internal.server.models
13+
14+
# revision identifiers, used by Alembic.
15+
revision = "652af7a3c9c9"
16+
down_revision = "620892d149b5"
17+
branch_labels = None
18+
depends_on = None
19+
20+
21+
def upgrade() -> None:
22+
with op.batch_alter_table("jobs", schema=None) as batch_op:
23+
batch_op.add_column(
24+
sa.Column("running_at", dstack._internal.server.models.NaiveDateTime(), nullable=True)
25+
)
26+
27+
28+
def downgrade() -> None:
29+
with op.batch_alter_table("jobs", schema=None) as batch_op:
30+
batch_op.drop_column("running_at")

‎src/dstack/_internal/server/models.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -535,6 +535,12 @@ class JobModel(PipelineModelMixin, BaseModel):
535535
job_name: Mapped[str] = mapped_column(String(100))
536536
submission_num: Mapped[int] = mapped_column(Integer)
537537
submitted_at: Mapped[datetime] = mapped_column(NaiveDateTime)
538+
running_at: Mapped[Optional[datetime]] = mapped_column(NaiveDateTime)
539+
"""`running_at` stores when the job entered the `RUNNING` status, that is, when the workload
540+
started, excluding provisioning and pulling. It is the reference point for server-side
541+
`max_duration` enforcement. `None` for jobs that never started running and for jobs that were
542+
already running before the server was upgraded.
543+
"""
538544
last_processed_at: Mapped[datetime] = mapped_column(NaiveDateTime)
539545
skip_min_processing_interval: Mapped[bool] = mapped_column(
540546
Boolean, default=False, server_default=false()

‎src/dstack/_internal/server/testing/common.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -432,6 +432,7 @@ async def create_job(
432432
submission_num: int = 0,
433433
status: JobStatus = JobStatus.SUBMITTED,
434434
submitted_at: datetime = datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc),
435+
running_at: Optional[datetime] = None,
435436
last_processed_at: datetime = datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc),
436437
termination_reason: Optional[JobTerminationReason] = None,
437438
job_provisioning_data: Optional[JobProvisioningData] = None,
@@ -475,6 +476,7 @@ async def create_job(
475476
deployment_num=deployment_num,
476477
submission_num=submission_num,
477478
submitted_at=submitted_at,
479+
running_at=running_at,
478480
last_processed_at=last_processed_at,
479481
status=status,
480482
termination_reason=termination_reason,

‎src/tests/_internal/server/background/pipeline_tasks/test_running_jobs.py‎

Lines changed: 124 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626
from dstack._internal.core.models.duration import Duration
2727
from dstack._internal.core.models.gateways import GatewayReplicaStatus, GatewayStatus
2828
from dstack._internal.core.models.instances import InstanceStatus
29-
from dstack._internal.core.models.profiles import StartupOrder, UtilizationPolicy
29+
from dstack._internal.core.models.profiles import Profile, StartupOrder, UtilizationPolicy
3030
from dstack._internal.core.models.runs import (
3131
ClusterInfo,
3232
ImagePullProgress,
@@ -44,6 +44,7 @@
4444
from dstack._internal.server import settings as server_settings
4545
from dstack._internal.server.background.pipeline_tasks.jobs_running import (
4646
JOB_DISCONNECTED_RETRY_TIMEOUT,
47+
MAX_DURATION_ENFORCEMENT_GRACE,
4748
ROUTER_PROVISIONING_WAIT_TIMEOUT_SECONDS,
4849
JobRunningFetcher,
4950
JobRunningPipeline,
@@ -1752,6 +1753,128 @@ async def test_inactivity_duration(
17521753
assert job.termination_reason == expected_termination_reason
17531754
assert job.inactivity_secs == expected_inactivity_secs
17541755

1756+
@pytest.mark.parametrize(
1757+
(
1758+
"max_duration",
1759+
"running_for",
1760+
"stamp_running_at",
1761+
"expected_status",
1762+
"expected_termination_reason",
1763+
"expect_pull",
1764+
),
1765+
[
1766+
pytest.param(
1767+
600,
1768+
timedelta(seconds=600),
1769+
True,
1770+
JobStatus.RUNNING,
1771+
None,
1772+
True,
1773+
id="deadline-reached-but-runner-still-within-grace",
1774+
),
1775+
pytest.param(
1776+
600,
1777+
timedelta(seconds=600) + MAX_DURATION_ENFORCEMENT_GRACE - timedelta(seconds=1),
1778+
True,
1779+
JobStatus.RUNNING,
1780+
None,
1781+
True,
1782+
id="grace-not-elapsed",
1783+
),
1784+
pytest.param(
1785+
600,
1786+
timedelta(seconds=600) + MAX_DURATION_ENFORCEMENT_GRACE,
1787+
True,
1788+
JobStatus.TERMINATING,
1789+
JobTerminationReason.MAX_DURATION_EXCEEDED,
1790+
False,
1791+
id="grace-elapsed",
1792+
),
1793+
pytest.param(
1794+
600,
1795+
timedelta(days=1),
1796+
False,
1797+
JobStatus.RUNNING,
1798+
None,
1799+
True,
1800+
id="job-started-before-upgrade",
1801+
),
1802+
pytest.param(
1803+
"off",
1804+
timedelta(days=1),
1805+
True,
1806+
JobStatus.RUNNING,
1807+
None,
1808+
True,
1809+
id="max-duration-off",
1810+
),
1811+
],
1812+
)
1813+
async def test_max_duration_enforced_by_server(
1814+
self,
1815+
test_db,
1816+
session: AsyncSession,
1817+
worker: JobRunningWorker,
1818+
max_duration,
1819+
running_for: timedelta,
1820+
stamp_running_at: bool,
1821+
expected_status: JobStatus,
1822+
expected_termination_reason: Optional[JobTerminationReason],
1823+
expect_pull: bool,
1824+
) -> None:
1825+
now = datetime(2023, 1, 2, 5, 0, tzinfo=timezone.utc)
1826+
project = await create_project(session=session)
1827+
user = await create_user(session=session)
1828+
repo = await create_repo(session=session, project_id=project.id)
1829+
run = await create_run(
1830+
session=session,
1831+
project=project,
1832+
repo=repo,
1833+
user=user,
1834+
status=RunStatus.RUNNING,
1835+
run_name="test-run",
1836+
run_spec=get_run_spec(
1837+
run_name="test-run",
1838+
repo_id=repo.name,
1839+
profile=Profile(name="default", max_duration=max_duration),
1840+
),
1841+
)
1842+
instance = await create_instance(
1843+
session=session,
1844+
project=project,
1845+
status=InstanceStatus.BUSY,
1846+
)
1847+
job = await create_job(
1848+
session=session,
1849+
run=run,
1850+
status=JobStatus.RUNNING,
1851+
running_at=(now - running_for) if stamp_running_at else None,
1852+
job_provisioning_data=get_job_provisioning_data(),
1853+
instance=instance,
1854+
instance_assigned=True,
1855+
)
1856+
with (
1857+
patch("dstack._internal.server.services.runner.pool.SSHTunnel"),
1858+
patch(
1859+
"dstack._internal.server.services.runner.client.RunnerClient.from_address"
1860+
) as runner_client_cls,
1861+
freeze_time(now),
1862+
):
1863+
runner_client_mock = runner_client_cls.return_value
1864+
runner_client_mock.pull.return_value = PullResponse(
1865+
job_states=[],
1866+
job_logs=[],
1867+
runner_logs=[],
1868+
last_updated=0,
1869+
no_connections_secs=0,
1870+
)
1871+
await _process_job(session, worker, job)
1872+
assert runner_client_mock.pull.called == expect_pull
1873+
1874+
await session.refresh(job)
1875+
assert job.status == expected_status
1876+
assert job.termination_reason == expected_termination_reason
1877+
17551878
@pytest.mark.parametrize(
17561879
["samples", "expected_status"],
17571880
[

0 commit comments

Comments
 (0)