SciMesh/scimesh/worker/daemon.py

"""The worker state machine and its safe failure handling."""

from __future__ import annotations

import logging
from pathlib import Path
import random
import shutil
import threading
import time
from datetime import datetime, timezone

from .artifacts import ArtifactClient, sha256_file
from .config import WorkerConfig
from .coordinator import CoordinatorClient, CoordinatorTransientError
from .models import ClaimedTask
from .runners import Runner


class LeaseHeartbeat:
    """Renews a claimed task lease while local work is in progress."""

    def __init__(self, task: ClaimedTask, coordinator: CoordinatorClient, config: WorkerConfig) -> None:
        self.task, self.coordinator, self.config = task, coordinator, config
        self._stop = threading.Event()
        self._error: Exception | None = None
        self._thread: threading.Thread | None = None
        self._lease_expires_at = task.lease_expires_at

    def start(self) -> None:
        # Verify ownership before expensive download or calculation begins.
        self._lease_expires_at = self.coordinator.heartbeat(
            self.task, self.config.worker_id
        )
        self._thread = threading.Thread(target=self._run, name=f"lease-{self.task.task_id}", daemon=True)
        self._thread.start()

    def stop(self) -> None:
        self._stop.set()
        if self._thread:
            self._thread.join()

    def raise_if_failed(self) -> None:
        if self._error:
            raise self._error

    def _run(self) -> None:
        delay = min(self.config.heartbeat_interval, self._seconds_until_expiry() / 2)
        while not self._stop.wait(max(delay, 0.01)):
            try:
                self._lease_expires_at = self.coordinator.heartbeat(
                    self.task, self.config.worker_id
                )
            except Exception as error:  # Surface the lease loss in the main state machine.
                self._error = error
                return
            delay = min(
                self.config.heartbeat_interval, self._seconds_until_expiry() / 2
            )

    def _seconds_until_expiry(self) -> float:
        try:
            expiry = datetime.fromisoformat(self._lease_expires_at.replace("Z", "+00:00"))
        except ValueError as error:
            raise ValueError("invalid lease_expires_at") from error
        seconds = (expiry - datetime.now(timezone.utc)).total_seconds()
        if seconds <= 0:
            raise ValueError("claimed task lease has already expired")
        return seconds


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
        self.log = logging.getLogger("scimesh.worker")

    def run_forever(self) -> None:
        failures = 0
        while True:
            try:
                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))

    def run_once(self) -> bool:
        self._log("claiming")
        task = self.coordinator.claim(self.config.worker_id, self.config.capabilities)
        if task is None:
            self._log("idle")
            return False
        started = time.monotonic()
        task_dir = self.config.work_dir / task.task_id / str(task.attempt)
        task_dir.mkdir(parents=True, exist_ok=False)
        heartbeat = LeaseHeartbeat(task, self.coordinator, self.config)
        try:
            heartbeat.start()
            self._log("downloading", task)
            input_path = task_dir / "input"
            self.artifacts.download(task.input.uri, input_path)
            if sha256_file(input_path).lower() != task.input.sha256.lower():
                raise ValueError("input checksum mismatch")
            self._log("running", task)
            result = self.runner.run(task, task_dir)
            heartbeat.raise_if_failed()
            manifests = [
                {
                    "uri": self.artifacts.upload(task, self.config.worker_id, artifact),
                    "sha256": sha256_file(artifact.path),
                    "content_type": artifact.content_type,
                }
                for artifact in result.artifacts
            ]
            if not manifests:
                raise ValueError("runner produced no artifacts")
            self._log("submitting", task)
            heartbeat.raise_if_failed()
            self.coordinator.submit(task, {"worker_id": self.config.worker_id, "attempt": task.attempt, "status": "completed", "result": manifests[0], "artifacts": manifests, "metrics": {**result.metrics, "elapsed_seconds": round(time.monotonic() - started, 3)}})
            self._log("idle", task, elapsed_seconds=round(time.monotonic() - started, 3))
        except Exception as error:
            self._log("failed", task, error_type=type(error).__name__)
            self._report_failure(task, error)
        finally:
            heartbeat.stop()
        return True

    def _report_failure(self, task: ClaimedTask, error: Exception) -> None:
        message = str(error).replace(str(self.config.work_dir), "<worker-dir>")[:300]
        try:
            self.coordinator.fail(task, {"worker_id": self.config.worker_id, "attempt": task.attempt, "error_code": type(error).__name__, "error_message": message})
        except CoordinatorTransientError:
            raise
        except Exception:
            self._log("failed", task, error_type="FailureReportError")

    def _log(self, state: str, task: ClaimedTask | None = None, **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)

    def _cleanup_expired_directories(self) -> None:
        """Remove only old task attempt directories when retention was configured."""
        if self.config.cleanup_after_seconds is None or not self.config.work_dir.exists():
            return
        cutoff = time.time() - self.config.cleanup_after_seconds
        for task_dir in self.config.work_dir.iterdir():
            if not task_dir.is_dir():
                continue
            for attempt_dir in task_dir.iterdir():
                if attempt_dir.is_dir() and attempt_dir.stat().st_mtime < cutoff:
                    shutil.rmtree(attempt_dir)
            if not any(task_dir.iterdir()):
                task_dir.rmdir()

    @staticmethod
    def _sleep(delay: float) -> None:
        time.sleep(delay * random.uniform(0.75, 1.25))