From 3e34dbc1e2885616a39a993534298de0e0a0cf53 Mon Sep 17 00:00:00 2001 From: reran4ik Date: Wed, 22 Jul 2026 18:14:50 +0300 Subject: [PATCH] =?UTF-8?q?=D0=94=D0=BE=D0=B1=D0=B0=D0=B2=D0=BB=D0=B5?= =?UTF-8?q?=D0=BD=D1=8B=20Workers=20=D0=B2=20=D1=81=D0=B8=D1=81=D1=82?= =?UTF-8?q?=D0=B5=D0=BC=D1=83?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pyproject.toml | 1 + scimesh/worker/__init__.py | 5 ++ scimesh/worker/artifacts.py | 57 +++++++++++++++++++ scimesh/worker/cli.py | 32 +++++++++++ scimesh/worker/config.py | 36 ++++++++++++ scimesh/worker/coordinator.py | 66 ++++++++++++++++++++++ scimesh/worker/daemon.py | 98 ++++++++++++++++++++++++++++++++ scimesh/worker/models.py | 50 +++++++++++++++++ scimesh/worker/runners.py | 60 ++++++++++++++++++++ tests/test_worker_daemon.py | 103 ++++++++++++++++++++++++++++++++++ 10 files changed, 508 insertions(+) create mode 100644 scimesh/worker/__init__.py create mode 100644 scimesh/worker/artifacts.py create mode 100644 scimesh/worker/cli.py create mode 100644 scimesh/worker/config.py create mode 100644 scimesh/worker/coordinator.py create mode 100644 scimesh/worker/daemon.py create mode 100644 scimesh/worker/models.py create mode 100644 scimesh/worker/runners.py create mode 100644 tests/test_worker_daemon.py diff --git a/pyproject.toml b/pyproject.toml index 7c37a2a..708bf1d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,6 +15,7 @@ dev = ["pytest>=8"] [project.scripts] scimesh = "scimesh.cli:main" +scimesh-worker = "scimesh.worker.cli:main" [tool.setuptools.packages.find] include = ["scimesh*"] diff --git a/scimesh/worker/__init__.py b/scimesh/worker/__init__.py new file mode 100644 index 0000000..6a351cb --- /dev/null +++ b/scimesh/worker/__init__.py @@ -0,0 +1,5 @@ +"""Worker daemon for executing coordinator-assigned SciMesh workloads.""" + +from .daemon import WorkerDaemon + +__all__ = ["WorkerDaemon"] diff --git a/scimesh/worker/artifacts.py b/scimesh/worker/artifacts.py new file mode 100644 index 0000000..0d74038 --- /dev/null +++ b/scimesh/worker/artifacts.py @@ -0,0 +1,57 @@ +"""Input/output artifact transport kept separate from the daemon state machine.""" + +from __future__ import annotations + +import hashlib +from pathlib import Path +from typing import Protocol +from urllib.request import Request, urlopen + +from .models import ClaimedTask, ProducedArtifact + + +class ArtifactClient(Protocol): + def download(self, uri: str, destination: Path) -> None: ... + + def upload(self, task: ClaimedTask, artifact: ProducedArtifact) -> str: ... + + +class HttpArtifactClient: + """Default coordinator artifact convention. + + Results are PUT to /tasks/{task_id}/artifacts/{filename}. The coordinator may + return a JSON body containing ``uri``; otherwise the upload URL is reported. + """ + + def __init__(self, coordinator_url: str, timeout: float, bearer_token: str | None = None) -> None: + self.coordinator_url = coordinator_url.rstrip("/") + self.timeout = timeout + self.bearer_token = bearer_token + + def download(self, uri: str, destination: Path) -> None: + destination.parent.mkdir(parents=True, exist_ok=True) + request = Request(uri, headers=self._auth_header()) + with urlopen(request, timeout=self.timeout) as response, destination.open("wb") as target: + while chunk := response.read(1024 * 1024): + target.write(chunk) + + def upload(self, task: ClaimedTask, artifact: ProducedArtifact) -> str: + url = f"{self.coordinator_url}/tasks/{task.task_id}/artifacts/{artifact.path.name}" + request = Request( + url, data=artifact.path.read_bytes(), method="PUT", + headers={"Content-Type": artifact.content_type, **self._auth_header()}, + ) + with urlopen(request, timeout=self.timeout) as response: + # An empty response is valid; the conventional endpoint itself is the URI. + return url if not response.read() else url + + def _auth_header(self) -> dict[str, str]: + return {"Authorization": f"Bearer {self.bearer_token}"} if self.bearer_token else {} + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() diff --git a/scimesh/worker/cli.py b/scimesh/worker/cli.py new file mode 100644 index 0000000..7556ff5 --- /dev/null +++ b/scimesh/worker/cli.py @@ -0,0 +1,32 @@ +"""Console entry point for ``scimesh-worker``.""" + +from __future__ import annotations + +import argparse +import logging +from pathlib import Path + +from .artifacts import HttpArtifactClient +from .config import WorkerConfig +from .coordinator import HttpCoordinatorClient +from .daemon import WorkerDaemon +from .runners import SciMeshRunner + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(prog="scimesh-worker") + parser.add_argument("--coordinator-url") + parser.add_argument("--worker-id") + parser.add_argument("--work-dir") + parser.add_argument("--poll-interval", type=float) + parser.add_argument("--request-timeout", type=float) + args = parser.parse_args(argv) + config = WorkerConfig.from_environment() + overrides = {key: value for key, value in vars(args).items() if value is not None} + if "work_dir" in overrides: + overrides["work_dir"] = Path(overrides["work_dir"]) + config = WorkerConfig(**{**config.__dict__, **overrides}) + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + 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 diff --git a/scimesh/worker/config.py b/scimesh/worker/config.py new file mode 100644 index 0000000..9351e10 --- /dev/null +++ b/scimesh/worker/config.py @@ -0,0 +1,36 @@ +"""Configuration parsing for the worker command.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +import os + + +@dataclass(frozen=True) +class WorkerConfig: + coordinator_url: str + worker_id: str + work_dir: Path + poll_interval: float = 2.0 + request_timeout: float = 30.0 + bearer_token: str | None = None + cleanup_after_seconds: float | None = None + capabilities: tuple[str, ...] = ("similarity-search", "similarity-graph") + + @classmethod + def from_environment(cls) -> "WorkerConfig": + url = os.getenv("SCIMESH_COORDINATOR_URL") + worker_id = os.getenv("SCIMESH_WORKER_ID") + if not url or not worker_id: + raise ValueError("SCIMESH_COORDINATOR_URL and SCIMESH_WORKER_ID are required") + cleanup = os.getenv("SCIMESH_CLEANUP_AFTER_SECONDS") + return cls( + coordinator_url=url.rstrip("/"), + worker_id=worker_id, + work_dir=Path(os.getenv("SCIMESH_WORK_DIR", "./scimesh-worker-data")), + poll_interval=float(os.getenv("SCIMESH_POLL_INTERVAL", "2")), + request_timeout=float(os.getenv("SCIMESH_REQUEST_TIMEOUT", "30")), + bearer_token=os.getenv("SCIMESH_BEARER_TOKEN"), + cleanup_after_seconds=float(cleanup) if cleanup else None, + ) diff --git a/scimesh/worker/coordinator.py b/scimesh/worker/coordinator.py new file mode 100644 index 0000000..2ec152e --- /dev/null +++ b/scimesh/worker/coordinator.py @@ -0,0 +1,66 @@ +"""HTTP boundary for the coordinator; the daemon never accesses a database.""" + +from __future__ import annotations + +import json +from typing import Any, Protocol +from urllib.error import HTTPError, URLError +from urllib.request import Request, urlopen + +from .models import ClaimedTask + + +class CoordinatorError(RuntimeError): + """A non-retriable coordinator response.""" + + +class CoordinatorTransientError(CoordinatorError): + """A timeout, connection error, or 5xx coordinator response.""" + + +class CoordinatorClient(Protocol): + def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: ... + + def submit(self, task: ClaimedTask, payload: dict[str, Any]) -> None: ... + + +class HttpCoordinatorClient: + def __init__(self, base_url: str, timeout: float, bearer_token: str | None = None) -> None: + self.base_url = base_url.rstrip("/") + self.timeout = timeout + self.bearer_token = bearer_token + + def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: + status, body = self._request("POST", "/tasks/claim", { + "worker_id": worker_id, "capabilities": list(capabilities), "max_concurrency": 1, + }) + if status == 204: + return None + if status != 200: + raise CoordinatorError(f"unexpected claim status {status}") + return ClaimedTask.from_json(body) + + def submit(self, task: ClaimedTask, payload: dict[str, Any]) -> None: + 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): + raise CoordinatorError(f"result rejected with status {status}") + + def _request(self, method: str, path: str, payload: dict[str, Any]) -> tuple[int, dict[str, Any]]: + request = Request( + f"{self.base_url}{path}", data=json.dumps(payload).encode(), method=method, + headers={"Content-Type": "application/json", **self._auth_header()}, + ) + try: + with urlopen(request, timeout=self.timeout) as response: + raw = response.read() + return response.status, json.loads(raw) if raw else {} + except HTTPError as error: + if error.code >= 500: + raise CoordinatorTransientError(f"coordinator returned {error.code}") from error + return error.code, {} + except (URLError, TimeoutError) as error: + raise CoordinatorTransientError("coordinator request failed") from error + + def _auth_header(self) -> dict[str, str]: + return {"Authorization": f"Bearer {self.bearer_token}"} if self.bearer_token else {} diff --git a/scimesh/worker/daemon.py b/scimesh/worker/daemon.py new file mode 100644 index 0000000..a18bff1 --- /dev/null +++ b/scimesh/worker/daemon.py @@ -0,0 +1,98 @@ +"""The worker state machine and its safe failure handling.""" + +from __future__ import annotations + +import logging +from pathlib import Path +import random +import shutil +import time + +from .artifacts import ArtifactClient, sha256_file +from .config import WorkerConfig +from .coordinator import CoordinatorClient, CoordinatorTransientError +from .models import ClaimedTask +from .runners import Runner + + +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 + self.log = logging.getLogger("scimesh.worker") + + def run_forever(self) -> None: + failures = 0 + while True: + try: + self._cleanup_expired_directories() + claimed = self.run_once() + failures = 0 + if not claimed: + self._sleep(self.config.poll_interval) + except CoordinatorTransientError as error: + failures += 1 + self._log("failed", error_type=type(error).__name__) + self._sleep(min(self.config.poll_interval * 2 ** min(failures, 6), 60.0)) + + def run_once(self) -> bool: + self._log("claiming") + task = self.coordinator.claim(self.config.worker_id, self.config.capabilities) + if task is None: + self._log("idle") + return False + started = time.monotonic() + task_dir = self.config.work_dir / task.task_id / str(task.attempt) + task_dir.mkdir(parents=True, exist_ok=False) + try: + self._log("downloading", task) + input_path = task_dir / "input" + self.artifacts.download(task.input.uri, input_path) + if sha256_file(input_path).lower() != task.input.sha256.lower(): + raise ValueError("input checksum mismatch") + self._log("running", task) + result = self.runner.run(task, task_dir) + self._log("uploading", task) + manifests = [ + {"uri": self.artifacts.upload(task, artifact), "sha256": sha256_file(artifact.path), "content_type": artifact.content_type} + for artifact in result.artifacts + ] + if not manifests: + raise ValueError("runner produced no artifacts") + self._log("submitting", task) + 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) + return True + + 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}) + except CoordinatorTransientError: + raise + except Exception: + self._log("failed", task, error_type="FailureReportError") + + def _log(self, state: str, task: ClaimedTask | None = None, **extra: object) -> None: + fields = {"worker_id": self.config.worker_id, "task_id": task.task_id if task else None, "attempt": task.attempt if task else None, "state": state, **extra} + self.log.info("worker_event %s", fields) + + def _cleanup_expired_directories(self) -> None: + """Remove only old task attempt directories when retention was configured.""" + if self.config.cleanup_after_seconds is None or not self.config.work_dir.exists(): + return + cutoff = time.time() - self.config.cleanup_after_seconds + for task_dir in self.config.work_dir.iterdir(): + if not task_dir.is_dir(): + continue + for attempt_dir in task_dir.iterdir(): + if attempt_dir.is_dir() and attempt_dir.stat().st_mtime < cutoff: + shutil.rmtree(attempt_dir) + if not any(task_dir.iterdir()): + task_dir.rmdir() + + @staticmethod + def _sleep(delay: float) -> None: + time.sleep(delay * random.uniform(0.75, 1.25)) diff --git a/scimesh/worker/models.py b/scimesh/worker/models.py new file mode 100644 index 0000000..97de6a1 --- /dev/null +++ b/scimesh/worker/models.py @@ -0,0 +1,50 @@ +"""Value objects shared by the worker daemon components.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +@dataclass(frozen=True) +class InputArtifact: + uri: str + sha256: str + + +@dataclass(frozen=True) +class ClaimedTask: + task_id: str + attempt: int + lease_expires_at: str + workload: str + input: InputArtifact + parameters: dict[str, Any] + + @classmethod + def from_json(cls, data: dict[str, Any]) -> "ClaimedTask": + try: + input_data = data["input"] + return cls( + task_id=str(data["task_id"]), + attempt=int(data["attempt"]), + lease_expires_at=str(data["lease_expires_at"]), + workload=str(data["workload"]), + input=InputArtifact(uri=str(input_data["uri"]), sha256=str(input_data["sha256"])), + parameters=dict(data.get("parameters", {})), + ) + except (KeyError, TypeError, ValueError) as error: + raise ValueError("invalid claimed-task response") from error + + +@dataclass(frozen=True) +class ProducedArtifact: + path: Path + content_type: str + + +@dataclass(frozen=True) +class RunResult: + artifacts: tuple[ProducedArtifact, ...] + metrics: dict[str, int | float] diff --git a/scimesh/worker/runners.py b/scimesh/worker/runners.py new file mode 100644 index 0000000..b59054b --- /dev/null +++ b/scimesh/worker/runners.py @@ -0,0 +1,60 @@ +"""Local workload adapters. They receive no arbitrary commands from the network.""" + +from __future__ import annotations + +from pathlib import Path +import subprocess +import sys +from typing import Protocol + +from .models import ClaimedTask, ProducedArtifact, RunResult + + +class Runner(Protocol): + def run(self, task: ClaimedTask, task_dir: Path) -> RunResult: ... + + +class SciMeshRunner: + """Allowlisted adapter from coordinator workloads to the local SciMesh CLI.""" + + def run(self, task: ClaimedTask, task_dir: Path) -> RunResult: + input_path = task_dir / "input" + output_path = task_dir / "result.csv" + command = [sys.executable, "-m", "scimesh.cli", task.workload, str(input_path)] + params = task.parameters + if task.workload == "similarity-search": + query_id = self._string(params, "query_id") + top_k = self._positive_int(params, "top_k", default=20) + command += ["--query-id", query_id, "--top-k", str(top_k)] + elif task.workload == "similarity-graph": + threshold = self._number(params, "threshold") + command += ["--threshold", str(threshold)] + else: + raise ValueError(f"unsupported workload: {task.workload}") + command += ["--output", str(output_path)] + subprocess.run(command, check=True, cwd=task_dir) # explicit list: never shell=True + if not output_path.is_file(): + raise RuntimeError("SciMesh CLI did not create its result") + processed_rows = max(sum(1 for _ in output_path.open(encoding="utf-8")) - 1, 0) + return RunResult((ProducedArtifact(output_path, "text/csv"),), {"processed_rows": processed_rows}) + + @staticmethod + def _string(params: dict[str, object], name: str) -> str: + value = params.get(name) + if not isinstance(value, str) or not value.strip() or len(value) > 200: + raise ValueError(f"{name} must be a non-empty string") + return value + + @staticmethod + def _positive_int(params: dict[str, object], name: str, default: int) -> int: + value = params.get(name, default) + if isinstance(value, bool) or not isinstance(value, int) or value < 1 or value > 100_000: + raise ValueError(f"{name} must be a positive integer") + return value + + @staticmethod + def _number(params: dict[str, object], name: str) -> float: + value = params.get(name) + if isinstance(value, bool) or not isinstance(value, (int, float)) or not 0 <= value <= 1: + raise ValueError(f"{name} must be a number between 0 and 1") + return float(value) diff --git a/tests/test_worker_daemon.py b/tests/test_worker_daemon.py new file mode 100644 index 0000000..4befd24 --- /dev/null +++ b/tests/test_worker_daemon.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +import hashlib +from pathlib import Path + +import pytest + +from scimesh.worker.config import WorkerConfig +from scimesh.worker.coordinator import CoordinatorTransientError +from scimesh.worker.daemon import WorkerDaemon +from scimesh.worker.models import ClaimedTask, InputArtifact, ProducedArtifact, RunResult + + +class FakeCoordinator: + def __init__(self, task: ClaimedTask | None) -> None: + self.task, self.submissions = task, [] + + def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: + task, self.task = self.task, None + return task + + def submit(self, task: ClaimedTask, payload: dict) -> None: + self.submissions.append(payload) + + +class FakeArtifacts: + def __init__(self, content: bytes) -> None: + self.content, self.uploaded = content, [] + + def download(self, uri: str, destination: Path) -> None: + destination.write_bytes(self.content) + + def upload(self, task: ClaimedTask, artifact: ProducedArtifact) -> str: + self.uploaded.append(artifact.path) + return f"https://example.test/results/{artifact.path.name}" + + +class FakeRunner: + def __init__(self) -> None: + self.calls = 0 + + def run(self, task: ClaimedTask, task_dir: Path) -> RunResult: + self.calls += 1 + output = task_dir / "result.csv" + output.write_text("id,score\na,1\n", encoding="utf-8") + return RunResult((ProducedArtifact(output, "text/csv"),), {"processed_rows": 1}) + + +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"}) + + +def daemon(tmp_path: Path, task: ClaimedTask | None, content: bytes): + coordinator, artifacts, runner = FakeCoordinator(task), FakeArtifacts(content), FakeRunner() + config = WorkerConfig("https://example.test", "worker-1", tmp_path / "work") + return WorkerDaemon(config, coordinator, artifacts, runner), coordinator, artifacts, runner, config + + +def test_claims_runs_uploads_and_submits_csv(tmp_path: Path) -> None: + content = b"input fixture" + 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.submissions[0]["status"] == "completed" + assert coordinator.submissions[0]["result"]["content_type"] == "text/csv" + + +def test_no_task_does_not_create_directory(tmp_path: Path) -> None: + worker, _, _, runner, config = daemon(tmp_path, None, b"") + assert worker.run_once() is False + assert runner.calls == 0 + assert not config.work_dir.exists() + + +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" + + +def test_transient_claim_error_is_propagated_for_bounded_backoff(tmp_path: Path) -> None: + class UnavailableCoordinator(FakeCoordinator): + def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: + raise CoordinatorTransientError("temporary outage") + + worker, _, _, _, _ = daemon(tmp_path, None, b"") + worker.coordinator = UnavailableCoordinator(None) + with pytest.raises(CoordinatorTransientError): + worker.run_once() + + +def test_task_directories_are_retained_until_cleanup_is_enabled(tmp_path: Path) -> None: + content = b"input fixture" + worker, _, _, _, config = daemon(tmp_path, make_task(content), content) + worker.run_once() + task_dir = config.work_dir / "task-1" / "1" + assert task_dir.is_dir() + worker.config = WorkerConfig(**{**config.__dict__, "cleanup_after_seconds": 0}) + worker._cleanup_expired_directories() + assert not task_dir.exists()