Fix worker interruption handling

docs/building-workers.md

@@ -210,13 +210,14 @@ For a bounded manual check, use one of these lifecycle modes:
 # Make exactly one claim; exit immediately when no task is available.
 scimesh-worker --work-dir "$PWD/worker-data-check" --once
 
-# Keep polling until two tasks have been claimed and handled, then exit.
+# Keep polling until two tasks complete successfully, then exit.
 scimesh-worker --work-dir "$PWD/worker-data-check" --max-tasks 2
 ```
 
 `SCIMESH_MAX_TASKS` provides the same limit through the environment. Pressing
-`Ctrl+C` stops the reference worker cleanly; it reports a concise `stopped`
+`Ctrl+C` stops the reference worker cleanly. If it interrupts an active task,
-event rather than a traceback.
+the worker reports a sanitized retriable failure first, emits no traceback, and
+exits with status `130`.
 
 ## Generate a client from the spec
 

scimesh/worker/cli.py

@@ -69,13 +69,13 @@ def main(argv: list[str] | None = None) -> int:
         parser.error(str(error))
     logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
     client = HttpCoordinatorClient(config.coordinator_url, config.request_timeout, config.bearer_token)
-    WorkerDaemon(
+    completed_without_interruption = WorkerDaemon(
         config,
         client,
         HttpArtifactClient(config.coordinator_url, config.request_timeout, config.bearer_token),
         SciMeshRunner(),
     ).run_forever()
-    return 0
+    return 0 if completed_without_interruption else 130
 
 
 if __name__ == "__main__":

scimesh/worker/daemon.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import logging
 from dataclasses import replace
+from dataclasses import dataclass
 from pathlib import Path
 import random
 import shutil
@@ -72,6 +73,14 @@ class LeaseHeartbeat:
         return seconds
 
 
+@dataclass(frozen=True)
+class RunOnceOutcome:
+    """Whether a claim was made and whether that claimed task completed."""
+
+    claimed: bool
+    completed: bool
+
+
 class WorkerDaemon:
     def __init__(self, config: WorkerConfig, coordinator: CoordinatorClient, artifacts: ArtifactClient, runner: Runner) -> None:
         self.config, self.coordinator, self.artifacts, self.runner = config, coordinator, artifacts, runner
@@ -79,9 +88,10 @@ class WorkerDaemon:
         self._registered = False
         self.log = logging.getLogger("scimesh.worker")
 
-    def run_forever(self) -> None:
+    def run_forever(self) -> bool:
+        """Run until stopped; return false only when interrupted by the operator."""
         failures = 0
-        processed_tasks = 0
+        completed_tasks = 0
         self._log(
             "started",
             max_tasks=self.config.max_tasks,
@@ -93,16 +103,32 @@ class WorkerDaemon:
                     if not self._registered:
                         self._register_worker()
                     self._cleanup_expired_directories()
-                    claimed = self.run_once()
+                    outcome = self.run_once()
                     failures = 0
-                    if claimed:
+                    if outcome.claimed:
-                        processed_tasks += 1
+                        if outcome.completed:
-                        if self.config.max_tasks is not None and processed_tasks >= self.config.max_tasks:
+                            completed_tasks += 1
-                            self._log("stopped", reason="max_tasks_reached", processed_tasks=processed_tasks)
+                        if self.config.exit_when_idle:
-                            return
+                            self._log(
+                                "stopped",
+                                reason="one_claim_processed",
+                                completed_tasks=completed_tasks,
+                            )
+                            return True
+                        if (
+                            outcome.completed
+                            and self.config.max_tasks is not None
+                            and completed_tasks >= self.config.max_tasks
+                        ):
+                            self._log(
+                                "stopped",
+                                reason="max_tasks_reached",
+                                completed_tasks=completed_tasks,
+                            )
+                            return True
                     elif self.config.exit_when_idle:
-                        self._log("stopped", reason="queue_empty", processed_tasks=processed_tasks)
+                        self._log("stopped", reason="queue_empty", completed_tasks=completed_tasks)
-                        return
+                        return True
                     else:
                         self._sleep(self.config.poll_interval)
                 except CoordinatorTransientError as error:
@@ -110,18 +136,20 @@ class WorkerDaemon:
                     self._log("failed", error_type=type(error).__name__)
                     self._sleep(min(self.config.poll_interval * 2 ** min(failures, 6), 60.0))
         except KeyboardInterrupt:
-            self._log("stopped", reason="interrupted", processed_tasks=processed_tasks)
+            self._log("stopped", reason="interrupted", completed_tasks=completed_tasks)
+            return False
 
-    def run_once(self) -> bool:
+    def run_once(self) -> RunOnceOutcome:
         worker_id = self._worker_id()
         self._log("claiming", log_level=logging.DEBUG)
         task = self.coordinator.claim(worker_id, self.config.capabilities)
         if task is None:
             self._log("idle", log_level=logging.DEBUG)
-            return False
+            return RunOnceOutcome(claimed=False, completed=False)
         started = time.monotonic()
         task_dir = self.config.work_dir / task.task_id / str(task.attempt)
         heartbeat = LeaseHeartbeat(task, self.coordinator, self.config)
+        completed = False
         try:
             task_dir.mkdir(parents=True, exist_ok=False)
             heartbeat.start()
@@ -152,7 +180,15 @@ class WorkerDaemon:
                     },
                 },
             )
+            completed = True
             self._log("completed", task, elapsed_seconds=round(time.monotonic() - started, 3))
+        except KeyboardInterrupt:
+            self._log("interrupted", task)
+            try:
+                self._report_failure(task, InterruptedError("worker interrupted by operator"))
+            except CoordinatorTransientError:
+                self._log("failed", task, error_type="FailureReportError")
+            raise
         except CoordinatorConflictError as error:
             self._log("lease_lost", task, error_type=type(error).__name__)
         except Exception as error:
@@ -160,7 +196,7 @@ class WorkerDaemon:
             self._report_failure(task, error)
         finally:
             heartbeat.stop()
-        return True
+        return RunOnceOutcome(claimed=True, completed=completed)
 
     def _report_failure(self, task: ClaimedTask, error: Exception) -> None:
         message = str(error).replace(str(self.config.work_dir), "<worker-dir>")[:300]

scripts/two-worker-smoke.sh

@@ -95,7 +95,7 @@ WORKER_TWO_PID=$STARTED_WORKER_PID
 
 for _ in $(seq 1 30); do
     registered=$(docker compose -p "$COMPOSE_PROJECT" -f "$COORDINATOR_DIR/docker-compose.yml" \
-        exec -T postgres psql -U scimesh -d scimesh -Atc "SELECT count(*) FROM workers")
+        exec -T postgres psql -U scimesh -d scimesh -Atc "SELECT count(*) FROM workers" 2>/dev/null || printf '0')
     [[ "$registered" == "2" ]] && break
     sleep 1
 done

tests/test_worker_daemon.py

@@ -10,9 +10,10 @@ from urllib.request import Request
 import pytest
 
 from scimesh.worker.config import WorkerConfig
+from scimesh.worker import cli as worker_cli
 from scimesh.worker.cli import build_parser
 from scimesh.worker.coordinator import CoordinatorTransientError
-from scimesh.worker.daemon import LeaseHeartbeat, WorkerDaemon
+from scimesh.worker.daemon import LeaseHeartbeat, RunOnceOutcome, WorkerDaemon
 from scimesh.worker.models import (
     ClaimedTask,
     InputArtifact,
@@ -94,7 +95,7 @@ def daemon(tmp_path: Path, task: ClaimedTask | None, content: bytes):
 def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None:
     content = b"input fixture"
     worker, coordinator, artifacts, runner, _ = daemon(tmp_path, make_task(content), content)
-    assert worker.run_once() is True
+    assert worker.run_once() == RunOnceOutcome(claimed=True, completed=True)
     assert runner.calls == 1
     assert len(artifacts.uploaded) == 1
     assert coordinator.heartbeats == [("task-1", 1, "worker-1")]
@@ -106,7 +107,7 @@ def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None:
 
 def test_no_task_does_not_create_directory(tmp_path: Path) -> None:
     worker, _, _, runner, config = daemon(tmp_path, None, b"")
-    assert worker.run_once() is False
+    assert worker.run_once() == RunOnceOutcome(claimed=False, completed=False)
     assert runner.calls == 0
     assert not config.work_dir.exists()
 
@@ -115,7 +116,7 @@ def test_once_worker_exits_after_an_empty_claim(tmp_path: Path, caplog: pytest.L
     caplog.set_level(logging.INFO, logger="scimesh.worker")
     worker, _, _, runner, _ = daemon(tmp_path, None, b"")
     worker.config = WorkerConfig(**{**worker.config.__dict__, "exit_when_idle": True, "max_tasks": 1})
-    worker.run_forever()
+    assert worker.run_forever() is True
     assert runner.calls == 0
     assert "queue_empty" in caplog.text
 
@@ -125,7 +126,7 @@ def test_worker_stops_after_the_configured_number_of_claims(tmp_path: Path, capl
     content = b"input fixture"
     worker, _, _, runner, _ = daemon(tmp_path, make_task(content), content)
     worker.config = WorkerConfig(**{**worker.config.__dict__, "max_tasks": 1})
-    worker.run_forever()
+    assert worker.run_forever() is True
     assert runner.calls == 1
     assert "max_tasks_reached" in caplog.text
 
@@ -138,10 +139,65 @@ def test_keyboard_interrupt_stops_worker_without_propagating(tmp_path: Path, cap
 
     worker, _, _, _, _ = daemon(tmp_path, None, b"")
     worker.coordinator = InterruptingCoordinator(None)
-    worker.run_forever()
+    assert worker.run_forever() is False
     assert "interrupted" in caplog.text
 
 
+def test_interrupting_an_active_task_reports_a_sanitized_failure(tmp_path: Path) -> None:
+    content = b"input fixture"
+    worker, coordinator, _, _, _ = daemon(tmp_path, make_task(content), content)
+
+    class InterruptingRunner(FakeRunner):
+        def run(self, task: ClaimedTask, task_dir: Path) -> RunResult:
+            raise KeyboardInterrupt
+
+    worker.runner = InterruptingRunner()
+    with pytest.raises(KeyboardInterrupt):
+        worker.run_once()
+    assert coordinator.failures == [
+        {
+            "worker_id": "worker-1",
+            "attempt": 1,
+            "error_code": "InterruptedError",
+            "error_message": "worker interrupted by operator",
+        }
+    ]
+
+
+def test_max_tasks_counts_successes_not_failed_claims(tmp_path: Path) -> None:
+    successful_content = b"successful input"
+
+    class SequencedCoordinator(FakeCoordinator):
+        def __init__(self) -> None:
+            super().__init__(None)
+            self.tasks = [
+                make_task(b"bad input", "wrong-checksum"),
+                ClaimedTask(
+                    "task-2",
+                    1,
+                    (datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat(),
+                    "similarity-search",
+                    InputArtifact(
+                        "https://example.test/input",
+                        hashlib.sha256(successful_content).hexdigest(),
+                    ),
+                    {"query_id": "CHEMBL1"},
+                ),
+            ]
+
+        def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None:
+            return self.tasks.pop(0) if self.tasks else None
+
+    coordinator = SequencedCoordinator()
+    artifacts, runner = FakeArtifacts(successful_content), FakeRunner()
+    config = WorkerConfig("https://example.test", "worker-1", tmp_path / "work", max_tasks=1)
+    worker = WorkerDaemon(config, coordinator, artifacts, runner)
+    assert worker.run_forever() is True
+    assert len(coordinator.failures) == 1
+    assert len(coordinator.submissions) == 1
+    assert runner.calls == 1
+
+
 def test_worker_cli_lifecycle_options_are_explicit_and_exclusive() -> None:
     parser = build_parser()
     assert parser.parse_args(["--once"]).once is True
@@ -150,6 +206,22 @@ def test_worker_cli_lifecycle_options_are_explicit_and_exclusive() -> None:
         parser.parse_args(["--once", "--max-tasks", "2"])
 
 
+def test_worker_cli_uses_a_nonzero_exit_code_for_interruption(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: Path
+) -> None:
+    class InterruptedDaemon:
+        def __init__(self, *_: object) -> None:
+            pass
+
+        def run_forever(self) -> bool:
+            return False
+
+    monkeypatch.setattr(worker_cli, "WorkerDaemon", InterruptedDaemon)
+    assert worker_cli.main(
+        ["--coordinator-url", "https://example.test", "--work-dir", str(tmp_path)]
+    ) == 130
+
+
 @pytest.mark.parametrize("value", [0, -1, True])
 def test_max_tasks_must_be_positive(value: object, tmp_path: Path) -> None:
     with pytest.raises(ValueError, match="max_tasks"):
@@ -158,7 +230,7 @@ def test_max_tasks_must_be_positive(value: object, tmp_path: Path) -> None:
 
 def test_bad_checksum_reports_failure_without_running(tmp_path: Path) -> None:
     worker, coordinator, _, runner, _ = daemon(tmp_path, make_task(b"actual", "not-the-hash"), b"actual")
-    assert worker.run_once() is True
+    assert worker.run_once() == RunOnceOutcome(claimed=True, completed=False)
     assert runner.calls == 0
     assert coordinator.failures[0]["error_code"] == "ValueError"
     assert not coordinator.submissions
@@ -168,7 +240,7 @@ def test_directory_creation_failure_is_reported(tmp_path: Path) -> None:
     content = b"input fixture"
     worker, coordinator, _, _, config = daemon(tmp_path, make_task(content), content)
     (config.work_dir / "task-1" / "1").mkdir(parents=True)
-    assert worker.run_once() is True
+    assert worker.run_once() == RunOnceOutcome(claimed=True, completed=False)
     assert coordinator.failures[0]["error_code"] == "FileExistsError"