From 9ec8f503136607a0eaf77c3be23d0477ab23212d Mon Sep 17 00:00:00 2001 From: Emil Date: Fri, 24 Jul 2026 13:19:40 +0300 Subject: [PATCH] Fix worker interruption handling --- docs/building-workers.md | 7 +-- scimesh/worker/cli.py | 4 +- scimesh/worker/daemon.py | 64 +++++++++++++++++++++------ scripts/two-worker-smoke.sh | 2 +- tests/test_worker_daemon.py | 88 +++++++++++++++++++++++++++++++++---- 5 files changed, 137 insertions(+), 28 deletions(-) diff --git a/docs/building-workers.md b/docs/building-workers.md index db4127d..8e23884 100644 --- a/docs/building-workers.md +++ b/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` -event rather than a traceback. +`Ctrl+C` stops the reference worker cleanly. If it interrupts an active task, +the worker reports a sanitized retriable failure first, emits no traceback, and +exits with status `130`. ## Generate a client from the spec diff --git a/scimesh/worker/cli.py b/scimesh/worker/cli.py index b74443b..422c178 100644 --- a/scimesh/worker/cli.py +++ b/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__": diff --git a/scimesh/worker/daemon.py b/scimesh/worker/daemon.py index 50897bd..3ecb742 100644 --- a/scimesh/worker/daemon.py +++ b/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: - processed_tasks += 1 - if self.config.max_tasks is not None and processed_tasks >= self.config.max_tasks: - self._log("stopped", reason="max_tasks_reached", processed_tasks=processed_tasks) - return + if outcome.claimed: + if outcome.completed: + completed_tasks += 1 + if self.config.exit_when_idle: + 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) - return + self._log("stopped", reason="queue_empty", completed_tasks=completed_tasks) + 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), "")[:300] diff --git a/scripts/two-worker-smoke.sh b/scripts/two-worker-smoke.sh index e551bc7..c8d96a5 100755 --- a/scripts/two-worker-smoke.sh +++ b/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 diff --git a/tests/test_worker_daemon.py b/tests/test_worker_daemon.py index c1bab39..031098b 100644 --- a/tests/test_worker_daemon.py +++ b/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"