Emil_Shanaty/SciMesh
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improve worker lifecycle controls

Browse Source
This commit is contained in:
Emil
2026-07-24 13:08:10 +03:00
parent 43ceec1f77
commit bde6cdb4ba
6 changed files with 158 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scimesh/worker/cli.py
+29 -4
View File
@@ -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:
scimesh/worker/config.py
+14
View File
@@ -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
+42 -18
View File
@@ -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."""