From 99dc321c693c33a39e8a9f9cb3132d6f3e4b94c3 Mon Sep 17 00:00:00 2001 From: reran4ik Date: Wed, 22 Jul 2026 18:40:18 +0300 Subject: [PATCH] =?UTF-8?q?=D0=98=D1=81=D0=BF=D1=80=D0=B0=D0=B2=D0=BB?= =?UTF-8?q?=D0=B5=D0=BD=D1=8B=20=D0=B7=D0=B0=D0=BC=D0=B5=D1=87=D0=B0=D0=BD?= =?UTF-8?q?=D0=B8=D1=8F=20=D1=80=D0=B5=D0=B2=D1=8C=D1=8E=20Workers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- scimesh/worker/artifacts.py | 56 ++++++++++++++++----------- scimesh/worker/cli.py | 1 + scimesh/worker/config.py | 2 + scimesh/worker/coordinator.py | 10 +++++ scimesh/worker/daemon.py | 63 ++++++++++++++++++++++++++++++- scimesh/worker/runners.py | 39 ++++++++++++++++++- tests/test_worker_daemon.py | 71 +++++++++++++++++++++++++++++++---- 7 files changed, 208 insertions(+), 34 deletions(-) diff --git a/scimesh/worker/artifacts.py b/scimesh/worker/artifacts.py index 0d74038..4497341 100644 --- a/scimesh/worker/artifacts.py +++ b/scimesh/worker/artifacts.py @@ -5,49 +5,61 @@ from __future__ import annotations import hashlib from pathlib import Path from typing import Protocol -from urllib.request import Request, urlopen +from urllib.parse import 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) + scheme = parsed.scheme.lower() + default_port = {"http": 80, "https": 443}.get(scheme) + return scheme, (parsed.hostname or "").lower(), parsed.port or default_port + + +class _SameOriginAuthRedirectHandler(HTTPRedirectHandler): + """Do not forward the coordinator token when a download changes origin.""" + + def __init__(self, coordinator_origin: tuple[str, str, int | None]) -> None: + super().__init__() + self.coordinator_origin = coordinator_origin + + def redirect_request(self, req: Request, fp: object, code: int, msg: str, headers: object, newurl: str) -> Request | None: + redirected = super().redirect_request(req, fp, code, msg, headers, newurl) + if redirected and _origin(newurl) != self.coordinator_origin: + redirected.remove_header("Authorization") + return redirected 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. + """Downloads task inputs without exposing credentials to external storage. - Results are PUT to /tasks/{task_id}/artifacts/{filename}. The coordinator may - return a JSON body containing ``uri``; otherwise the upload URL is reported. + 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: self.coordinator_url = coordinator_url.rstrip("/") self.timeout = timeout self.bearer_token = bearer_token + self.coordinator_origin = _origin(coordinator_url) + self._opener = build_opener(_SameOriginAuthRedirectHandler(self.coordinator_origin)) 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: + request = Request(uri, headers=self._auth_headers_for(uri)) + with self._opener.open(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 _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: + return {"Authorization": f"Bearer {self.bearer_token}"} + return {} def sha256_file(path: Path) -> str: digest = hashlib.sha256() diff --git a/scimesh/worker/cli.py b/scimesh/worker/cli.py index 7556ff5..e0d5414 100644 --- a/scimesh/worker/cli.py +++ b/scimesh/worker/cli.py @@ -20,6 +20,7 @@ def main(argv: list[str] | None = None) -> int: parser.add_argument("--work-dir") parser.add_argument("--poll-interval", type=float) parser.add_argument("--request-timeout", type=float) + parser.add_argument("--heartbeat-interval", 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} diff --git a/scimesh/worker/config.py b/scimesh/worker/config.py index 9351e10..cd38cb0 100644 --- a/scimesh/worker/config.py +++ b/scimesh/worker/config.py @@ -14,6 +14,7 @@ class WorkerConfig: work_dir: Path poll_interval: float = 2.0 request_timeout: float = 30.0 + heartbeat_interval: float = 15.0 bearer_token: str | None = None cleanup_after_seconds: float | None = None capabilities: tuple[str, ...] = ("similarity-search", "similarity-graph") @@ -31,6 +32,7 @@ class WorkerConfig: 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")), + heartbeat_interval=float(os.getenv("SCIMESH_HEARTBEAT_INTERVAL", "15")), 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 index 2ec152e..4061e9a 100644 --- a/scimesh/worker/coordinator.py +++ b/scimesh/worker/coordinator.py @@ -23,6 +23,8 @@ class CoordinatorClient(Protocol): def submit(self, task: ClaimedTask, payload: dict[str, Any]) -> None: ... + def heartbeat(self, task: ClaimedTask, worker_id: str) -> None: ... + class HttpCoordinatorClient: def __init__(self, base_url: str, timeout: float, bearer_token: str | None = None) -> None: @@ -46,6 +48,14 @@ 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( + "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}") + 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, diff --git a/scimesh/worker/daemon.py b/scimesh/worker/daemon.py index a18bff1..4480918 100644 --- a/scimesh/worker/daemon.py +++ b/scimesh/worker/daemon.py @@ -6,7 +6,10 @@ import logging from pathlib import Path import random 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 @@ -15,6 +18,51 @@ from .models import ClaimedTask from .runners import Runner +class LeaseHeartbeat: + """Renews a claimed task lease while local work is in progress.""" + + def __init__(self, task: ClaimedTask, coordinator: CoordinatorClient, config: WorkerConfig) -> None: + self.task, self.coordinator, self.config = task, coordinator, config + self._stop = threading.Event() + self._error: Exception | None = None + self._thread: threading.Thread | None = None + + def start(self) -> None: + # Verify ownership before expensive download or calculation begins. + 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() + + def stop(self) -> None: + self._stop.set() + if self._thread: + self._thread.join() + + def raise_if_failed(self) -> None: + if self._error: + raise self._error + + def _run(self) -> None: + 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) + except Exception as error: # Surface the lease loss in the main state machine. + self._error = error + return + delay = self.config.heartbeat_interval + + def _seconds_until_expiry(self) -> float: + try: + expiry = datetime.fromisoformat(self.task.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() + if seconds <= 0: + raise ValueError("claimed task lease has already expired") + return seconds + + 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 @@ -43,7 +91,9 @@ class WorkerDaemon: started = time.monotonic() task_dir = self.config.work_dir / task.task_id / str(task.attempt) task_dir.mkdir(parents=True, exist_ok=False) + heartbeat = LeaseHeartbeat(task, self.coordinator, self.config) try: + heartbeat.start() self._log("downloading", task) input_path = task_dir / "input" self.artifacts.download(task.input.uri, input_path) @@ -51,21 +101,30 @@ class WorkerDaemon: raise ValueError("input checksum mismatch") self._log("running", task) result = self.runner.run(task, task_dir) - self._log("uploading", task) + heartbeat.raise_if_failed() manifests = [ - {"uri": self.artifacts.upload(task, artifact), "sha256": sha256_file(artifact.path), "content_type": artifact.content_type} + self._manifest_uri(task, artifact.path.name, artifact.content_type, sha256_file(artifact.path)) for artifact in result.artifacts ] if not manifests: raise ValueError("runner produced no artifacts") self._log("submitting", task) + heartbeat.raise_if_failed() 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) + finally: + 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: diff --git a/scimesh/worker/runners.py b/scimesh/worker/runners.py index b59054b..46d74af 100644 --- a/scimesh/worker/runners.py +++ b/scimesh/worker/runners.py @@ -23,12 +23,21 @@ class SciMeshRunner: 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") + self._reject_unknown(params, {"query_id", "query_smiles", "top_k", "threshold", "threshold_direction", "max_rows", "progress_every"}) + query_id, query_smiles = params.get("query_id"), params.get("query_smiles") + if (query_id is None) == (query_smiles is None): + raise ValueError("exactly one of query_id or query_smiles is required") top_k = self._positive_int(params, "top_k", default=20) - command += ["--query-id", query_id, "--top-k", str(top_k)] + command += ["--query-id", self._string(params, "query_id")] if query_id is not None else ["--query-smiles", self._string(params, "query_smiles")] + command += ["--top-k", str(top_k)] + self._append_common_options(command, params) elif task.workload == "similarity-graph": + self._reject_unknown(params, {"threshold", "threshold_direction", "block_size", "max_rows", "progress_every"}) threshold = self._number(params, "threshold") command += ["--threshold", str(threshold)] + self._append_common_options(command, params, include_threshold=False) + if "block_size" in params: + command += ["--block-size", str(self._positive_int(params, "block_size", default=1_000))] else: raise ValueError(f"unsupported workload: {task.workload}") command += ["--output", str(output_path)] @@ -38,6 +47,25 @@ class SciMeshRunner: 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}) + def _append_common_options(self, command: list[str], params: dict[str, object], *, include_threshold: bool = True) -> None: + if include_threshold and "threshold" in params: + command += ["--threshold", str(self._number(params, "threshold"))] + if "threshold_direction" in params: + direction = params["threshold_direction"] + if direction not in ("greater", "less"): + raise ValueError("threshold_direction must be 'greater' or 'less'") + command += ["--threshold-direction", str(direction)] + if "max_rows" in params: + command += ["--max-rows", str(self._positive_int(params, "max_rows", default=1))] + if "progress_every" in params: + command += ["--progress-every", str(self._nonnegative_int(params, "progress_every"))] + + @staticmethod + def _reject_unknown(params: dict[str, object], allowed: set[str]) -> None: + unknown = set(params) - allowed + if unknown: + raise ValueError(f"unsupported parameters: {', '.join(sorted(unknown))}") + @staticmethod def _string(params: dict[str, object], name: str) -> str: value = params.get(name) @@ -52,6 +80,13 @@ class SciMeshRunner: raise ValueError(f"{name} must be a positive integer") return value + @staticmethod + def _nonnegative_int(params: dict[str, object], name: str) -> int: + value = params[name] + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise ValueError(f"{name} must be a non-negative integer") + return value + @staticmethod def _number(params: dict[str, object], name: str) -> float: value = params.get(name) diff --git a/tests/test_worker_daemon.py b/tests/test_worker_daemon.py index 4befd24..f88e3d8 100644 --- a/tests/test_worker_daemon.py +++ b/tests/test_worker_daemon.py @@ -2,6 +2,8 @@ from __future__ import annotations import hashlib from pathlib import Path +import time +from urllib.request import Request import pytest @@ -9,11 +11,13 @@ 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 +from scimesh.worker.artifacts import HttpArtifactClient, _SameOriginAuthRedirectHandler, _origin +from scimesh.worker.runners import SciMeshRunner class FakeCoordinator: def __init__(self, task: ClaimedTask | None) -> None: - self.task, self.submissions = task, [] + self.task, self.submissions, self.heartbeats = task, [], [] def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: task, self.task = self.task, None @@ -22,19 +26,17 @@ class FakeCoordinator: def submit(self, task: ClaimedTask, payload: dict) -> None: self.submissions.append(payload) + def heartbeat(self, task: ClaimedTask, worker_id: str) -> None: + self.heartbeats.append((task.task_id, task.attempt, worker_id)) + class FakeArtifacts: def __init__(self, content: bytes) -> None: - self.content, self.uploaded = content, [] + self.content = 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 @@ -61,9 +63,10 @@ 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/") def test_no_task_does_not_create_directory(tmp_path: Path) -> None: @@ -101,3 +104,55 @@ def test_task_directories_are_retained_until_cleanup_is_enabled(tmp_path: Path) worker.config = WorkerConfig(**{**config.__dict__, "cleanup_after_seconds": 0}) worker._cleanup_expired_directories() assert not task_dir.exists() + + +def test_input_token_is_sent_only_to_the_coordinator_origin() -> None: + client = HttpArtifactClient("https://coordinator.example/api", 10, "secret") + assert client._auth_headers_for("https://coordinator.example/tasks/1/input") == {"Authorization": "Bearer secret"} + assert client._auth_headers_for("https://bucket.example/presigned") == {} + + +def test_redirect_to_external_storage_strips_authorization() -> None: + handler = _SameOriginAuthRedirectHandler(_origin("https://coordinator.example")) + source = Request( + "https://coordinator.example/tasks/1/input", headers={"Authorization": "Bearer secret"} + ) + redirected = handler.redirect_request(source, None, 302, "Found", {}, "https://bucket.example/presigned") + assert redirected is not None + assert redirected.get_header("Authorization") is None + + +def test_lease_is_renewed_while_a_runner_is_still_working(tmp_path: Path) -> None: + content = b"input fixture" + worker, coordinator, _, _, config = daemon(tmp_path, make_task(content), content) + + class SlowRunner(FakeRunner): + def run(self, task: ClaimedTask, task_dir: Path) -> RunResult: + time.sleep(0.05) + return super().run(task, task_dir) + + worker.runner = SlowRunner() + worker.config = WorkerConfig(**{**config.__dict__, "heartbeat_interval": 0.01}) + worker.run_once() + assert len(coordinator.heartbeats) >= 2 + + +def test_runner_maps_graph_and_smiles_search_parameters(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + commands: list[list[str]] = [] + + def fake_run(command: list[str], **_: object) -> None: + commands.append(command) + output = Path(command[command.index("--output") + 1]) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text("a,b\n", encoding="utf-8") + + monkeypatch.setattr("scimesh.worker.runners.subprocess.run", fake_run) + runner = SciMeshRunner() + graph = ClaimedTask("graph", 1, "2026-07-30T00:00:00Z", "similarity-graph", InputArtifact("https://example/input", "x"), {"threshold": 0.2, "threshold_direction": "less", "block_size": 42, "max_rows": 7, "progress_every": 0}) + search = ClaimedTask("search", 1, "2026-07-30T00:00:00Z", "similarity-search", InputArtifact("https://example/input", "x"), {"query_smiles": "CCO", "top_k": 3}) + runner.run(graph, tmp_path / "graph") + runner.run(search, tmp_path / "search") + assert "--threshold-direction" in commands[0] and "less" in commands[0] + assert "--block-size" in commands[0] and "42" in commands[0] + assert "--max-rows" in commands[0] and "7" in commands[0] + assert "--query-smiles" in commands[1] and "CCO" in commands[1]