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 c5cd9c7b6..4b2ddf6e2 100644 --- a/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py +++ b/src/dstack/_internal/server/background/pipeline_tasks/jobs_terminating.py @@ -11,7 +11,7 @@ from dstack._internal.core.backends.base.backend import Backend from dstack._internal.core.backends.base.compute import ComputeWithVolumeSupport -from dstack._internal.core.consts import DSTACK_SHIM_HTTP_PORT +from dstack._internal.core.consts import DSTACK_RUNNER_HTTP_PORT, DSTACK_SHIM_HTTP_PORT from dstack._internal.core.errors import BackendError from dstack._internal.core.models.instances import InstanceStatus, InstanceTerminationReason from dstack._internal.core.models.runs import ( @@ -50,6 +50,7 @@ ) from dstack._internal.server.services import backends as backends_services from dstack._internal.server.services import events +from dstack._internal.server.services import logs as logs_services from dstack._internal.server.services.instances import ( emit_instance_status_change_event, get_instance_ssh_private_keys, @@ -81,6 +82,12 @@ logger = get_logger(__name__) +# How long a job is given to finish and hand over its last logs before its container is killed. +# Covers the runner's own termination staging -- SIGHUP after 5s, SIGKILL after 10s -- and the +# final log flush that follows it. +JOB_TERMINATION_DEADLINE = timedelta(seconds=30) + + @dataclass class JobTerminatingPipelineItem(PipelineItem): volumes_detached_at: Optional[datetime] @@ -173,10 +180,6 @@ async def fetch(self, limit: int) -> list[JobTerminatingPipelineItem]: select(JobModel) .where( JobModel.status == JobStatus.TERMINATING, - or_( - JobModel.remove_at.is_(None), - JobModel.remove_at < now, - ), or_( # Processing volumes detach can be less frequent since it may take time. and_( @@ -302,6 +305,7 @@ class _JobUpdateMap(ItemUpdateMap, total=False): termination_reason_message: Optional[str] instance_id: Optional[uuid.UUID] graceful_termination_attempts: int + runner_timestamp: Optional[int] volumes_detached_at: UpdateMapDateTime registered: bool remove_at: UpdateMapDateTime @@ -659,9 +663,28 @@ async def _process_terminating_job( result.job_update_map["status"] = _get_job_termination_status(job_model) return result - if job_model.graceful_termination_attempts == 0 and job_model.remove_at is None: - result.job_update_map = await _stop_job_gracefully(job_model, instance_model) - result.graceful_stop_event_message = "Graceful job stop requested" + if job_model.remove_at is None: + # The first terminating pass. `graceful_termination_attempts` decides whether the runner + # is asked to stop the job; `remove_at` decides when the container is killed regardless. + # The two are independent: a job that has already finished on its own is not asked to + # stop, but its logs are still drained below. + graceful = job_model.graceful_termination_attempts == 0 + if graceful: + await stop_runner(job_model=job_model, instance_model=instance_model) + result.job_update_map["graceful_termination_attempts"] = 1 + result.graceful_stop_event_message = "Graceful job stop requested" + if graceful or _has_logs_to_drain(job_model): + result.job_update_map["remove_at"] = get_current_datetime() + JOB_TERMINATION_DEADLINE + return result + # Nothing to wait for, stop the container right away + elif get_current_datetime() < job_model.remove_at and not await _drain_job_logs( + job_model=job_model, + instance_model=instance_model, + job_update_map=result.job_update_map, + ): + # The runner still has logs to hand over and there is time left to collect them. They + # must be collected before the container is stopped, since that destroys the runner + # along with everything it has buffered. return result jrd = get_job_runtime_data(job_model) @@ -718,18 +741,82 @@ async def _process_terminating_job( return result -async def _stop_job_gracefully( - job_model: JobModel, instance_model: InstanceModel -) -> _JobUpdateMap: +def _has_logs_to_drain(job_model: JobModel) -> bool: + """ + Whether the runner may still be holding logs for this job. + + `running_at` says the workload started, but it is only stamped by servers new enough to have + the column, so `runner_timestamp` -- advanced on every successful pull -- covers jobs that + were already running before the upgrade. A job that never started running has neither, and + waiting for logs it cannot have would only delay its termination. + """ + return job_model.running_at is not None or job_model.runner_timestamp is not None + + +async def _drain_job_logs( + job_model: JobModel, + instance_model: InstanceModel, + job_update_map: _JobUpdateMap, +) -> bool: """ - Tells the runner to stop the job's command. Records the first graceful-stop attempt and - sets `remove_at` so `_process_terminating_job()` stops the container on a later iteration. + Collects the logs the runner has buffered since the last pull. + + Returns whether the runner has nothing left to hand over, or cannot be asked at all -- the + caller keeps the container alive until then, or until `remove_at` passes. """ - job_update_map = _JobUpdateMap() - await stop_runner(job_model=job_model, instance_model=instance_model) - job_update_map["graceful_termination_attempts"] = 1 - job_update_map["remove_at"] = get_current_datetime() + timedelta(seconds=10) - return job_update_map + jpd = get_job_provisioning_data(job_model) + if jpd is None: + return True + jrd = get_job_runtime_data(job_model) + ssh_private_keys = get_instance_ssh_private_keys(instance_model) + try: + return await common.run_async( + _pull_job_logs, + ssh_private_keys, + jpd, + jrd, + job_model.run, + job_model, + job_update_map, + ) + except client.PeerConnectionError as e: + # An unreachable runner has nothing left to give, and waiting out the deadline would + # only delay terminating a job whose instance is likely gone already. + logger.debug("%s: can't collect the last logs: %s", fmt(job_model), e) + return True + except client.RunnerError as e: + # The runner answered, but not usefully. Collecting the logs is best-effort: the job is + # being terminated either way, and retrying until the deadline would only delay it. + logger.warning("%s: runner failed to hand over the last logs: %s", fmt(job_model), e) + return True + + +@runner_ssh_tunnel +def _pull_job_logs( + addresses: Mapping[int, client.LocalAddress], + run_model: RunModel, + job_model: JobModel, + job_update_map: _JobUpdateMap, +) -> bool: + runner_client = client.RunnerClient.from_address(addresses[DSTACK_RUNNER_HTTP_PORT]) + resp = runner_client.pull(job_model.runner_timestamp or 0) + try: + logs_services.write_logs( + project=run_model.project, + run_name=run_model.run_name, + job_submission_id=job_model.id, + runner_logs=resp.runner_logs, + job_logs=resp.job_logs, + ) + except logs_services.LogStorageError as e: + # `runner_timestamp` is not advanced, so the same logs are pulled again on the next pass + # instead of being lost. + logger.error("%s: failed to write the last logs: %s", fmt(job_model), e) + return False + job_update_map["runner_timestamp"] = resp.last_updated + # An old runner does not report `has_more`, and then there is no way to tell when the logs + # are exhausted -- take what this pull returned and stop. + return not resp.has_more async def _process_job_volumes_detaching( diff --git a/src/dstack/_internal/server/models.py b/src/dstack/_internal/server/models.py index 8fb9113a9..a80f4922a 100644 --- a/src/dstack/_internal/server/models.py +++ b/src/dstack/_internal/server/models.py @@ -568,8 +568,9 @@ class JobModel(PipelineModelMixin, BaseModel): * `>= 1` means at least one graceful stop attempt was sent. """ remove_at: Mapped[Optional[datetime]] = mapped_column(NaiveDateTime) - """`remove_at` is used to ensure the container/instance is killed after the job is gracefully finished. - Cannot kill the container/instance until `remove_at` is set. + """`remove_at` is when the job's container is killed, whether or not the runner has handed + over its last logs. `None` until the job starts terminating, and only set for jobs that are + given time to finish -- the rest are stopped on their first terminating pass. """ volumes_detached_at: Mapped[Optional[datetime]] = mapped_column(NaiveDateTime) instance_assigned: Mapped[bool] = mapped_column(Boolean, default=False) diff --git a/src/dstack/_internal/server/schemas/runner.py b/src/dstack/_internal/server/schemas/runner.py index 532a8e93f..a9aa40492 100644 --- a/src/dstack/_internal/server/schemas/runner.py +++ b/src/dstack/_internal/server/schemas/runner.py @@ -47,6 +47,10 @@ class PullResponse(CoreModel): last_updated: int no_connections_secs: Optional[int] = None """`no_connections_secs` is optional for compatibility with old runners.""" + has_more: Optional[bool] = None + """`has_more` tells whether the runner may still have logs to hand over. + It is optional for compatibility with runners that do not report it. + """ class JobInfoResponse(CoreModel): diff --git a/src/dstack/_internal/server/services/jobs/__init__.py b/src/dstack/_internal/server/services/jobs/__init__.py index 4e3d26a5b..ffd3078c2 100644 --- a/src/dstack/_internal/server/services/jobs/__init__.py +++ b/src/dstack/_internal/server/services/jobs/__init__.py @@ -1,7 +1,6 @@ import itertools import json from collections.abc import Mapping -from datetime import timedelta from typing import Dict, Iterable, List, Optional, Tuple from uuid import UUID @@ -331,10 +330,6 @@ def job_spec_updatable_in_place(old_job_spec: JobSpec, new_job_spec: JobSpec) -> return old_job_spec == new_job_spec -def delay_job_instance_termination(job_model: JobModel): - job_model.remove_at = common.get_current_datetime() + timedelta(seconds=15) - - def is_multinode_job(job: Job) -> bool: return job.job_spec.jobs_per_replica > 1 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 f6f93613d..80e062246 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,8 @@ import asyncio import uuid from datetime import datetime, timedelta, timezone -from unittest.mock import AsyncMock, Mock, patch +from typing import Optional +from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest from sqlalchemy import select @@ -13,6 +14,7 @@ from dstack._internal.core.models.instances import InstanceStatus from dstack._internal.core.models.runs import JobStatus, JobTerminationReason 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 ( JobTerminatingFetcher, JobTerminatingPipeline, @@ -21,6 +23,12 @@ _get_related_instance_lock_owner, ) from dstack._internal.server.models import InstanceModel, JobModel, VolumeAttachmentModel +from dstack._internal.server.schemas.runner import LogEvent, PullResponse +from dstack._internal.server.services.runner.client import ( + PeerConnectionError, + RunnerClient, + RunnerError, +) from dstack._internal.server.testing.common import ( ComputeMockSpec, create_instance, @@ -46,6 +54,23 @@ def worker() -> JobTerminatingWorker: return JobTerminatingWorker(queue=Mock(), heartbeater=Mock(), pipeline_hinter=Mock()) +@pytest.fixture +def ssh_tunnel_mock(monkeypatch: pytest.MonkeyPatch) -> Mock: + mock = MagicMock(spec_set=SSHTunnel) + monkeypatch.setattr("dstack._internal.server.services.runner.pool.SSHTunnel", mock) + return mock + + +@pytest.fixture +def runner_client_mock(monkeypatch: pytest.MonkeyPatch) -> Mock: + mock = Mock(spec_set=RunnerClient) + monkeypatch.setattr( + "dstack._internal.server.services.runner.client.RunnerClient.from_address", + Mock(return_value=mock), + ) + return mock + + @pytest.fixture def fetcher() -> JobTerminatingFetcher: return JobTerminatingFetcher( @@ -107,6 +132,8 @@ async def test_fetch_selects_eligible_jobs_and_sets_lock_fields( past_remove_at.remove_at = stale past_remove_at.volumes_detached_at = stale - timedelta(seconds=30) + # `remove_at` is when the container is killed, not a condition for processing the job: + # the job is still fetched so that its logs can be collected in the meantime future_remove_at = await create_job( session=session, run=run, @@ -168,12 +195,14 @@ async def test_fetch_selects_eligible_jobs_and_sets_lock_fields( assert [item.id for item in items] == [ terminating.id, past_remove_at.id, + future_remove_at.id, expired_same_owner.id, recent_skip.id, ] assert {(item.id, item.volumes_detached_at) for item in items} == { (terminating.id, None), (past_remove_at.id, past_remove_at.volumes_detached_at), + (future_remove_at.id, None), (expired_same_owner.id, None), (recent_skip.id, None), } @@ -190,14 +219,19 @@ async def test_fetch_selects_eligible_jobs_and_sets_lock_fields( ]: await session.refresh(job) - fetched_jobs = [terminating, past_remove_at, expired_same_owner, recent_skip] + fetched_jobs = [ + terminating, + past_remove_at, + future_remove_at, + expired_same_owner, + recent_skip, + ] assert all(job.lock_owner == JobTerminatingPipeline.__name__ for job in fetched_jobs) assert all(job.lock_expires_at is not None for job in fetched_jobs) assert all(job.lock_token is not None for job in fetched_jobs) assert all(not job.skip_min_processing_interval for job in fetched_jobs) assert len({job.lock_token for job in fetched_jobs}) == 1 - assert future_remove_at.lock_owner is None assert non_terminating.lock_owner is None assert recent.lock_owner is None assert locked.lock_owner == "OtherPipeline" @@ -344,6 +378,219 @@ async def test_stops_job_gracefully_before_terminating_container( events = await list_events(session) assert any(event.message == "Graceful job stop requested" for event in events) + async def test_gives_finished_job_time_to_hand_over_logs( + self, test_db, session: AsyncSession, worker: JobTerminatingWorker + ): + """A job that finished on its own is not asked to stop, but its logs are still waited for.""" + 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) + job = await create_job( + session=session, + run=run, + status=JobStatus.TERMINATING, + termination_reason=JobTerminationReason.DONE_BY_RUNNER, + job_provisioning_data=get_job_provisioning_data(dockerized=True), + instance=instance, + running_at=get_current_datetime() - timedelta(minutes=1), + ) + _lock_job(job) + await session.commit() + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_terminating.stop_runner", + new=AsyncMock(), + ) as stop_runner, + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_terminating._stop_container", + new=AsyncMock(return_value=True), + ) as stop_container, + ): + await worker.process(_job_to_pipeline_item(job)) + + stop_runner.assert_not_awaited() + stop_container.assert_not_awaited() + + await session.refresh(job) + assert job.status == JobStatus.TERMINATING + assert job.graceful_termination_attempts is None + assert job.remove_at is not None + + async def test_terminates_job_that_never_ran_without_waiting( + self, test_db, session: AsyncSession, worker: JobTerminatingWorker + ): + """There are no logs to wait for if the workload never started.""" + 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) + job = await create_job( + session=session, + run=run, + status=JobStatus.TERMINATING, + termination_reason=JobTerminationReason.TERMINATED_BY_SERVER, + job_provisioning_data=get_job_provisioning_data(dockerized=True), + instance=instance, + ) + _lock_job(job) + await session.commit() + assert job.running_at is None and job.runner_timestamp is None + + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_terminating.stop_runner", + new=AsyncMock(), + ) as stop_runner, + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_terminating._stop_container", + new=AsyncMock(return_value=True), + ) as stop_container, + ): + await worker.process(_job_to_pipeline_item(job)) + + stop_runner.assert_not_awaited() + stop_container.assert_awaited_once() + + await session.refresh(job) + assert job.remove_at is None + + @pytest.mark.parametrize( + ("has_more", "container_stopped"), + [ + # The runner still has logs buffered, the container must stay up to hand them over + pytest.param(True, False, id="has_more"), + pytest.param(False, True, id="drained"), + # An old runner does not report `has_more`, so there is no way to wait for the end + pytest.param(None, True, id="not_reported"), + ], + ) + async def test_collects_last_logs_before_terminating_container( + self, + test_db, + session: AsyncSession, + worker: JobTerminatingWorker, + ssh_tunnel_mock: Mock, + runner_client_mock: Mock, + has_more: Optional[bool], + container_stopped: bool, + ): + 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) + 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=True), + instance=instance, + running_at=get_current_datetime() - timedelta(minutes=1), + ) + job.graceful_termination_attempts = 1 + job.remove_at = get_current_datetime() + timedelta(seconds=30) + job.runner_timestamp = 1 + _lock_job(job) + await session.commit() + + runner_client_mock.pull.return_value = PullResponse( + job_states=[], + job_logs=[LogEvent(timestamp=2, message=b"the tail")], + runner_logs=[LogEvent(timestamp=3, message=b"Job state changed")], + last_updated=3, + has_more=has_more, + ) + with ( + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_terminating.logs_services.write_logs" + ) as write_logs, + patch( + "dstack._internal.server.background.pipeline_tasks.jobs_terminating._stop_container", + new=AsyncMock(return_value=True), + ) as stop_container, + ): + await worker.process(_job_to_pipeline_item(job)) + + runner_client_mock.pull.assert_called_once_with(1) + write_logs.assert_called_once() + assert write_logs.call_args.kwargs["job_logs"][0].message == b"the tail" + assert stop_container.await_count == int(container_stopped) + + await session.refresh(job) + # Advanced, so the next pass does not collect the same logs again + assert job.runner_timestamp == 3 + + @pytest.mark.parametrize( + "error", + [ + pytest.param(RunnerError("runner is confused"), id="runner_error"), + pytest.param(PeerConnectionError("instance is gone"), id="unreachable"), + ], + ) + async def test_terminates_container_when_last_logs_cannot_be_collected( + self, + test_db, + session: AsyncSession, + worker: JobTerminatingWorker, + ssh_tunnel_mock: Mock, + runner_client_mock: Mock, + error: Exception, + ): + """Collecting the last logs is best-effort: a failure must not hold up the termination.""" + 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) + 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=True), + instance=instance, + running_at=get_current_datetime() - timedelta(minutes=1), + ) + job.graceful_termination_attempts = 1 + job.remove_at = get_current_datetime() + timedelta(seconds=30) + _lock_job(job) + await session.commit() + + runner_client_mock.pull.side_effect = error + with patch( + "dstack._internal.server.background.pipeline_tasks.jobs_terminating._stop_container", + new=AsyncMock(return_value=True), + ) as stop_container: + await worker.process(_job_to_pipeline_item(job)) + + stop_container.assert_awaited_once() + + await session.refresh(job) + # The job must not be left locked for the pipeline to trip over + assert job.lock_owner is None + assert job.lock_token is None + async def test_terminates_gracefully_stopped_job_after_remove_at( self, test_db, session: AsyncSession, worker: JobTerminatingWorker ):