fix: harden job cleanup and stuck job recovery

Author: lkstrp

app/backends/local.py

@@ -320,6 +320,8 @@ async def cleanup(
             try:
                 pid = int(pid_path.read_text().strip())
                 os.kill(pid, signal.SIGTERM)
+                await asyncio.sleep(5)
+                os.kill(pid, signal.SIGKILL)
             except ProcessLookupError as exc:
                 logger.debug("Process already gone for job %s: %s", job_id, exc)
             except (ValueError, OSError) as exc:

app/backends/slurm_ssh.py

@@ -19,7 +19,6 @@
 from app.config import SlurmSSHConfig
 from app.models import JobStatus, WorkflowFileInfo
 from app.utils import (
-    build_snkmt_setup_commands,
     build_wrapper_script,
     enforce_error_limit,
     rename_with_cleanup,
@@ -527,9 +526,17 @@ async def cleanup(
         work_dir: str,
     ) -> None:
         # 1. Kill the wrapper process if a PID file exists
+        pid_file = f"{shlex.quote(work_dir)}/.pid"
         await self._run_ssh(
-            f"test -f {shlex.quote(work_dir)}/.pid && "
-            f"kill $(cat {shlex.quote(work_dir)}/.pid) 2>/dev/null || true",
+            f"test -f {pid_file} && "
+            f"kill $(cat {pid_file}) 2>/dev/null || true",
+            check=False,
+        )
+        # Wait before sending SIGKILL in case SIGTERM wasn't enough
+        await asyncio.sleep(5)
+        await self._run_ssh(
+            f"test -f {pid_file} && "
+            f"kill -9 $(cat {pid_file}) 2>/dev/null || true",
             check=False,
         )
         # 2. Cancel only SLURM jobs launched from this work_dir

app/tasks.py

@@ -179,8 +179,9 @@ def log_callback(line: str) -> None:
         )
 
     except asyncio.CancelledError:
-        logger.info("Job %s was cancelled", job_id)
-        store._mark_cancelled(job_id)
+        if record.status not in TERMINAL_STATUSES:
+            logger.info("Job %s was cancelled", job_id)
+            store._mark_cancelled(job_id)
         raise
 
     except Exception as exc:
@@ -224,27 +225,27 @@ async def sync_job_data_loop(
             if not record.work_dir:
                 continue
 
-            # Recover orphaned RUNNING jobs (no active monitor task)
-            if record.task is None or record.task.done():
-                try:
-                    exit_code = await backend.check_job_status(
-                        job_id, record.work_dir
-                    )
-                    if exit_code is not None:
-                        logger.warning(
-                            "Recovering orphaned job %s (exit code %d)",
-                            job_id,
-                            exit_code,
-                        )
-                        store.mark_finished(job_id, exit_code)
-                        store.persist(record)
-                        continue
-                except Exception:
+            try:
+                exit_code = await backend.check_job_status(
+                    job_id, record.work_dir
+                )
+                if exit_code is not None:
                     logger.warning(
-                        "Failed to check orphaned job %s",
+                        "Recovering stuck job %s (exit code %d)",
                         job_id,
-                        exc_info=True,
+                        exit_code,
                     )
+                    store.mark_finished(job_id, exit_code)
+                    store.persist(record)
+                    if record.task and not record.task.done():
+                        record.task.cancel()
+                    continue
+            except Exception:
+                logger.warning(
+                    "Failed to check job %s",
+                    job_id,
+                    exc_info=True,
+                )
 
             _flush_logs(store, job_id)
             snkmt_db_path = store.get_snkmt_db_path(job_id)

tests/test_slurm_ssh.py

@@ -527,7 +527,7 @@ async def mock_connect(*args, **kwargs):
 
 
 class TestCleanupSlurm:
-    async def test_cleanup_issues_three_separate_ssh_commands(self, backend):
+    async def test_cleanup_issues_four_separate_ssh_commands(self, backend):
         ssh_calls: list[str] = []
 
         async def run_ssh(cmd, **kwargs):
@@ -540,7 +540,8 @@ async def run_ssh(cmd, **kwargs):
         with patch.object(backend, "_run_ssh", side_effect=run_ssh):
             await backend.cleanup("job-1", "/scratch/jobs/job-1")
 
-        assert len(ssh_calls) == 3
-        assert any("kill" in cmd for cmd in ssh_calls)
+        assert len(ssh_calls) == 4
+        assert any("kill" in cmd and "kill -9" not in cmd for cmd in ssh_calls)
+        assert any("kill -9" in cmd for cmd in ssh_calls)
         assert any("scancel" in cmd for cmd in ssh_calls)
         assert any("rm -rf" in cmd for cmd in ssh_calls)