Improve worker lifecycle controls

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:

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_HEARTBEAT_INTERVAL, SCIMESH_CLEANUP_AFTER_SECONDS, "
-            "SCIMESH_BEARER_TOKEN. SCIMESH_WORKER_ID is a legacy/test override."
+            "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:

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

scimesh/worker/daemon.py

@@ -81,26 +81,43 @@ class WorkerDaemon:
 
     def run_forever(self) -> None:
         failures = 0
-        while True:
+        processed_tasks = 0
-            try:
+        self._log(
-                if not self._registered:
+            "started",
-                    self._register_worker()
+            max_tasks=self.config.max_tasks,
-                self._cleanup_expired_directories()
+            exit_when_idle=self.config.exit_when_idle,
-                claimed = self.run_once()
+        )
-                failures = 0
+        try:
-                if not claimed:
+            while True:
-                    self._sleep(self.config.poll_interval)
+                try:
-            except CoordinatorTransientError as error:
+                    if not self._registered:
-                failures += 1
+                        self._register_worker()
-                self._log("failed", error_type=type(error).__name__)
+                    self._cleanup_expired_directories()
-                self._sleep(min(self.config.poll_interval * 2 ** min(failures, 6), 60.0))
+                    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."""

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=$!
 }
 

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