Исправлены замечания ревью Workers

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
+    The present coordinator contract persists a result *manifest* at ``/result``
-    return a JSON body containing ``uri``; otherwise the upload URL is reported.
+    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())
+        request = Request(uri, headers=self._auth_headers_for(uri))
-        with urlopen(request, timeout=self.timeout) as response, destination.open("wb") as target:
+        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:
+    def _auth_headers_for(self, uri: str) -> dict[str, str]:
-        url = f"{self.coordinator_url}/tasks/{task.task_id}/artifacts/{artifact.path.name}"
+        """Only coordinator-owned URLs receive the coordinator bearer token."""
-        request = Request(
+        if self.bearer_token and _origin(uri) == self.coordinator_origin:
-            url, data=artifact.path.read_bytes(), method="PUT",
+            return {"Authorization": f"Bearer {self.bearer_token}"}
-            headers={"Content-Type": artifact.content_type, **self._auth_header()},
+        return {}
-        )
-        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()

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}

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,
         )

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,

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), "<worker-dir>")[:300]
         try:

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)

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]