Исправлены замечания ревью Workers

scimesh/worker/daemon.py

@@ -6,7 +6,10 @@ import logging
 from pathlib import Path
 import random
 import shutil
+import threading
 import time
+from datetime import datetime, timezone
+from urllib.parse import quote
 
 from .artifacts import ArtifactClient, sha256_file
 from .config import WorkerConfig
@@ -15,6 +18,51 @@ 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
+
+    def start(self) -> None:
+        # Verify ownership before expensive download or calculation begins.
+        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.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 = self.config.heartbeat_interval
+
+    def _seconds_until_expiry(self) -> float:
+        try:
+            expiry = datetime.fromisoformat(self.task.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
@@ -43,7 +91,9 @@ class WorkerDaemon:
         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)
@@ -51,21 +101,30 @@ class WorkerDaemon:
                 raise ValueError("input checksum mismatch")
             self._log("running", task)
             result = self.runner.run(task, task_dir)
-            self._log("uploading", task)
+            heartbeat.raise_if_failed()
             manifests = [
-                {"uri": self.artifacts.upload(task, artifact), "sha256": sha256_file(artifact.path), "content_type": artifact.content_type}
+                self._manifest_uri(task, artifact.path.name, artifact.content_type, sha256_file(artifact.path))
                 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 _manifest_uri(self, task: ClaimedTask, filename: str, content_type: str, checksum: str) -> dict[str, str]:
+        """A stable logical URI until a shared artifact store is introduced."""
+        worker = quote(self.config.worker_id, safe="")
+        name = quote(filename, safe="")
+        return {"uri": f"worker://{worker}/{task.task_id}/{task.attempt}/{name}", "sha256": checksum, "content_type": content_type}
+
     def _report_failure(self, task: ClaimedTask, error: Exception) -> None:
         message = str(error).replace(str(self.config.work_dir), "<worker-dir>")[:300]
         try: