Fix worker result and lease contracts
This commit is contained in:
@@ -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
|
- HTTP server: standard `net/http` is sufficient; do not introduce a framework
|
||||||
unless it solves a concrete requirement.
|
unless it solves a concrete requirement.
|
||||||
- Schema migrations: versioned SQL migrations using `golang-migrate`.
|
- 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.
|
- Identifiers: UUID.
|
||||||
- Times: timezone-aware UTC timestamps.
|
- Times: timezone-aware UTC timestamps.
|
||||||
- API boundary: the Go coordinator owns all database access. Python workers
|
- 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 /jobs` | Validate submission, create job and pending tasks transactionally |
|
||||||
| `POST /tasks/claim` | Atomically lease one compatible task; `204` when none exists |
|
| `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}/result` | Idempotently persist a completed result manifest |
|
||||||
| `POST /tasks/{task_id}/failure` | Record a safe failure or retryable state |
|
| `POST /tasks/{task_id}/failure` | Record a safe failure or retryable state |
|
||||||
| `GET /jobs/{job_id}` | Return aggregate job/task progress |
|
| `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,
|
request ID in error responses and emit structured logs with the request ID,
|
||||||
task ID, worker ID, and operation.
|
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
|
## Tests and acceptance criteria
|
||||||
|
|
||||||
- `golang-migrate` upgrades an empty PostgreSQL database to the current schema
|
- `golang-migrate` upgrades an empty PostgreSQL database to the current schema
|
||||||
@@ -190,7 +200,7 @@ task ID, worker ID, and operation.
|
|||||||
|
|
||||||
## Out of scope
|
## 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;
|
- video chunk generation and trajectory stitching;
|
||||||
- authentication provider, UI, object-storage implementation, and PDF report;
|
- authentication provider, UI, object-storage implementation, and PDF report;
|
||||||
- network/distributed coordinator deployment beyond the local coordinator
|
- network/distributed coordinator deployment beyond the local coordinator
|
||||||
|
|||||||
@@ -12,7 +12,7 @@ The target flow is:
|
|||||||
Worker Daemon -> coordinator: claim task
|
Worker Daemon -> coordinator: claim task
|
||||||
coordinator -> Worker Daemon: task metadata + input location
|
coordinator -> Worker Daemon: task metadata + input location
|
||||||
Worker Daemon -> local SciMesh Core / CV runner: execute
|
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
|
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
|
The `result.uri` must be the durable URI returned by the artifact upload
|
||||||
`error_code` and `error_message`. Never send a Python traceback, access token,
|
endpoint below; a worker-local `file://` or `worker://` path is invalid.
|
||||||
or local path outside the worker directory.
|
|
||||||
|
### Upload a result artifact
|
||||||
|
|
||||||
|
```http
|
||||||
|
PUT /tasks/{task_id}/artifacts/{filename}
|
||||||
|
Content-Type: text/csv
|
||||||
|
X-Worker-ID: worker-01
|
||||||
|
X-Task-Attempt: 1
|
||||||
|
|
||||||
|
<CSV bytes>
|
||||||
|
```
|
||||||
|
|
||||||
|
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
|
## Required state machine
|
||||||
|
|
||||||
@@ -121,6 +143,8 @@ idle -> claiming -> downloading -> running -> uploading -> submitting -> idle
|
|||||||
- Create one isolated task directory: `<work-dir>/<task-id>/<attempt>/`.
|
- Create one isolated task directory: `<work-dir>/<task-id>/<attempt>/`.
|
||||||
- Invoke the runner with an explicit argument list, never `shell=True`.
|
- Invoke the runner with an explicit argument list, never `shell=True`.
|
||||||
- Upload/submit exactly the produced result files listed by the runner.
|
- 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
|
- A timeout, network error, or rejected submission must leave the local task
|
||||||
directory available for diagnostics until a configurable cleanup period.
|
directory available for diagnostics until a configurable cleanup period.
|
||||||
- Treat a duplicate successful submission as success when the coordinator
|
- Treat a duplicate successful submission as success when the coordinator
|
||||||
|
|||||||
@@ -3,11 +3,15 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
import hashlib
|
import hashlib
|
||||||
|
import http.client
|
||||||
|
import json
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Protocol
|
from typing import Protocol
|
||||||
from urllib.parse import urlsplit
|
from urllib.parse import quote, urlsplit
|
||||||
from urllib.request import HTTPRedirectHandler, Request, build_opener
|
from urllib.request import HTTPRedirectHandler, Request, build_opener
|
||||||
|
|
||||||
|
from .models import ClaimedTask, ProducedArtifact
|
||||||
|
|
||||||
|
|
||||||
def _origin(uri: str) -> tuple[str, str, int | None]:
|
def _origin(uri: str) -> tuple[str, str, int | None]:
|
||||||
parsed = urlsplit(uri)
|
parsed = urlsplit(uri)
|
||||||
@@ -32,14 +36,13 @@ class _SameOriginAuthRedirectHandler(HTTPRedirectHandler):
|
|||||||
class ArtifactClient(Protocol):
|
class ArtifactClient(Protocol):
|
||||||
def download(self, uri: str, destination: Path) -> None: ...
|
def download(self, uri: str, destination: Path) -> None: ...
|
||||||
|
|
||||||
|
def upload(
|
||||||
|
self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact
|
||||||
|
) -> str: ...
|
||||||
|
|
||||||
|
|
||||||
class HttpArtifactClient:
|
class HttpArtifactClient:
|
||||||
"""Downloads task inputs without exposing credentials to external storage.
|
"""Transfers artifacts through the coordinator without leaking credentials."""
|
||||||
|
|
||||||
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.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, coordinator_url: str, timeout: float, bearer_token: str | None = None) -> None:
|
def __init__(self, coordinator_url: str, timeout: float, bearer_token: str | None = None) -> None:
|
||||||
self.coordinator_url = coordinator_url.rstrip("/")
|
self.coordinator_url = coordinator_url.rstrip("/")
|
||||||
@@ -55,6 +58,48 @@ class HttpArtifactClient:
|
|||||||
while chunk := response.read(1024 * 1024):
|
while chunk := response.read(1024 * 1024):
|
||||||
target.write(chunk)
|
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]:
|
def _auth_headers_for(self, uri: str) -> dict[str, str]:
|
||||||
"""Only coordinator-owned URLs receive the coordinator bearer token."""
|
"""Only coordinator-owned URLs receive the coordinator bearer token."""
|
||||||
if self.bearer_token and _origin(uri) == self.coordinator_origin:
|
if self.bearer_token and _origin(uri) == self.coordinator_origin:
|
||||||
|
|||||||
@@ -31,3 +31,7 @@ def main(argv: list[str] | None = None) -> int:
|
|||||||
client = HttpCoordinatorClient(config.coordinator_url, config.request_timeout, config.bearer_token)
|
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()
|
WorkerDaemon(config, client, HttpArtifactClient(config.coordinator_url, config.request_timeout, config.bearer_token), SciMeshRunner()).run_forever()
|
||||||
return 0
|
return 0
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
raise SystemExit(main())
|
||||||
|
|||||||
@@ -19,6 +19,14 @@ class WorkerConfig:
|
|||||||
cleanup_after_seconds: float | None = None
|
cleanup_after_seconds: float | None = None
|
||||||
capabilities: tuple[str, ...] = ("similarity-search", "similarity-graph")
|
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
|
@classmethod
|
||||||
def from_environment(cls) -> "WorkerConfig":
|
def from_environment(cls) -> "WorkerConfig":
|
||||||
url = os.getenv("SCIMESH_COORDINATOR_URL")
|
url = os.getenv("SCIMESH_COORDINATOR_URL")
|
||||||
|
|||||||
@@ -23,7 +23,9 @@ class CoordinatorClient(Protocol):
|
|||||||
|
|
||||||
def submit(self, task: ClaimedTask, payload: dict[str, Any]) -> None: ...
|
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:
|
class HttpCoordinatorClient:
|
||||||
@@ -48,13 +50,22 @@ class HttpCoordinatorClient:
|
|||||||
if status not in (200, 201, 202):
|
if status not in (200, 201, 202):
|
||||||
raise CoordinatorError(f"result rejected with status {status}")
|
raise CoordinatorError(f"result rejected with status {status}")
|
||||||
|
|
||||||
def heartbeat(self, task: ClaimedTask, worker_id: str) -> None:
|
def fail(self, task: ClaimedTask, payload: dict[str, Any]) -> None:
|
||||||
status, _ = self._request(
|
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",
|
"POST", f"/tasks/{task.task_id}/heartbeat",
|
||||||
{"worker_id": worker_id, "attempt": task.attempt},
|
{"worker_id": worker_id, "attempt": task.attempt},
|
||||||
)
|
)
|
||||||
if status != 200:
|
if status != 200:
|
||||||
raise CoordinatorError(f"heartbeat rejected with status {status}")
|
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]]:
|
def _request(self, method: str, path: str, payload: dict[str, Any]) -> tuple[int, dict[str, Any]]:
|
||||||
request = Request(
|
request = Request(
|
||||||
|
|||||||
+17
-13
@@ -9,7 +9,6 @@ import shutil
|
|||||||
import threading
|
import threading
|
||||||
import time
|
import time
|
||||||
from datetime import datetime, timezone
|
from datetime import datetime, timezone
|
||||||
from urllib.parse import quote
|
|
||||||
|
|
||||||
from .artifacts import ArtifactClient, sha256_file
|
from .artifacts import ArtifactClient, sha256_file
|
||||||
from .config import WorkerConfig
|
from .config import WorkerConfig
|
||||||
@@ -26,10 +25,13 @@ class LeaseHeartbeat:
|
|||||||
self._stop = threading.Event()
|
self._stop = threading.Event()
|
||||||
self._error: Exception | None = None
|
self._error: Exception | None = None
|
||||||
self._thread: threading.Thread | None = None
|
self._thread: threading.Thread | None = None
|
||||||
|
self._lease_expires_at = task.lease_expires_at
|
||||||
|
|
||||||
def start(self) -> None:
|
def start(self) -> None:
|
||||||
# Verify ownership before expensive download or calculation begins.
|
# 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 = threading.Thread(target=self._run, name=f"lease-{self.task.task_id}", daemon=True)
|
||||||
self._thread.start()
|
self._thread.start()
|
||||||
|
|
||||||
@@ -46,15 +48,19 @@ class LeaseHeartbeat:
|
|||||||
delay = min(self.config.heartbeat_interval, self._seconds_until_expiry() / 2)
|
delay = min(self.config.heartbeat_interval, self._seconds_until_expiry() / 2)
|
||||||
while not self._stop.wait(max(delay, 0.01)):
|
while not self._stop.wait(max(delay, 0.01)):
|
||||||
try:
|
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.
|
except Exception as error: # Surface the lease loss in the main state machine.
|
||||||
self._error = error
|
self._error = error
|
||||||
return
|
return
|
||||||
delay = self.config.heartbeat_interval
|
delay = min(
|
||||||
|
self.config.heartbeat_interval, self._seconds_until_expiry() / 2
|
||||||
|
)
|
||||||
|
|
||||||
def _seconds_until_expiry(self) -> float:
|
def _seconds_until_expiry(self) -> float:
|
||||||
try:
|
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:
|
except ValueError as error:
|
||||||
raise ValueError("invalid lease_expires_at") from error
|
raise ValueError("invalid lease_expires_at") from error
|
||||||
seconds = (expiry - datetime.now(timezone.utc)).total_seconds()
|
seconds = (expiry - datetime.now(timezone.utc)).total_seconds()
|
||||||
@@ -103,7 +109,11 @@ class WorkerDaemon:
|
|||||||
result = self.runner.run(task, task_dir)
|
result = self.runner.run(task, task_dir)
|
||||||
heartbeat.raise_if_failed()
|
heartbeat.raise_if_failed()
|
||||||
manifests = [
|
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
|
for artifact in result.artifacts
|
||||||
]
|
]
|
||||||
if not manifests:
|
if not manifests:
|
||||||
@@ -119,16 +129,10 @@ class WorkerDaemon:
|
|||||||
heartbeat.stop()
|
heartbeat.stop()
|
||||||
return True
|
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:
|
def _report_failure(self, task: ClaimedTask, error: Exception) -> None:
|
||||||
message = str(error).replace(str(self.config.work_dir), "<worker-dir>")[:300]
|
message = str(error).replace(str(self.config.work_dir), "<worker-dir>")[:300]
|
||||||
try:
|
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:
|
except CoordinatorTransientError:
|
||||||
raise
|
raise
|
||||||
except Exception:
|
except Exception:
|
||||||
|
|||||||
@@ -3,13 +3,14 @@ from __future__ import annotations
|
|||||||
import hashlib
|
import hashlib
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
import time
|
import time
|
||||||
|
from datetime import datetime, timedelta, timezone
|
||||||
from urllib.request import Request
|
from urllib.request import Request
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from scimesh.worker.config import WorkerConfig
|
from scimesh.worker.config import WorkerConfig
|
||||||
from scimesh.worker.coordinator import CoordinatorTransientError
|
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.models import ClaimedTask, InputArtifact, ProducedArtifact, RunResult
|
||||||
from scimesh.worker.artifacts import HttpArtifactClient, _SameOriginAuthRedirectHandler, _origin
|
from scimesh.worker.artifacts import HttpArtifactClient, _SameOriginAuthRedirectHandler, _origin
|
||||||
from scimesh.worker.runners import SciMeshRunner
|
from scimesh.worker.runners import SciMeshRunner
|
||||||
@@ -17,7 +18,7 @@ from scimesh.worker.runners import SciMeshRunner
|
|||||||
|
|
||||||
class FakeCoordinator:
|
class FakeCoordinator:
|
||||||
def __init__(self, task: ClaimedTask | None) -> None:
|
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:
|
def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None:
|
||||||
task, self.task = self.task, None
|
task, self.task = self.task, None
|
||||||
@@ -26,17 +27,25 @@ class FakeCoordinator:
|
|||||||
def submit(self, task: ClaimedTask, payload: dict) -> None:
|
def submit(self, task: ClaimedTask, payload: dict) -> None:
|
||||||
self.submissions.append(payload)
|
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))
|
self.heartbeats.append((task.task_id, task.attempt, worker_id))
|
||||||
|
return (datetime.now(timezone.utc) + timedelta(seconds=1)).isoformat()
|
||||||
|
|
||||||
|
|
||||||
class FakeArtifacts:
|
class FakeArtifacts:
|
||||||
def __init__(self, content: bytes) -> None:
|
def __init__(self, content: bytes) -> None:
|
||||||
self.content = content
|
self.content, self.uploaded = content, []
|
||||||
|
|
||||||
def download(self, uri: str, destination: Path) -> None:
|
def download(self, uri: str, destination: Path) -> None:
|
||||||
destination.write_bytes(self.content)
|
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:
|
class FakeRunner:
|
||||||
def __init__(self) -> None:
|
def __init__(self) -> None:
|
||||||
self.calls = 0
|
self.calls = 0
|
||||||
@@ -49,7 +58,8 @@ class FakeRunner:
|
|||||||
|
|
||||||
|
|
||||||
def make_task(content: bytes, checksum: str | None = None) -> ClaimedTask:
|
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):
|
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)
|
worker, coordinator, artifacts, runner, _ = daemon(tmp_path, make_task(content), content)
|
||||||
assert worker.run_once() is True
|
assert worker.run_once() is True
|
||||||
assert runner.calls == 1
|
assert runner.calls == 1
|
||||||
|
assert len(artifacts.uploaded) == 1
|
||||||
assert coordinator.heartbeats == [("task-1", 1, "worker-1")]
|
assert coordinator.heartbeats == [("task-1", 1, "worker-1")]
|
||||||
assert coordinator.submissions[0]["status"] == "completed"
|
assert coordinator.submissions[0]["status"] == "completed"
|
||||||
assert coordinator.submissions[0]["result"]["content_type"] == "text/csv"
|
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:
|
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")
|
worker, coordinator, _, runner, _ = daemon(tmp_path, make_task(b"actual", "not-the-hash"), b"actual")
|
||||||
assert worker.run_once() is True
|
assert worker.run_once() is True
|
||||||
assert runner.calls == 0
|
assert runner.calls == 0
|
||||||
assert coordinator.submissions[0]["status"] == "failed"
|
assert coordinator.failures[0]["error_code"] == "ValueError"
|
||||||
assert coordinator.submissions[0]["error_code"] == "ValueError"
|
assert not coordinator.submissions
|
||||||
|
|
||||||
|
|
||||||
def test_transient_claim_error_is_propagated_for_bounded_backoff(tmp_path: Path) -> None:
|
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
|
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:
|
def test_runner_maps_graph_and_smiles_search_parameters(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||||
commands: list[list[str]] = []
|
commands: list[list[str]] = []
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user