diff --git a/docs/database-integration-task.md b/docs/database-integration-task.md index fc85aa7..e7b2b2e 100644 --- a/docs/database-integration-task.md +++ b/docs/database-integration-task.md @@ -17,6 +17,8 @@ sketch. It does not implement the Worker Daemon or scientific/CV computation. - HTTP server: standard `net/http` is sufficient; do not introduce a framework unless it solves a concrete requirement. - Schema migrations: versioned SQL migrations using `golang-migrate`. +- Result storage: a coordinator-managed local artifact directory in the first + version; keep its interface replaceable by object storage later. - Identifiers: UUID. - Times: timezone-aware UTC timestamps. - API boundary: the Go coordinator owns all database access. Python workers @@ -155,7 +157,8 @@ must be `POST` rather than `GET`. | --- | --- | | `POST /jobs` | Validate submission, create job and pending tasks transactionally | | `POST /tasks/claim` | Atomically lease one compatible task; `204` when none exists | -| `POST /tasks/{task_id}/heartbeat` | Renew the current worker's lease | +| `POST /tasks/{task_id}/heartbeat` | Renew the current worker's lease; return the new `lease_expires_at` | +| `PUT /tasks/{task_id}/artifacts/{filename}` | Stream one result artifact into coordinator storage; verify the current worker lease | | `POST /tasks/{task_id}/result` | Idempotently persist a completed result manifest | | `POST /tasks/{task_id}/failure` | Record a safe failure or retryable state | | `GET /jobs/{job_id}` | Return aggregate job/task progress | @@ -170,6 +173,13 @@ inputs in the handler. Do not return raw database errors to clients. Return a request ID in error responses and emit structured logs with the request ID, task ID, worker ID, and operation. +The artifact endpoint receives binary content with `Content-Type`, +`X-Worker-ID`, and `X-Task-Attempt` headers. It must stream the request body to +the coordinator artifact directory instead of buffering it in memory, calculate +or verify its SHA-256, and return `201` with the durable artifact `uri`. Only +the worker holding the current lease may upload. `POST /result` accepts only +URIs returned by this endpoint for the same task and attempt. + ## Tests and acceptance criteria - `golang-migrate` upgrades an empty PostgreSQL database to the current schema @@ -190,7 +200,7 @@ task ID, worker ID, and operation. ## Out of scope -- Python Worker Daemon execution, input download, and result upload; +- Python Worker Daemon execution and input download; - video chunk generation and trajectory stitching; - authentication provider, UI, object-storage implementation, and PDF report; - network/distributed coordinator deployment beyond the local coordinator diff --git a/docs/worker-daemon-task.md b/docs/worker-daemon-task.md index e0900ce..82c7538 100644 --- a/docs/worker-daemon-task.md +++ b/docs/worker-daemon-task.md @@ -12,7 +12,7 @@ The target flow is: Worker Daemon -> coordinator: claim task coordinator -> Worker Daemon: task metadata + input location Worker Daemon -> local SciMesh Core / CV runner: execute -Worker Daemon -> coordinator: submit result +Worker Daemon -> coordinator: upload result artifact, then submit result manifest ``` The architecture sketch uses video chunks and CV, while the current SciMesh @@ -103,9 +103,31 @@ Content-Type: application/json } ``` -For a failed execution, send `status: "failed"` with a short, sanitized -`error_code` and `error_message`. Never send a Python traceback, access token, -or local path outside the worker directory. +The `result.uri` must be the durable URI returned by the artifact upload +endpoint below; a worker-local `file://` or `worker://` path is invalid. + +### Upload a result artifact + +```http +PUT /tasks/{task_id}/artifacts/{filename} +Content-Type: text/csv +X-Worker-ID: worker-01 +X-Task-Attempt: 1 + + +``` + +The coordinator streams the artifact to its configured storage and responds: + +```json +{ + "uri": "https://coordinator.example/tasks/0d2d/artifacts/result.csv" +} +``` + +For a failed execution, send a short, sanitized `error_code` and +`error_message` to `POST /tasks/{task_id}/failure`. Never send a Python +traceback, access token, or local path outside the worker directory. ## Required state machine @@ -121,6 +143,8 @@ idle -> claiming -> downloading -> running -> uploading -> submitting -> idle - Create one isolated task directory: `///`. - Invoke the runner with an explicit argument list, never `shell=True`. - Upload/submit exactly the produced result files listed by the runner. +- Do not mark a task completed until every submitted artifact has a durable + coordinator-provided URI. - A timeout, network error, or rejected submission must leave the local task directory available for diagnostics until a configurable cleanup period. - Treat a duplicate successful submission as success when the coordinator diff --git a/scimesh/worker/artifacts.py b/scimesh/worker/artifacts.py index 4497341..b075c42 100644 --- a/scimesh/worker/artifacts.py +++ b/scimesh/worker/artifacts.py @@ -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: diff --git a/scimesh/worker/cli.py b/scimesh/worker/cli.py index e0d5414..7d3bf1f 100644 --- a/scimesh/worker/cli.py +++ b/scimesh/worker/cli.py @@ -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()) diff --git a/scimesh/worker/config.py b/scimesh/worker/config.py index cd38cb0..0b5ca91 100644 --- a/scimesh/worker/config.py +++ b/scimesh/worker/config.py @@ -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") diff --git a/scimesh/worker/coordinator.py b/scimesh/worker/coordinator.py index 4061e9a..5fccd33 100644 --- a/scimesh/worker/coordinator.py +++ b/scimesh/worker/coordinator.py @@ -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( diff --git a/scimesh/worker/daemon.py b/scimesh/worker/daemon.py index 4480918..3200f6b 100644 --- a/scimesh/worker/daemon.py +++ b/scimesh/worker/daemon.py @@ -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), "")[: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: diff --git a/tests/test_worker_daemon.py b/tests/test_worker_daemon.py index f88e3d8..583ee93 100644 --- a/tests/test_worker_daemon.py +++ b/tests/test_worker_daemon.py @@ -3,13 +3,14 @@ from __future__ import annotations import hashlib from pathlib import Path import time +from datetime import datetime, timedelta, timezone from urllib.request import Request import pytest from scimesh.worker.config import WorkerConfig from scimesh.worker.coordinator import CoordinatorTransientError -from scimesh.worker.daemon import WorkerDaemon +from scimesh.worker.daemon import LeaseHeartbeat, WorkerDaemon from scimesh.worker.models import ClaimedTask, InputArtifact, ProducedArtifact, RunResult from scimesh.worker.artifacts import HttpArtifactClient, _SameOriginAuthRedirectHandler, _origin from scimesh.worker.runners import SciMeshRunner @@ -17,7 +18,7 @@ from scimesh.worker.runners import SciMeshRunner class FakeCoordinator: def __init__(self, task: ClaimedTask | None) -> None: - self.task, self.submissions, self.heartbeats = task, [], [] + self.task, self.submissions, self.failures, self.heartbeats = task, [], [], [] def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: task, self.task = self.task, None @@ -26,17 +27,25 @@ class FakeCoordinator: def submit(self, task: ClaimedTask, payload: dict) -> None: self.submissions.append(payload) - def heartbeat(self, task: ClaimedTask, worker_id: str) -> None: + def fail(self, task: ClaimedTask, payload: dict) -> None: + self.failures.append(payload) + + def heartbeat(self, task: ClaimedTask, worker_id: str) -> str: self.heartbeats.append((task.task_id, task.attempt, worker_id)) + return (datetime.now(timezone.utc) + timedelta(seconds=1)).isoformat() class FakeArtifacts: def __init__(self, content: bytes) -> None: - self.content = content + self.content, self.uploaded = content, [] def download(self, uri: str, destination: Path) -> None: destination.write_bytes(self.content) + def upload(self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact) -> str: + self.uploaded.append((task.task_id, worker_id, artifact.path)) + return f"https://example.test/tasks/{task.task_id}/artifacts/{artifact.path.name}" + class FakeRunner: def __init__(self) -> None: self.calls = 0 @@ -49,7 +58,8 @@ class FakeRunner: def make_task(content: bytes, checksum: str | None = None) -> ClaimedTask: - return ClaimedTask("task-1", 1, "2026-07-30T00:00:00Z", "similarity-search", InputArtifact("https://example.test/input", checksum or hashlib.sha256(content).hexdigest()), {"query_id": "CHEMBL1"}) + lease = (datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat() + return ClaimedTask("task-1", 1, lease, "similarity-search", InputArtifact("https://example.test/input", checksum or hashlib.sha256(content).hexdigest()), {"query_id": "CHEMBL1"}) def daemon(tmp_path: Path, task: ClaimedTask | None, content: bytes): @@ -63,10 +73,11 @@ def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None: worker, coordinator, artifacts, runner, _ = daemon(tmp_path, make_task(content), content) assert worker.run_once() is True assert runner.calls == 1 + assert len(artifacts.uploaded) == 1 assert coordinator.heartbeats == [("task-1", 1, "worker-1")] assert coordinator.submissions[0]["status"] == "completed" assert coordinator.submissions[0]["result"]["content_type"] == "text/csv" - assert coordinator.submissions[0]["result"]["uri"].startswith("worker://worker-1/") + assert coordinator.submissions[0]["result"]["uri"].startswith("https://example.test/tasks/task-1/artifacts/") def test_no_task_does_not_create_directory(tmp_path: Path) -> None: @@ -80,8 +91,8 @@ def test_bad_checksum_reports_failure_without_running(tmp_path: Path) -> None: worker, coordinator, _, runner, _ = daemon(tmp_path, make_task(b"actual", "not-the-hash"), b"actual") assert worker.run_once() is True assert runner.calls == 0 - assert coordinator.submissions[0]["status"] == "failed" - assert coordinator.submissions[0]["error_code"] == "ValueError" + assert coordinator.failures[0]["error_code"] == "ValueError" + assert not coordinator.submissions def test_transient_claim_error_is_propagated_for_bounded_backoff(tmp_path: Path) -> None: @@ -137,6 +148,23 @@ def test_lease_is_renewed_while_a_runner_is_still_working(tmp_path: Path) -> Non assert len(coordinator.heartbeats) >= 2 +def test_heartbeat_reschedules_from_the_renewed_lease(tmp_path: Path) -> None: + class ShortLeaseCoordinator(FakeCoordinator): + def heartbeat(self, task: ClaimedTask, worker_id: str) -> str: + self.heartbeats.append((task.task_id, task.attempt, worker_id)) + return (datetime.now(timezone.utc) + timedelta(seconds=0.02)).isoformat() + + config = WorkerConfig( + "https://example.test", "worker-1", tmp_path / "work", heartbeat_interval=1 + ) + coordinator = ShortLeaseCoordinator(None) + heartbeat = LeaseHeartbeat(make_task(b"fixture"), coordinator, config) + heartbeat.start() + time.sleep(0.06) + heartbeat.stop() + assert len(coordinator.heartbeats) >= 3 + + def test_runner_maps_graph_and_smiles_search_parameters(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: commands: list[list[str]] = []