Fail background jobs interrupted by shutdown

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",

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,

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

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: