From bde6cdb4ba0c7aac87ebc03a27ba4b3f40567356 Mon Sep 17 00:00:00 2001 From: Emil Date: Fri, 24 Jul 2026 13:08:10 +0300 Subject: [PATCH] Improve worker lifecycle controls --- docs/building-workers.md | 25 ++++++++++++++++ scimesh/worker/cli.py | 33 +++++++++++++++++--- scimesh/worker/config.py | 14 +++++++++ scimesh/worker/daemon.py | 60 ++++++++++++++++++++++++++----------- scripts/two-worker-smoke.sh | 2 +- tests/test_worker_daemon.py | 47 +++++++++++++++++++++++++++++ 6 files changed, 158 insertions(+), 23 deletions(-) diff --git a/docs/building-workers.md b/docs/building-workers.md index 4fa7bbd..db4127d 100644 --- a/docs/building-workers.md +++ b/docs/building-workers.md @@ -193,6 +193,31 @@ Per the worker contract, at minimum: - poll interval and request timeout - a working directory for downloaded inputs and generated outputs +## Run the reference worker locally + +Use one terminal per worker and a distinct work directory for each process: + +```sh +SCIMESH_COORDINATOR_URL=http://localhost:8080 \ +SCIMESH_BEARER_TOKEN=dev-token \ +SCIMESH_WORKER_NAME=worker-1 \ +scimesh-worker --work-dir "$PWD/worker-data-1" +``` + +For a bounded manual check, use one of these lifecycle modes: + +```sh +# 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. +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. + ## Generate a client from the spec Instead of hand-writing request code, generate it: diff --git a/scimesh/worker/cli.py b/scimesh/worker/cli.py index 13e1fb5..b74443b 100644 --- a/scimesh/worker/cli.py +++ b/scimesh/worker/cli.py @@ -13,15 +13,17 @@ from .daemon import WorkerDaemon from .runners import SciMeshRunner -def main(argv: list[str] | None = None) -> int: +def build_parser() -> argparse.ArgumentParser: + """Build the worker CLI parser for command-line use and focused tests.""" parser = argparse.ArgumentParser( prog="scimesh-worker", epilog=( "Environment: SCIMESH_COORDINATOR_URL, SCIMESH_WORK_DIR, " "SCIMESH_WORKER_NAME, SCIMESH_CPU_COUNT, SCIMESH_MEMORY_MB, " "SCIMESH_POLL_INTERVAL, SCIMESH_REQUEST_TIMEOUT, " - "SCIMESH_HEARTBEAT_INTERVAL, SCIMESH_CLEANUP_AFTER_SECONDS, and " - "SCIMESH_BEARER_TOKEN. SCIMESH_WORKER_ID is a legacy/test override." + "SCIMESH_HEARTBEAT_INTERVAL, SCIMESH_CLEANUP_AFTER_SECONDS, " + "SCIMESH_MAX_TASKS, and SCIMESH_BEARER_TOKEN. " + "SCIMESH_WORKER_ID is a legacy/test override." ), ) parser.add_argument("--coordinator-url") @@ -34,8 +36,31 @@ def main(argv: list[str] | None = None) -> int: parser.add_argument("--request-timeout", type=float) parser.add_argument("--heartbeat-interval", type=float) parser.add_argument("--cleanup-after-seconds", type=float) + lifecycle = parser.add_mutually_exclusive_group() + lifecycle.add_argument( + "--once", + action="store_true", + help="Claim at most one task, then exit; exit immediately when the queue is empty", + ) + lifecycle.add_argument( + "--max-tasks", + type=int, + help="Process this many claimed tasks, then exit", + ) + return parser + + +def main(argv: list[str] | None = None) -> int: + parser = build_parser() args = parser.parse_args(argv) - overrides = {key: value for key, value in vars(args).items() if value is not None} + overrides = { + key: value + for key, value in vars(args).items() + if value is not None and key != "once" + } + if args.once: + overrides["max_tasks"] = 1 + overrides["exit_when_idle"] = True if "work_dir" in overrides: overrides["work_dir"] = Path(overrides["work_dir"]) try: diff --git a/scimesh/worker/config.py b/scimesh/worker/config.py index a0420eb..c0d5c9a 100644 --- a/scimesh/worker/config.py +++ b/scimesh/worker/config.py @@ -36,6 +36,8 @@ class WorkerConfig: heartbeat_interval: float = 15.0 bearer_token: str | None = None cleanup_after_seconds: float | None = None + max_tasks: int | None = None + exit_when_idle: bool = False # The local CLI uses hyphens; the first coordinator contract used # underscores. Advertise both stable spellings while jobs are migrated. capabilities: tuple[str, ...] = ( @@ -66,6 +68,15 @@ class WorkerConfig: _positive_number(self.heartbeat_interval, "heartbeat_interval") if self.cleanup_after_seconds is not None: _positive_number(self.cleanup_after_seconds, "cleanup_after_seconds", allow_zero=True) + if self.max_tasks is not None: + if ( + isinstance(self.max_tasks, bool) + or not isinstance(self.max_tasks, int) + or self.max_tasks < 1 + ): + raise ValueError("max_tasks must be positive when set") + if not isinstance(self.exit_when_idle, bool): + raise ValueError("exit_when_idle must be a boolean") if not self.capabilities: raise ValueError("capabilities cannot be empty") # Runner subprocesses use a task directory as their cwd. Keep the @@ -90,6 +101,7 @@ class WorkerConfig: cleanup = value("cleanup_after_seconds", "SCIMESH_CLEANUP_AFTER_SECONDS") cpu_count = value("cpu_count", "SCIMESH_CPU_COUNT", os.cpu_count() or 1) memory_mb = value("memory_mb", "SCIMESH_MEMORY_MB") + max_tasks = value("max_tasks", "SCIMESH_MAX_TASKS") return cls( coordinator_url=url.rstrip("/"), worker_id=value("worker_id", "SCIMESH_WORKER_ID"), @@ -102,4 +114,6 @@ class WorkerConfig: heartbeat_interval=float(value("heartbeat_interval", "SCIMESH_HEARTBEAT_INTERVAL", "15")), bearer_token=value("bearer_token", "SCIMESH_BEARER_TOKEN"), cleanup_after_seconds=float(cleanup) if cleanup else None, + max_tasks=int(max_tasks) if max_tasks is not None else None, + exit_when_idle=bool(values.get("exit_when_idle", False)), ) diff --git a/scimesh/worker/daemon.py b/scimesh/worker/daemon.py index dbf035b..50897bd 100644 --- a/scimesh/worker/daemon.py +++ b/scimesh/worker/daemon.py @@ -81,26 +81,43 @@ class WorkerDaemon: def run_forever(self) -> None: failures = 0 - while True: - try: - if not self._registered: - self._register_worker() - self._cleanup_expired_directories() - claimed = self.run_once() - failures = 0 - if not claimed: - self._sleep(self.config.poll_interval) - except CoordinatorTransientError as error: - failures += 1 - self._log("failed", error_type=type(error).__name__) - self._sleep(min(self.config.poll_interval * 2 ** min(failures, 6), 60.0)) + processed_tasks = 0 + self._log( + "started", + max_tasks=self.config.max_tasks, + exit_when_idle=self.config.exit_when_idle, + ) + try: + while True: + try: + if not self._registered: + self._register_worker() + self._cleanup_expired_directories() + claimed = 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 + elif self.config.exit_when_idle: + self._log("stopped", reason="queue_empty", processed_tasks=processed_tasks) + return + else: + self._sleep(self.config.poll_interval) + except CoordinatorTransientError as error: + failures += 1 + 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) def run_once(self) -> bool: worker_id = self._worker_id() - self._log("claiming") + self._log("claiming", log_level=logging.DEBUG) task = self.coordinator.claim(worker_id, self.config.capabilities) if task is None: - self._log("idle") + self._log("idle", log_level=logging.DEBUG) return False started = time.monotonic() task_dir = self.config.work_dir / task.task_id / str(task.attempt) @@ -135,7 +152,7 @@ class WorkerDaemon: }, }, ) - self._log("idle", task, elapsed_seconds=round(time.monotonic() - started, 3)) + self._log("completed", task, elapsed_seconds=round(time.monotonic() - started, 3)) except CoordinatorConflictError as error: self._log("lease_lost", task, error_type=type(error).__name__) except Exception as error: @@ -180,9 +197,16 @@ class WorkerDaemon: """Keep completion payload exact: coordinator owns all artifact metadata.""" return {"artifact_id": uploaded.artifact_id} - def _log(self, state: str, task: ClaimedTask | None = None, **extra: object) -> None: + def _log( + self, + state: str, + task: ClaimedTask | None = None, + *, + log_level: int = logging.INFO, + **extra: object, + ) -> None: fields = {"worker_id": self.config.worker_id, "task_id": task.task_id if task else None, "attempt": task.attempt if task else None, "state": state, **extra} - self.log.info("worker_event %s", fields) + self.log.log(log_level, "worker_event %s", fields) def _cleanup_expired_directories(self) -> None: """Remove only old task attempt directories when retention was configured.""" diff --git a/scripts/two-worker-smoke.sh b/scripts/two-worker-smoke.sh index 9cb000f..e551bc7 100755 --- a/scripts/two-worker-smoke.sh +++ b/scripts/two-worker-smoke.sh @@ -84,7 +84,7 @@ start_worker() { SCIMESH_BEARER_TOKEN="$TOKEN" \ SCIMESH_WORKER_NAME="$worker_name" \ SCIMESH_POLL_INTERVAL=0.2 \ - "$WORKER_PYTHON" -m scimesh.worker.cli --work-dir "$worker_dir" >"$worker_dir.log" 2>&1 & + "$WORKER_PYTHON" -m scimesh.worker.cli --work-dir "$worker_dir" --max-tasks 2 >"$worker_dir.log" 2>&1 & STARTED_WORKER_PID=$! } diff --git a/tests/test_worker_daemon.py b/tests/test_worker_daemon.py index 36f6cf6..c1bab39 100644 --- a/tests/test_worker_daemon.py +++ b/tests/test_worker_daemon.py @@ -1,6 +1,7 @@ from __future__ import annotations import hashlib +import logging from pathlib import Path import time from datetime import datetime, timedelta, timezone @@ -9,6 +10,7 @@ from urllib.request import Request import pytest from scimesh.worker.config import WorkerConfig +from scimesh.worker.cli import build_parser from scimesh.worker.coordinator import CoordinatorTransientError from scimesh.worker.daemon import LeaseHeartbeat, WorkerDaemon from scimesh.worker.models import ( @@ -109,6 +111,51 @@ def test_no_task_does_not_create_directory(tmp_path: Path) -> None: assert not config.work_dir.exists() +def test_once_worker_exits_after_an_empty_claim(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: + 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 runner.calls == 0 + assert "queue_empty" in caplog.text + + +def test_worker_stops_after_the_configured_number_of_claims(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: + caplog.set_level(logging.INFO, logger="scimesh.worker") + 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 runner.calls == 1 + assert "max_tasks_reached" in caplog.text + + +def test_keyboard_interrupt_stops_worker_without_propagating(tmp_path: Path, caplog: pytest.LogCaptureFixture) -> None: + caplog.set_level(logging.INFO, logger="scimesh.worker") + class InterruptingCoordinator(FakeCoordinator): + def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: + raise KeyboardInterrupt + + worker, _, _, _, _ = daemon(tmp_path, None, b"") + worker.coordinator = InterruptingCoordinator(None) + worker.run_forever() + assert "interrupted" in caplog.text + + +def test_worker_cli_lifecycle_options_are_explicit_and_exclusive() -> None: + parser = build_parser() + assert parser.parse_args(["--once"]).once is True + assert parser.parse_args(["--max-tasks", "2"]).max_tasks == 2 + with pytest.raises(SystemExit): + parser.parse_args(["--once", "--max-tasks", "2"]) + + +@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"): + WorkerConfig("https://example.test", None, tmp_path, max_tasks=value) # type: ignore[arg-type] + + 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