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

Fix worker result and lease contracts

Browse Source
This commit is contained in:
Emil
2026-07-22 18:48:52 +03:00
parent 167b9e2aaa
commit bfdd234c2c
8 changed files with 171 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scimesh/worker/artifacts.py
+52 -7
View File
@@ -3,11 +3,15 @@
from __future__ import annotations
import hashlib
import http.client
import json
from pathlib import Path
from typing import Protocol
from urllib.parse import urlsplit
from urllib.parse import quote, urlsplit
from urllib.request import HTTPRedirectHandler, Request, build_opener
from .models import ClaimedTask, ProducedArtifact
def _origin(uri: str) -> tuple[str, str, int | None]:
parsed = urlsplit(uri)
@@ -32,14 +36,13 @@ class _SameOriginAuthRedirectHandler(HTTPRedirectHandler):
class ArtifactClient(Protocol):
def download(self, uri: str, destination: Path) -> None: ...
def upload(
self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact
) -> str: ...
class HttpArtifactClient:
"""Downloads task inputs without exposing credentials to external storage.
The present coordinator contract persists a result *manifest* at ``/result``
and deliberately defines no artifact-upload endpoint. Output storage can be
added later as a separate ArtifactClient implementation.
"""
"""Transfers artifacts through the coordinator without leaking credentials."""
def __init__(self, coordinator_url: str, timeout: float, bearer_token: str | None = None) -> None:
self.coordinator_url = coordinator_url.rstrip("/")
@@ -55,6 +58,48 @@ class HttpArtifactClient:
while chunk := response.read(1024 * 1024):
target.write(chunk)
def upload(self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact) -> str:
"""Stream one result artifact to the coordinator and return its stable URI."""
url = (
f"{self.coordinator_url}/tasks/{quote(task.task_id, safe='')}/artifacts/"
f"{quote(artifact.path.name, safe='')}"
)
parsed = urlsplit(url)
if parsed.scheme not in {"http", "https"} or not parsed.hostname:
raise ValueError("coordinator URL must be an absolute HTTP(S) URL")
connection_class = (
http.client.HTTPSConnection if parsed.scheme == "https" else http.client.HTTPConnection
)
connection = connection_class(parsed.hostname, parsed.port, timeout=self.timeout)
try:
path = parsed.path + (f"?{parsed.query}" if parsed.query else "")
connection.putrequest("PUT", path)
connection.putheader("Content-Type", artifact.content_type)
connection.putheader("Content-Length", str(artifact.path.stat().st_size))
connection.putheader("X-Worker-ID", worker_id)
connection.putheader("X-Task-Attempt", str(task.attempt))
for name, value in self._auth_headers_for(url).items():
connection.putheader(name, value)
connection.endheaders()
with artifact.path.open("rb") as source:
while chunk := source.read(1024 * 1024):
connection.send(chunk)
response = connection.getresponse()
body = response.read()
if not 200 <= response.status < 300:
raise RuntimeError(f"artifact upload rejected with status {response.status}")
if body:
try:
response_data = json.loads(body)
except json.JSONDecodeError as error:
raise RuntimeError("artifact upload returned invalid JSON") from error
response_uri = response_data.get("uri") if isinstance(response_data, dict) else None
if isinstance(response_uri, str) and response_uri:
return response_uri
return url
finally:
connection.close()
def _auth_headers_for(self, uri: str) -> dict[str, str]:
"""Only coordinator-owned URLs receive the coordinator bearer token."""
if self.bearer_token and _origin(uri) == self.coordinator_origin:
scimesh/worker/cli.py
+4
View File
@@ -31,3 +31,7 @@ def main(argv: list[str] | None = None) -> int:
client = HttpCoordinatorClient(config.coordinator_url, config.request_timeout, config.bearer_token)
WorkerDaemon(config, client, HttpArtifactClient(config.coordinator_url, config.request_timeout, config.bearer_token), SciMeshRunner()).run_forever()
return 0
if __name__ == "__main__":
raise SystemExit(main())
scimesh/worker/config.py
+8
View File
@@ -19,6 +19,14 @@ class WorkerConfig:
cleanup_after_seconds: float | None = None
capabilities: tuple[str, ...] = ("similarity-search", "similarity-graph")
def __post_init__(self) -> None:
if self.poll_interval <= 0:
raise ValueError("poll_interval must be positive")
if self.request_timeout <= 0:
raise ValueError("request_timeout must be positive")
if self.heartbeat_interval <= 0:
raise ValueError("heartbeat_interval must be positive")
@classmethod
def from_environment(cls) -> "WorkerConfig":
url = os.getenv("SCIMESH_COORDINATOR_URL")
scimesh/worker/coordinator.py
+14 -3
View File
@@ -23,7 +23,9 @@ class CoordinatorClient(Protocol):
def submit(self, task: ClaimedTask, payload: dict[str, Any]) -> None: ...
def heartbeat(self, task: ClaimedTask, worker_id: str) -> None: ...
def fail(self, task: ClaimedTask, payload: dict[str, Any]) -> None: ...
def heartbeat(self, task: ClaimedTask, worker_id: str) -> str: ...
class HttpCoordinatorClient:
@@ -48,13 +50,22 @@ class HttpCoordinatorClient:
if status not in (200, 201, 202):
raise CoordinatorError(f"result rejected with status {status}")
def heartbeat(self, task: ClaimedTask, worker_id: str) -> None:
status, _ = self._request(
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):
raise CoordinatorError(f"failure report rejected with status {status}")
def heartbeat(self, task: ClaimedTask, worker_id: str) -> str:
status, body = self._request(
"POST", f"/tasks/{task.task_id}/heartbeat",
{"worker_id": worker_id, "attempt": task.attempt},
)
if status != 200:
raise CoordinatorError(f"heartbeat rejected with status {status}")
lease_expires_at = body.get("lease_expires_at")
if not isinstance(lease_expires_at, str):
raise CoordinatorError("heartbeat response is missing lease_expires_at")
return lease_expires_at
def _request(self, method: str, path: str, payload: dict[str, Any]) -> tuple[int, dict[str, Any]]:
request = Request(
scimesh/worker/daemon.py
+17 -13
View File
@@ -9,7 +9,6 @@ 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
@@ -26,10 +25,13 @@ class LeaseHeartbeat:
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.coordinator.heartbeat(self.task, self.config.worker_id)
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()
@@ -46,15 +48,19 @@ class LeaseHeartbeat:
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)
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 = self.config.heartbeat_interval
delay = min(
self.config.heartbeat_interval, self._seconds_until_expiry() / 2
)
def _seconds_until_expiry(self) -> float:
try:
expiry = datetime.fromisoformat(self.task.lease_expires_at.replace("Z", "+00:00"))
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()
@@ -103,7 +109,11 @@ class WorkerDaemon:
result = self.runner.run(task, task_dir)
heartbeat.raise_if_failed()
manifests = [
self._manifest_uri(task, artifact.path.name, artifact.content_type, sha256_file(artifact.path))
{
"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:
@@ -119,16 +129,10 @@ class WorkerDaemon:
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:
self.coordinator.submit(task, {"worker_id": self.config.worker_id, "attempt": task.attempt, "status": "failed", "error_code": type(error).__name__, "error_message": message})
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: