Fix coordinator worker integration

scimesh/worker/coordinator.py

@@ -5,9 +5,10 @@ from __future__ import annotations
 import json
 from typing import Any, Protocol
 from urllib.error import HTTPError, URLError
-from urllib.request import Request, urlopen
+from urllib.request import Request, build_opener
 
-from .models import ClaimedTask
+from .models import ClaimedTask, RegisteredWorker
+from .transport import NoRedirectHandler
 
 
 class CoordinatorError(RuntimeError):
@@ -18,7 +19,15 @@ class CoordinatorTransientError(CoordinatorError):
     """A timeout, connection error, or 5xx coordinator response."""
 
 
+class CoordinatorConflictError(CoordinatorError):
+    """The worker no longer owns the task lease or attempted a conflicting mutation."""
+
+
 class CoordinatorClient(Protocol):
+    def register(
+        self, name: str, capabilities: tuple[str, ...], cpu_count: int, memory_mb: int | None
+    ) -> RegisteredWorker: ...
+
     def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: ...
 
     def submit(self, task: ClaimedTask, payload: dict[str, Any]) -> None: ...
@@ -33,6 +42,25 @@ class HttpCoordinatorClient:
         self.base_url = base_url.rstrip("/")
         self.timeout = timeout
         self.bearer_token = bearer_token
+        self._opener = build_opener(NoRedirectHandler())
+
+    def register(
+        self, name: str, capabilities: tuple[str, ...], cpu_count: int, memory_mb: int | None
+    ) -> RegisteredWorker:
+        payload: dict[str, Any] = {
+            "name": name,
+            "capabilities": list(capabilities),
+            "cpu_count": cpu_count,
+        }
+        if memory_mb is not None:
+            payload["memory_mb"] = memory_mb
+        status, body = self._request("POST", "/workers/register", payload)
+        if status != 201:
+            raise CoordinatorError(f"worker registration rejected with status {status}")
+        try:
+            return RegisteredWorker.from_json(body)
+        except ValueError as error:
+            raise CoordinatorError("invalid worker registration response") from error
 
     def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None:
         status, body = self._request("POST", "/tasks/claim", {
@@ -48,11 +76,15 @@ class HttpCoordinatorClient:
         status, _ = self._request("POST", f"/tasks/{task.task_id}/result", payload)
         # 200/201/202 include a successful or idempotent duplicate result response.
         if status not in (200, 201, 202):
+            if status == 409:
+                raise CoordinatorConflictError("result rejected because the task lease was lost")
             raise CoordinatorError(f"result rejected with status {status}")
 
     def fail(self, task: ClaimedTask, payload: dict[str, Any]) -> None:
         status, _ = self._request("POST", f"/tasks/{task.task_id}/failure", payload)
         if status not in (200, 201, 202):
+            if status == 409:
+                raise CoordinatorConflictError("failure rejected because the task lease was lost")
             raise CoordinatorError(f"failure report rejected with status {status}")
 
     def heartbeat(self, task: ClaimedTask, worker_id: str) -> str:
@@ -61,6 +93,8 @@ class HttpCoordinatorClient:
             {"worker_id": worker_id, "attempt": task.attempt},
         )
         if status != 200:
+            if status == 409:
+                raise CoordinatorConflictError("heartbeat rejected because the task lease was lost")
             raise CoordinatorError(f"heartbeat rejected with status {status}")
         lease_expires_at = body.get("lease_expires_at")
         if not isinstance(lease_expires_at, str):
@@ -73,9 +107,12 @@ class HttpCoordinatorClient:
             headers={"Content-Type": "application/json", **self._auth_header()},
         )
         try:
-            with urlopen(request, timeout=self.timeout) as response:
+            with self._opener.open(request, timeout=self.timeout) as response:
                 raw = response.read()
-                return response.status, json.loads(raw) if raw else {}
+                try:
+                    return response.status, json.loads(raw) if raw else {}
+                except json.JSONDecodeError as error:
+                    raise CoordinatorError("coordinator returned invalid JSON") from error
         except HTTPError as error:
             if error.code >= 500:
                 raise CoordinatorTransientError(f"coordinator returned {error.code}") from error