diff --git a/src/libkernelbot/background_submission_manager.py b/src/libkernelbot/background_submission_manager.py index ed3a9cd2..92cfddea 100644 --- a/src/libkernelbot/background_submission_manager.py +++ b/src/libkernelbot/background_submission_manager.py @@ -276,6 +276,23 @@ async def _mark_job_failed_after_worker_crash(self, item: JobItem): exc_info=True, ) + async def _mark_job_failed_after_shutdown(self, item: JobItem): + ts = dt.datetime.now(dt.timezone.utc) + try: + with self.backend.db as db: + db.fail_submission_job_if_active( + item.sub_id, + "job interrupted while kernelbot was shutting down; please resubmit", + ts, + ) + except Exception: + logger.error( + "[Background Job][worker %r] failed to mark interrupted submission job `%s`", + id(asyncio.current_task()), + item.sub_id, + exc_info=True, + ) + async def _worker_loop(self): """ A worker will keep listening to the queue, and process the job in the queue. @@ -307,6 +324,14 @@ async def _worker_loop(self): self._live_tasks.add(t) try: await t # wait submission job to finish + except asyncio.CancelledError: + logger.info( + "[Background Job][worker %r] submission job `%s` interrupted by shutdown", + id(asyncio.current_task()), + item.sub_id, + ) + await self._mark_job_failed_after_shutdown(item) + raise except Exception: logger.error( "[Background Job][worker %r] submission job `%s` crashed", diff --git a/src/libkernelbot/leaderboard_db.py b/src/libkernelbot/leaderboard_db.py index 9464f861..2ed15acb 100644 --- a/src/libkernelbot/leaderboard_db.py +++ b/src/libkernelbot/leaderboard_db.py @@ -425,6 +425,33 @@ def update_heartbeat_if_active(self, sub_id: int, ts: datetime.datetime) -> None logger.error("Failed to upsert submission job status. sub_id: '%s'", sub_id, exc_info=e) raise KernelBotError("Error updating job status") from e + def fail_submission_job_if_active( + self, sub_id: int, error: str, ts: datetime.datetime + ) -> bool: + """Fail a pending/running job without overwriting a terminal status.""" + try: + self.cursor.execute( + """ + UPDATE leaderboard.submission_job_status + SET status = 'failed', error = %s, last_heartbeat = %s + WHERE submission_id = %s + AND status IN ('pending','running') + RETURNING id + """, + (error, ts, sub_id), + ) + updated = self.cursor.fetchone() is not None + self.connection.commit() + return updated + except psycopg2.Error as e: + self.connection.rollback() + logger.error( + "Failed to mark active submission job as failed. sub_id: '%s'", + sub_id, + exc_info=e, + ) + raise KernelBotError("Error failing active job") from e + def upsert_submission_job_status( self, sub_id: int, diff --git a/tests/test_background_submission_manager.py b/tests/test_background_submission_manager.py index c0fd6374..c824b892 100644 --- a/tests/test_background_submission_manager.py +++ b/tests/test_background_submission_manager.py @@ -10,6 +10,7 @@ from libkernelbot.kernelguard import KernelGuardRejected from libkernelbot.run_eval import FullResult from libkernelbot.submission import ProcessedSubmissionRequest +from libkernelbot.task import make_task_definition @pytest.fixture @@ -225,6 +226,84 @@ async def test_stop_rejects_new_jobs(mock_backend): await manager.enqueue(req, SubmissionMode.TEST, 99) +@pytest.mark.asyncio +async def test_stop_marks_running_job_failed(mock_backend): + db_context = mock_backend.db + db_context.upsert_submission_job_status = mock.Mock(side_effect=lambda *args, **kwargs: args[0]) + db_context.update_heartbeat_if_active = mock.Mock() + db_context.fail_submission_job_if_active = mock.Mock(return_value=True) + started = asyncio.Event() + + async def fake_submit_full(req, mode, reporter, sub_id, skip_precheck=False): + started.set() + await asyncio.Event().wait() + + mock_backend.submit_full = fake_submit_full + + manager = BackgroundSubmissionManager( + mock_backend, min_workers=1, max_workers=1, idle_seconds=0.1 + ) + await manager.start() + await manager.enqueue(get_req(1), SubmissionMode.TEST, sub_id=99) + await asyncio.wait_for(started.wait(), timeout=1) + + await manager.stop() + + db_context.fail_submission_job_if_active.assert_called_once_with( + 99, + "job interrupted while kernelbot was shutting down; please resubmit", + mock.ANY, + ) + + +@pytest.mark.asyncio +async def test_stop_marks_real_database_job_failed(database, task_directory): + definition = make_task_definition(task_directory / "task.yml") + now = datetime.datetime.now(tz=datetime.timezone.utc) + with database as db: + db.create_leaderboard( + name="shutdown-test", + deadline=now + datetime.timedelta(days=1), + definition=definition, + creator_id=1, + forum_id=1, + gpu_types=["A100"], + ) + sub_id = db.create_submission( + "shutdown-test", + "submission.py", + 1, + "print('hello')", + now, + user_name="shutdown-test-user", + ) + + backend = mock.Mock() + backend.db = database + started = asyncio.Event() + + async def fake_submit_full(req, mode, reporter, sub_id, skip_precheck=False): + started.set() + await asyncio.Event().wait() + + backend.submit_full = fake_submit_full + manager = BackgroundSubmissionManager( + backend, min_workers=1, max_workers=1, idle_seconds=0.1 + ) + await manager.start() + await manager.enqueue(get_req(1), SubmissionMode.TEST, sub_id=sub_id) + await asyncio.wait_for(started.wait(), timeout=1) + + await manager.stop() + + with database as db: + submission = db.get_submission_by_id(sub_id) + assert submission["job_status"] == "failed" + assert submission["job_error"] == ( + "job interrupted while kernelbot was shutting down; please resubmit" + ) + + @pytest.mark.asyncio async def test_scale_up_and_down(mock_backend): db_context = mock_backend.db diff --git a/tests/test_leaderboard_db.py b/tests/test_leaderboard_db.py index 621b7391..6e5eb535 100644 --- a/tests/test_leaderboard_db.py +++ b/tests/test_leaderboard_db.py @@ -210,6 +210,21 @@ def test_leaderboard_submission_basic(database, submit_leaderboard): assert submission["job_error"] == "still working" assert submission["job_last_heartbeat"] == heartbeat + interrupted_at = datetime.datetime.now(tz=datetime.timezone.utc) + assert db.fail_submission_job_if_active( + sub_id, "interrupted by shutdown", interrupted_at + ) + submission = db.get_submission_by_id(sub_id) + assert submission["job_status"] == "failed" + assert submission["job_error"] == "interrupted by shutdown" + assert submission["job_last_heartbeat"] == interrupted_at + + assert not db.fail_submission_job_if_active( + sub_id, "must not replace a terminal status", interrupted_at + ) + submission = db.get_submission_by_id(sub_id) + assert submission["job_error"] == "interrupted by shutdown" + # add a submission run run_result = sample_run_result() with database as db: