Replace Python worker daemon with the Go worker agent
This commit is contained in:
Emil
@@ -1,221 +0,0 @@
|
||||
"""Tests for worker token strategies and the client's 401 refresh."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from urllib.error import HTTPError
|
||||
|
||||
import pytest
|
||||
|
||||
from scimesh.worker.auth import (
|
||||
StaticTokenProvider,
|
||||
TokenExchangeError,
|
||||
WorkerKeyTokenProvider,
|
||||
provider_from_config,
|
||||
)
|
||||
from scimesh.worker.config import WorkerConfig
|
||||
from scimesh.worker.coordinator import HttpCoordinatorClient
|
||||
|
||||
|
||||
class FakeResponse:
|
||||
def __init__(self, status: int, body: bytes) -> None:
|
||||
self.status = status
|
||||
self._body = body
|
||||
|
||||
def read(self) -> bytes:
|
||||
return self._body
|
||||
|
||||
def __enter__(self) -> "FakeResponse":
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
class SeqOpener:
|
||||
"""Returns/raises a scripted sequence of responses, recording each request."""
|
||||
|
||||
def __init__(self, actions: list) -> None:
|
||||
self.actions = list(actions)
|
||||
self.requests: list = []
|
||||
|
||||
def open(self, request, timeout=None):
|
||||
self.requests.append(request)
|
||||
action = self.actions.pop(0)
|
||||
if isinstance(action, Exception):
|
||||
raise action
|
||||
return action
|
||||
|
||||
|
||||
def _exchange_response(token: str, expires_in: int) -> FakeResponse:
|
||||
return FakeResponse(
|
||||
200, json.dumps({"token": token, "expires_in": expires_in}).encode()
|
||||
)
|
||||
|
||||
|
||||
def test_static_provider_returns_fixed_token_and_never_refreshes():
|
||||
provider = StaticTokenProvider("tok")
|
||||
assert provider.token() == "tok"
|
||||
provider.refresh()
|
||||
assert provider.token() == "tok"
|
||||
|
||||
|
||||
def test_static_provider_none_means_no_auth():
|
||||
assert StaticTokenProvider(None).token() is None
|
||||
|
||||
|
||||
def test_worker_key_provider_exchanges_once_then_caches():
|
||||
clock = {"t": 1000.0}
|
||||
provider = WorkerKeyTokenProvider(
|
||||
"http://users", "scimesh_wk_live_x", timeout=5, now=lambda: clock["t"]
|
||||
)
|
||||
provider._opener = SeqOpener([_exchange_response("jwt-1", 100)]) # type: ignore[reportAttributeAccessIssue]
|
||||
|
||||
# First call exchanges; a second call well within the TTL reuses the cache.
|
||||
assert provider.token() == "jwt-1"
|
||||
clock["t"] = 1050.0 # 50s later, TTL 100s with 0.2 leeway → refresh at +80s
|
||||
assert provider.token() == "jwt-1"
|
||||
assert len(provider._opener.requests) == 1 # type: ignore[reportAttributeAccessIssue]
|
||||
|
||||
|
||||
def test_worker_key_provider_refreshes_after_leeway():
|
||||
clock = {"t": 0.0}
|
||||
provider = WorkerKeyTokenProvider(
|
||||
"http://users", "k", timeout=5, now=lambda: clock["t"]
|
||||
)
|
||||
provider._opener = SeqOpener( # type: ignore[reportAttributeAccessIssue]
|
||||
[
|
||||
_exchange_response("jwt-1", 100),
|
||||
_exchange_response("jwt-2", 100),
|
||||
]
|
||||
)
|
||||
assert provider.token() == "jwt-1"
|
||||
clock["t"] = 85.0 # past the 80s refresh point
|
||||
assert provider.token() == "jwt-2"
|
||||
assert len(provider._opener.requests) == 2 # type: ignore[reportAttributeAccessIssue]
|
||||
|
||||
|
||||
def test_worker_key_provider_force_refresh():
|
||||
provider = WorkerKeyTokenProvider("http://users", "k", timeout=5, now=lambda: 0.0)
|
||||
provider._opener = SeqOpener( # type: ignore[reportAttributeAccessIssue]
|
||||
[
|
||||
_exchange_response("jwt-1", 100),
|
||||
_exchange_response("jwt-2", 100),
|
||||
]
|
||||
)
|
||||
assert provider.token() == "jwt-1"
|
||||
provider.refresh()
|
||||
assert provider.token() == "jwt-2"
|
||||
|
||||
|
||||
def test_worker_key_provider_raises_on_rejected_key():
|
||||
provider = WorkerKeyTokenProvider("http://users", "bad", timeout=5, now=lambda: 0.0)
|
||||
provider._opener = SeqOpener( # type: ignore[reportAttributeAccessIssue]
|
||||
[HTTPError("http://users", 401, "unauthorized", {}, None)] # type: ignore[reportArgumentType]
|
||||
)
|
||||
with pytest.raises(TokenExchangeError):
|
||||
provider.token()
|
||||
|
||||
|
||||
def test_worker_key_provider_raises_when_token_missing():
|
||||
provider = WorkerKeyTokenProvider("http://users", "k", timeout=5, now=lambda: 0.0)
|
||||
provider._opener = SeqOpener( # type: ignore[reportAttributeAccessIssue]
|
||||
[FakeResponse(200, json.dumps({"expires_in": 100}).encode())]
|
||||
)
|
||||
with pytest.raises(TokenExchangeError):
|
||||
provider.token()
|
||||
|
||||
|
||||
def test_provider_from_config_selects_worker_key_mode():
|
||||
provider = provider_from_config(
|
||||
worker_key="scimesh_wk_live_x",
|
||||
userservice_url="http://users",
|
||||
bearer_token="ignored",
|
||||
request_timeout=5,
|
||||
)
|
||||
assert isinstance(provider, WorkerKeyTokenProvider)
|
||||
|
||||
|
||||
def test_provider_from_config_falls_back_to_static():
|
||||
provider = provider_from_config(
|
||||
worker_key=None, userservice_url=None, bearer_token="tok", request_timeout=5
|
||||
)
|
||||
assert isinstance(provider, StaticTokenProvider)
|
||||
assert provider.token() == "tok"
|
||||
|
||||
|
||||
class RefreshCountingProvider:
|
||||
def __init__(self) -> None:
|
||||
self.tokens = ["stale", "fresh"]
|
||||
self.index = 0
|
||||
self.refreshes = 0
|
||||
|
||||
def token(self) -> str:
|
||||
return self.tokens[min(self.index, len(self.tokens) - 1)]
|
||||
|
||||
def refresh(self) -> None:
|
||||
self.refreshes += 1
|
||||
self.index += 1
|
||||
|
||||
|
||||
def test_coordinator_client_refreshes_and_retries_once_on_401():
|
||||
provider = RefreshCountingProvider()
|
||||
client = HttpCoordinatorClient("http://coord", timeout=5, token_provider=provider)
|
||||
client._opener = SeqOpener( # type: ignore[reportAttributeAccessIssue]
|
||||
[
|
||||
HTTPError("http://coord/tasks/claim", 401, "unauthorized", {}, None), # type: ignore[reportArgumentType]
|
||||
FakeResponse(204, b""),
|
||||
]
|
||||
)
|
||||
|
||||
status, _ = client._request("POST", "/tasks/claim", {"worker_id": "w"})
|
||||
|
||||
assert status == 204
|
||||
assert provider.refreshes == 1
|
||||
# The retry carried the refreshed token.
|
||||
assert provider.index == 1
|
||||
|
||||
|
||||
def test_coordinator_client_does_not_loop_on_persistent_401():
|
||||
provider = RefreshCountingProvider()
|
||||
client = HttpCoordinatorClient("http://coord", timeout=5, token_provider=provider)
|
||||
client._opener = SeqOpener( # type: ignore[reportAttributeAccessIssue]
|
||||
[
|
||||
HTTPError("http://coord/x", 401, "unauthorized", {}, None), # type: ignore[reportArgumentType]
|
||||
HTTPError("http://coord/x", 401, "unauthorized", {}, None), # type: ignore[reportArgumentType]
|
||||
]
|
||||
)
|
||||
|
||||
status, _ = client._request("POST", "/x", {})
|
||||
|
||||
# One refresh, one retry, then the second 401 is surfaced rather than retried.
|
||||
assert status == 401
|
||||
assert provider.refreshes == 1
|
||||
|
||||
|
||||
def _base_config(**extra) -> dict:
|
||||
return {
|
||||
"coordinator_url": "http://coord",
|
||||
"worker_id": None,
|
||||
"work_dir": Path("."),
|
||||
**extra,
|
||||
}
|
||||
|
||||
|
||||
def test_worker_key_requires_userservice_url():
|
||||
with pytest.raises(ValueError, match="userservice_url"):
|
||||
WorkerConfig(**_base_config(worker_key="scimesh_wk_live_x"))
|
||||
|
||||
|
||||
def test_worker_key_with_userservice_url_is_valid():
|
||||
cfg = WorkerConfig(
|
||||
**_base_config(worker_key="scimesh_wk_live_x", userservice_url="http://users")
|
||||
)
|
||||
assert cfg.worker_key == "scimesh_wk_live_x"
|
||||
assert cfg.userservice_url == "http://users"
|
||||
|
||||
|
||||
def test_userservice_url_must_be_absolute():
|
||||
with pytest.raises(ValueError, match="userservice_url"):
|
||||
WorkerConfig(**_base_config(userservice_url="not-a-url"))
|
||||
@@ -1,842 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import time
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from urllib.request import Request
|
||||
|
||||
import pytest
|
||||
|
||||
from scimesh.worker.config import WorkerConfig, _workload_allowlist
|
||||
from scimesh.worker import cli as worker_cli
|
||||
from scimesh.worker.cli import build_parser
|
||||
from scimesh.worker.coordinator import CoordinatorTransientError
|
||||
from scimesh.worker.daemon import LeaseHeartbeat, RunOnceOutcome, WorkerDaemon
|
||||
from scimesh.worker.models import (
|
||||
ClaimedTask,
|
||||
InputArtifact,
|
||||
ProducedArtifact,
|
||||
RegisteredWorker,
|
||||
RunResult,
|
||||
UploadedArtifact,
|
||||
)
|
||||
from scimesh.worker.artifacts import (
|
||||
HttpArtifactClient,
|
||||
_SameOriginAuthRedirectHandler,
|
||||
_origin,
|
||||
)
|
||||
from scimesh.worker.runners import SciMeshRunner
|
||||
from scimesh.worker.transport import NoRedirectHandler
|
||||
|
||||
|
||||
class FakeCoordinator:
|
||||
def __init__(self, task: ClaimedTask | None) -> None:
|
||||
self.task, self.submissions, self.failures, self.heartbeats = task, [], [], []
|
||||
|
||||
def claim(
|
||||
self, worker_id: str, capabilities: tuple[str, ...]
|
||||
) -> ClaimedTask | None:
|
||||
task, self.task = self.task, None
|
||||
return task
|
||||
|
||||
def register(
|
||||
self,
|
||||
name: str,
|
||||
capabilities: tuple[str, ...],
|
||||
cpu_count: int,
|
||||
memory_mb: int | None,
|
||||
) -> RegisteredWorker:
|
||||
return RegisteredWorker("11111111-1111-4111-8111-111111111111", 15)
|
||||
|
||||
def submit(self, task: ClaimedTask, payload: dict) -> None:
|
||||
self.submissions.append(payload)
|
||||
|
||||
def fail(self, task: ClaimedTask, payload: dict) -> None:
|
||||
self.failures.append(payload)
|
||||
|
||||
def heartbeat(self, task: ClaimedTask, worker_id: str) -> str:
|
||||
self.heartbeats.append((task.task_id, task.attempt, worker_id))
|
||||
return (datetime.now(timezone.utc) + timedelta(seconds=1)).isoformat()
|
||||
|
||||
|
||||
class FakeArtifacts:
|
||||
def __init__(self, content: bytes) -> None:
|
||||
self.content, self.uploaded = content, []
|
||||
|
||||
def download(self, uri: str, destination: Path) -> None:
|
||||
destination.write_bytes(self.content)
|
||||
|
||||
def upload(
|
||||
self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact
|
||||
) -> UploadedArtifact:
|
||||
self.uploaded.append((task.task_id, worker_id, artifact.path))
|
||||
content = artifact.path.read_bytes()
|
||||
return UploadedArtifact(
|
||||
"22222222-2222-4222-8222-222222222222",
|
||||
f"https://example.test/tasks/{task.task_id}/artifacts/{artifact.path.name}",
|
||||
hashlib.sha256(content).hexdigest(),
|
||||
len(content),
|
||||
)
|
||||
|
||||
|
||||
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:
|
||||
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):
|
||||
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() == RunOnceOutcome(claimed=True, completed=True)
|
||||
assert runner.calls == 1
|
||||
assert len(artifacts.uploaded) == 1
|
||||
assert coordinator.heartbeats == [("task-1", 1, "worker-1")]
|
||||
assert "status" not in coordinator.submissions[0]
|
||||
assert coordinator.submissions[0]["result"] == {
|
||||
"artifact_id": "22222222-2222-4222-8222-222222222222"
|
||||
}
|
||||
|
||||
|
||||
def test_worker_executes_a_resolved_similarity_search_shard(tmp_path: Path) -> None:
|
||||
content = (
|
||||
b"chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\nINVALID\tnot-a-smiles\n"
|
||||
)
|
||||
task = make_task(content)
|
||||
task = ClaimedTask(
|
||||
task.task_id,
|
||||
task.attempt,
|
||||
task.lease_expires_at,
|
||||
task.workload,
|
||||
task.input,
|
||||
{"query_smiles": "CCO", "top_k": 5, "progress_every": 0},
|
||||
)
|
||||
worker, coordinator, artifacts, _, _ = daemon(tmp_path, task, content)
|
||||
worker.runner = SciMeshRunner()
|
||||
|
||||
assert worker.run_once() == RunOnceOutcome(claimed=True, completed=True)
|
||||
output = artifacts.uploaded[0][2].read_text(encoding="utf-8")
|
||||
assert output.startswith("rank,chembl_id,canonical_smiles,similarity\n")
|
||||
metrics = coordinator.submissions[0]["metrics"]
|
||||
assert metrics["scanned_rows"] == 3
|
||||
assert metrics["valid_molecules"] == 2
|
||||
assert metrics["invalid_smiles"] == 1
|
||||
assert metrics["matches_emitted"] == 1
|
||||
assert isinstance(metrics["elapsed_seconds"], float)
|
||||
|
||||
|
||||
def test_two_workers_complete_resolved_shards_after_one_retry(tmp_path: Path) -> None:
|
||||
content = b"chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\n"
|
||||
first = make_task(content)
|
||||
first = ClaimedTask(
|
||||
"retry-task",
|
||||
1,
|
||||
first.lease_expires_at,
|
||||
"similarity-search",
|
||||
first.input,
|
||||
{"query_smiles": "CCO", "top_k": 5},
|
||||
)
|
||||
second = ClaimedTask(
|
||||
"other-task",
|
||||
1,
|
||||
first.lease_expires_at,
|
||||
"similarity-search",
|
||||
first.input,
|
||||
{"query_smiles": "CCO", "top_k": 5},
|
||||
)
|
||||
|
||||
class RetryCoordinator(FakeCoordinator):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(None)
|
||||
self.queue = [first, second]
|
||||
self.claimants: list[str] = []
|
||||
|
||||
def claim(
|
||||
self, worker_id: str, capabilities: tuple[str, ...]
|
||||
) -> ClaimedTask | None:
|
||||
self.claimants.append(worker_id)
|
||||
return self.queue.pop(0) if self.queue else None
|
||||
|
||||
def fail(self, task: ClaimedTask, payload: dict) -> None:
|
||||
self.failures.append(payload)
|
||||
if (
|
||||
task.task_id == "retry-task"
|
||||
and task.attempt == 1
|
||||
and payload["retryable"]
|
||||
):
|
||||
self.queue.append(
|
||||
ClaimedTask(
|
||||
task.task_id,
|
||||
2,
|
||||
task.lease_expires_at,
|
||||
task.workload,
|
||||
task.input,
|
||||
task.parameters,
|
||||
)
|
||||
)
|
||||
|
||||
class FailFirstAttempt:
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
self.delegate = SciMeshRunner()
|
||||
|
||||
def run(self, task: ClaimedTask, task_dir: Path) -> RunResult:
|
||||
self.calls += 1
|
||||
if self.calls == 1:
|
||||
raise RuntimeError("simulated retryable shard failure")
|
||||
return self.delegate.run(task, task_dir)
|
||||
|
||||
coordinator = RetryCoordinator()
|
||||
artifacts = FakeArtifacts(content)
|
||||
worker_a = WorkerDaemon(
|
||||
WorkerConfig("https://example.test", "worker-a", tmp_path / "worker-a"),
|
||||
coordinator,
|
||||
artifacts,
|
||||
FailFirstAttempt(),
|
||||
)
|
||||
worker_b = WorkerDaemon(
|
||||
WorkerConfig("https://example.test", "worker-b", tmp_path / "worker-b"),
|
||||
coordinator,
|
||||
artifacts,
|
||||
SciMeshRunner(),
|
||||
)
|
||||
|
||||
assert worker_a.run_once() == RunOnceOutcome(claimed=True, completed=False)
|
||||
assert worker_b.run_once() == RunOnceOutcome(claimed=True, completed=True)
|
||||
assert worker_a.run_once() == RunOnceOutcome(claimed=True, completed=True)
|
||||
assert coordinator.claimants == ["worker-a", "worker-b", "worker-a"]
|
||||
assert len(coordinator.failures) == 1
|
||||
assert coordinator.failures[0]["retryable"] is True
|
||||
assert len(coordinator.submissions) == 2
|
||||
|
||||
|
||||
def test_no_task_does_not_create_directory(tmp_path: Path) -> None:
|
||||
worker, _, _, runner, config = daemon(tmp_path, None, b"")
|
||||
assert worker.run_once() == RunOnceOutcome(claimed=False, completed=False)
|
||||
assert runner.calls == 0
|
||||
assert not config.work_dir.exists()
|
||||
|
||||
|
||||
def test_once_worker_exits_after_an_empty_claim(
|
||||
tmp_path: Path, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
caplog.set_level(logging.INFO, logger="scimesh.worker")
|
||||
worker, _, _, runner, _ = daemon(tmp_path, None, b"")
|
||||
worker.config = WorkerConfig(
|
||||
**{**worker.config.__dict__, "exit_when_idle": True, "max_tasks": 1}
|
||||
)
|
||||
assert worker.run_forever() is True
|
||||
assert runner.calls == 0
|
||||
assert "queue_empty" in caplog.text
|
||||
|
||||
|
||||
def test_worker_stops_after_the_configured_number_of_claims(
|
||||
tmp_path: Path, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
caplog.set_level(logging.INFO, logger="scimesh.worker")
|
||||
content = b"input fixture"
|
||||
worker, _, _, runner, _ = daemon(tmp_path, make_task(content), content)
|
||||
worker.config = WorkerConfig(**{**worker.config.__dict__, "max_tasks": 1})
|
||||
assert worker.run_forever() is True
|
||||
assert runner.calls == 1
|
||||
assert "max_tasks_reached" in caplog.text
|
||||
|
||||
|
||||
def test_keyboard_interrupt_stops_worker_without_propagating(
|
||||
tmp_path: Path, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
caplog.set_level(logging.INFO, logger="scimesh.worker")
|
||||
|
||||
class InterruptingCoordinator(FakeCoordinator):
|
||||
def claim(
|
||||
self, worker_id: str, capabilities: tuple[str, ...]
|
||||
) -> ClaimedTask | None:
|
||||
raise KeyboardInterrupt
|
||||
|
||||
worker, _, _, _, _ = daemon(tmp_path, None, b"")
|
||||
worker.coordinator = InterruptingCoordinator(None)
|
||||
assert worker.run_forever() is False
|
||||
assert "interrupted" in caplog.text
|
||||
|
||||
|
||||
def test_interrupting_an_active_task_reports_a_sanitized_failure(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
content = b"input fixture"
|
||||
worker, coordinator, _, _, _ = daemon(tmp_path, make_task(content), content)
|
||||
|
||||
class InterruptingRunner(FakeRunner):
|
||||
def run(self, task: ClaimedTask, task_dir: Path) -> RunResult:
|
||||
raise KeyboardInterrupt
|
||||
|
||||
worker.runner = InterruptingRunner()
|
||||
with pytest.raises(KeyboardInterrupt):
|
||||
worker.run_once()
|
||||
assert coordinator.failures == [
|
||||
{
|
||||
"worker_id": "worker-1",
|
||||
"attempt": 1,
|
||||
"error_code": "InterruptedError",
|
||||
"error_message": "worker interrupted by operator",
|
||||
"retryable": True,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_max_tasks_counts_successes_not_failed_claims(tmp_path: Path) -> None:
|
||||
successful_content = b"successful input"
|
||||
|
||||
class SequencedCoordinator(FakeCoordinator):
|
||||
def __init__(self) -> None:
|
||||
super().__init__(None)
|
||||
self.tasks = [
|
||||
make_task(b"bad input", "wrong-checksum"),
|
||||
ClaimedTask(
|
||||
"task-2",
|
||||
1,
|
||||
(datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat(),
|
||||
"similarity-search",
|
||||
InputArtifact(
|
||||
"https://example.test/input",
|
||||
hashlib.sha256(successful_content).hexdigest(),
|
||||
),
|
||||
{"query_id": "CHEMBL1"},
|
||||
),
|
||||
]
|
||||
|
||||
def claim(
|
||||
self, worker_id: str, capabilities: tuple[str, ...]
|
||||
) -> ClaimedTask | None:
|
||||
return self.tasks.pop(0) if self.tasks else None
|
||||
|
||||
coordinator = SequencedCoordinator()
|
||||
artifacts, runner = FakeArtifacts(successful_content), FakeRunner()
|
||||
config = WorkerConfig(
|
||||
"https://example.test", "worker-1", tmp_path / "work", max_tasks=1
|
||||
)
|
||||
worker = WorkerDaemon(config, coordinator, artifacts, runner)
|
||||
assert worker.run_forever() is True
|
||||
assert len(coordinator.failures) == 1
|
||||
assert len(coordinator.submissions) == 1
|
||||
assert runner.calls == 1
|
||||
|
||||
|
||||
def test_worker_cli_lifecycle_options_are_explicit_and_exclusive() -> None:
|
||||
parser = build_parser()
|
||||
assert parser.parse_args(["--once"]).once is True
|
||||
assert parser.parse_args(["--max-tasks", "2"]).max_tasks == 2
|
||||
with pytest.raises(SystemExit):
|
||||
parser.parse_args(["--once", "--max-tasks", "2"])
|
||||
|
||||
|
||||
def test_worker_cli_uses_a_nonzero_exit_code_for_interruption(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
class InterruptedDaemon:
|
||||
def __init__(self, *_: object) -> None:
|
||||
pass
|
||||
|
||||
def run_forever(self) -> bool:
|
||||
return False
|
||||
|
||||
monkeypatch.setattr(worker_cli, "WorkerDaemon", InterruptedDaemon)
|
||||
assert (
|
||||
worker_cli.main(
|
||||
["--coordinator-url", "https://example.test", "--work-dir", str(tmp_path)]
|
||||
)
|
||||
== 130
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", [0, -1, True])
|
||||
def test_max_tasks_must_be_positive(value: object, tmp_path: Path) -> None:
|
||||
with pytest.raises(ValueError, match="max_tasks"):
|
||||
WorkerConfig("https://example.test", None, tmp_path, max_tasks=value) # type: ignore[arg-type]
|
||||
|
||||
|
||||
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() == RunOnceOutcome(claimed=True, completed=False)
|
||||
assert runner.calls == 0
|
||||
assert coordinator.failures[0]["error_code"] == "ValueError"
|
||||
assert coordinator.failures[0]["retryable"] is False
|
||||
assert not coordinator.submissions
|
||||
|
||||
|
||||
def test_failure_reporting_removes_paths_outside_the_worker_directory(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
worker, coordinator, _, _, _ = daemon(tmp_path, make_task(b"input"), b"input")
|
||||
error = subprocess.CalledProcessError(
|
||||
1,
|
||||
["/home/alice/.venv/bin/python", "-m", "scimesh.cli", "/private/input.tsv"],
|
||||
)
|
||||
worker._report_failure(make_task(b"input"), error)
|
||||
message = coordinator.failures[0]["error_message"]
|
||||
assert "/home/alice" not in message
|
||||
assert "/private/input.tsv" not in message
|
||||
assert "<path>" in message
|
||||
|
||||
|
||||
def test_directory_creation_failure_is_reported(tmp_path: Path) -> None:
|
||||
content = b"input fixture"
|
||||
worker, coordinator, _, _, config = daemon(tmp_path, make_task(content), content)
|
||||
(config.work_dir / "task-1" / "1").mkdir(parents=True)
|
||||
assert worker.run_once() == RunOnceOutcome(claimed=True, completed=False)
|
||||
assert coordinator.failures[0]["error_code"] == "FileExistsError"
|
||||
|
||||
|
||||
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()
|
||||
|
||||
|
||||
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_relative_input_uri_is_resolved_against_the_coordinator() -> None:
|
||||
client = HttpArtifactClient("https://coordinator.example/api", 10, "secret")
|
||||
assert client._auth_headers_for("https://coordinator.example/tasks/1/input") == {
|
||||
"Authorization": "Bearer secret"
|
||||
}
|
||||
# The coordinator's contract returns root-relative artifact paths.
|
||||
assert client.coordinator_url == "https://coordinator.example/api"
|
||||
|
||||
|
||||
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" # type: ignore[arg-type]
|
||||
)
|
||||
assert redirected is not None
|
||||
assert redirected.get_header("Authorization") is None
|
||||
|
||||
|
||||
def test_api_requests_never_follow_redirects() -> None:
|
||||
handler = NoRedirectHandler()
|
||||
request = Request(
|
||||
"https://coordinator.example/tasks/claim",
|
||||
headers={"Authorization": "Bearer secret"},
|
||||
)
|
||||
assert (
|
||||
handler.redirect_request(
|
||||
request, None, 302, "Found", {}, "https://other.example"
|
||||
)
|
||||
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_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_executes_search_through_the_sdk_and_rejects_graph(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
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},
|
||||
)
|
||||
search_dir = tmp_path / "search"
|
||||
search_dir.mkdir()
|
||||
(search_dir / "input").write_text(
|
||||
"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8"
|
||||
)
|
||||
with pytest.raises(ValueError, match="unsupported workload"):
|
||||
runner.run(graph, tmp_path / "graph")
|
||||
result = runner.run(search, search_dir)
|
||||
assert result.metrics == {
|
||||
"scanned_rows": 2,
|
||||
"valid_molecules": 2,
|
||||
"invalid_smiles": 0,
|
||||
"matches_emitted": 1,
|
||||
}
|
||||
assert result.artifacts[0].content_type == "text/csv"
|
||||
assert (
|
||||
result.artifacts[0]
|
||||
.path.read_text(encoding="utf-8")
|
||||
.startswith("rank,chembl_id,canonical_smiles,similarity\n")
|
||||
)
|
||||
|
||||
|
||||
def test_runner_resolves_query_id_from_the_shard_and_rejects_plan_time_parameters(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
task_dir = tmp_path / "search"
|
||||
task_dir.mkdir()
|
||||
(task_dir / "input").write_text(
|
||||
"chembl_id\tcanonical_smiles\nQUERY\tCCO\nMATCH\tCCCO\n", encoding="utf-8"
|
||||
)
|
||||
task = ClaimedTask(
|
||||
"search",
|
||||
1,
|
||||
"2026-07-30T00:00:00Z",
|
||||
"similarity-search",
|
||||
InputArtifact("https://example/input", "a" * 64),
|
||||
{"query_id": "QUERY", "top_k": 5},
|
||||
)
|
||||
result = SciMeshRunner().run(task, task_dir)
|
||||
assert result.metrics["matches_emitted"] == 1
|
||||
assert (task_dir / "result.csv").is_file()
|
||||
|
||||
with_max_rows = ClaimedTask(
|
||||
"search",
|
||||
1,
|
||||
"2026-07-30T00:00:00Z",
|
||||
"similarity-search",
|
||||
InputArtifact("https://example/input", "a" * 64),
|
||||
{"query_smiles": "CCO", "max_rows": 1},
|
||||
)
|
||||
with pytest.raises(ValueError, match="outside the stage projection"):
|
||||
SciMeshRunner().run(with_max_rows, task_dir)
|
||||
|
||||
|
||||
def test_runner_accepts_coordinator_workload_names(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
task_dir = tmp_path / "search"
|
||||
task_dir.mkdir()
|
||||
(task_dir / "input").write_text(
|
||||
"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8"
|
||||
)
|
||||
task = ClaimedTask(
|
||||
"search",
|
||||
1,
|
||||
"2026-07-30T00:00:00Z",
|
||||
"similarity_search",
|
||||
InputArtifact("https://example/input", "a" * 64),
|
||||
{"query_smiles": "CCO"},
|
||||
)
|
||||
result = SciMeshRunner().run(task, task_dir)
|
||||
assert result.metrics["matches_emitted"] == 1
|
||||
assert (task_dir / "result.csv").is_file()
|
||||
|
||||
|
||||
def test_claimed_task_rejects_path_traversal_and_invalid_metadata() -> None:
|
||||
payload = {
|
||||
"task_id": "../outside",
|
||||
"attempt": 1,
|
||||
"lease_expires_at": "2026-07-30T00:00:00Z",
|
||||
"workload": "similarity-search",
|
||||
"input": {"uri": "https://example.test/input", "sha256": "a" * 64},
|
||||
"parameters": {},
|
||||
}
|
||||
with pytest.raises(ValueError, match="invalid claimed-task response"):
|
||||
ClaimedTask.from_json(payload)
|
||||
|
||||
payload["task_id"] = "11111111-1111-4111-8111-111111111111"
|
||||
payload["input"] = {"uri": "//outside.example/input", "sha256": "a" * 64}
|
||||
with pytest.raises(ValueError, match="invalid claimed-task response"):
|
||||
ClaimedTask.from_json(payload)
|
||||
|
||||
payload["input"] = {"uri": "/tasks/../outside/input", "sha256": "a" * 64}
|
||||
with pytest.raises(ValueError, match="invalid claimed-task response"):
|
||||
ClaimedTask.from_json(payload)
|
||||
|
||||
|
||||
def test_claimed_task_accepts_a_coordinator_relative_input_path() -> None:
|
||||
task = ClaimedTask.from_json(
|
||||
{
|
||||
"task_id": "11111111-1111-4111-8111-111111111111",
|
||||
"attempt": 1,
|
||||
"lease_expires_at": "2026-07-30T00:00:00Z",
|
||||
"workload": "similarity_search",
|
||||
"input": {
|
||||
"uri": "/tasks/11111111-1111-4111-8111-111111111111/input",
|
||||
"sha256": "a" * 64,
|
||||
},
|
||||
"parameters": {},
|
||||
}
|
||||
)
|
||||
assert task.input.uri.startswith("/tasks/")
|
||||
|
||||
|
||||
def test_uploaded_artifact_requires_complete_durable_metadata() -> None:
|
||||
artifact = UploadedArtifact.from_json(
|
||||
{
|
||||
"artifact_id": "22222222-2222-4222-8222-222222222222",
|
||||
"uri": "https://coordinator.example/artifacts/222/download",
|
||||
"sha256": "a" * 64,
|
||||
"size_bytes": 12,
|
||||
}
|
||||
)
|
||||
assert artifact.size_bytes == 12
|
||||
with pytest.raises(ValueError, match="artifact size_bytes"):
|
||||
UploadedArtifact.from_json({"artifact_id": "missing"})
|
||||
|
||||
|
||||
def test_environment_overrides_allow_cli_only_configuration(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
monkeypatch.delenv("SCIMESH_COORDINATOR_URL", raising=False)
|
||||
config = WorkerConfig.from_environment(
|
||||
{
|
||||
"coordinator_url": "https://coordinator.example",
|
||||
"work_dir": tmp_path,
|
||||
"worker_name": "test-worker",
|
||||
}
|
||||
)
|
||||
assert config.coordinator_url == "https://coordinator.example"
|
||||
assert config.worker_id is None
|
||||
assert "similarity-search" in config.capabilities
|
||||
assert "similarity_search" in config.capabilities
|
||||
assert "similarity-graph" not in config.capabilities
|
||||
|
||||
|
||||
def test_relative_work_dir_is_normalized_for_runner_subprocesses(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
monkeypatch.chdir(tmp_path)
|
||||
config = WorkerConfig("https://coordinator.example", None, Path("./worker-data"))
|
||||
assert config.work_dir == tmp_path / "worker-data"
|
||||
|
||||
task_dir = config.work_dir / "task" / "1"
|
||||
task_dir.mkdir(parents=True)
|
||||
(task_dir / "input").write_text(
|
||||
"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCO\n", encoding="utf-8"
|
||||
)
|
||||
task = ClaimedTask(
|
||||
"task",
|
||||
1,
|
||||
"2026-07-30T00:00:00Z",
|
||||
"similarity-search",
|
||||
InputArtifact("https://example.test/input", "a" * 64),
|
||||
{"query_smiles": "CCO"},
|
||||
)
|
||||
SciMeshRunner().run(task, task_dir)
|
||||
assert (task_dir / "result.csv").is_file()
|
||||
|
||||
|
||||
def test_worker_registration_sets_returned_identity(tmp_path: Path) -> None:
|
||||
worker, _, _, _, _ = daemon(tmp_path, None, b"")
|
||||
worker._register_worker()
|
||||
assert worker.worker_id == "11111111-1111-4111-8111-111111111111"
|
||||
assert worker.config.heartbeat_interval == 15
|
||||
|
||||
|
||||
def test_runner_executes_an_arbitrary_sdk_workload(tmp_path: Path) -> None:
|
||||
from scimesh.workloads.descriptors import descriptor_batch_sdk_definition
|
||||
|
||||
content = b"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\n"
|
||||
task = ClaimedTask(
|
||||
"task-1",
|
||||
1,
|
||||
(datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat(),
|
||||
"descriptor-batch",
|
||||
InputArtifact("https://example.test/input", hashlib.sha256(content).hexdigest()),
|
||||
{"skip_invalid": True},
|
||||
)
|
||||
task_dir = tmp_path / "task-1" / "1"
|
||||
task_dir.mkdir(parents=True)
|
||||
(task_dir / "input").write_bytes(content)
|
||||
runner = SciMeshRunner(
|
||||
definitions={
|
||||
"descriptor-batch": descriptor_batch_sdk_definition().definition()
|
||||
}
|
||||
)
|
||||
|
||||
result = runner.run(task, task_dir)
|
||||
header = result.artifacts[0].path.read_text(encoding="utf-8").splitlines()[0]
|
||||
assert header.startswith("chembl_id,canonical_smiles,ExactMolWt")
|
||||
assert result.metrics["rows_emitted"] == 2
|
||||
|
||||
|
||||
def test_runner_rejects_workloads_outside_the_v1_single_input_contract(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
from scimesh.workloads.graph import similarity_graph_sdk_definition
|
||||
|
||||
content = b"chembl_id\tcanonical_smiles\nA\tCCO\nB\tCCCC\n"
|
||||
task = ClaimedTask(
|
||||
"graph-task",
|
||||
1,
|
||||
(datetime.now(timezone.utc) + timedelta(seconds=60)).isoformat(),
|
||||
"similarity-graph",
|
||||
InputArtifact("https://example.test/input", hashlib.sha256(content).hexdigest()),
|
||||
{"threshold": 0.5},
|
||||
)
|
||||
task_dir = tmp_path / "graph"
|
||||
task_dir.mkdir(parents=True)
|
||||
(task_dir / "input").write_bytes(content)
|
||||
runner = SciMeshRunner(
|
||||
definitions={
|
||||
"similarity-graph": similarity_graph_sdk_definition().definition()
|
||||
}
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="v1 single-input contract"):
|
||||
runner.run(task, task_dir)
|
||||
|
||||
|
||||
def test_runner_for_worker_discovers_allowlisted_installed_workloads(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
from scimesh.sdk.registry import WorkloadRegistry
|
||||
from scimesh.workloads.search import similarity_search_sdk_definition
|
||||
|
||||
definition = similarity_search_sdk_definition().definition()
|
||||
|
||||
def fake_discover(self: WorkloadRegistry, allowlist) -> None:
|
||||
assert len(allowlist) == 1
|
||||
self.register(definition, enabled=True)
|
||||
|
||||
monkeypatch.setattr(WorkloadRegistry, "discover_installed", fake_discover)
|
||||
allowlist = _workload_allowlist(
|
||||
'[{"distribution": "scimesh", "name": "similarity-search", '
|
||||
'"version": "1.0.0", "digest": "sha256:' + "a" * 64 + '"}]'
|
||||
)
|
||||
config = WorkerConfig(
|
||||
"https://example.test", "worker-1", tmp_path / "work",
|
||||
capabilities=("similarity-search",),
|
||||
workload_allowlist=allowlist,
|
||||
)
|
||||
runner = SciMeshRunner.for_worker(config)
|
||||
assert set(runner._definitions) == {"similarity-search"}
|
||||
|
||||
|
||||
def test_worker_config_parses_capabilities_and_workload_allowlist(
|
||||
monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
monkeypatch.setenv("SCIMESH_CAPABILITIES", "similarity-search,descriptor-batch")
|
||||
config = WorkerConfig.from_environment(
|
||||
{"coordinator_url": "https://example.test", "work_dir": tmp_path}
|
||||
)
|
||||
assert config.capabilities == ("similarity-search", "descriptor-batch")
|
||||
assert config.workload_allowlist == ()
|
||||
|
||||
monkeypatch.setenv(
|
||||
"SCIMESH_WORKLOAD_ALLOWLIST",
|
||||
'[{"distribution": "scimesh", "name": "descriptor-batch", '
|
||||
'"version": "1.0.0", "digest": "sha256:' + "b" * 64 + '"}]',
|
||||
)
|
||||
config = WorkerConfig.from_environment(
|
||||
{"coordinator_url": "https://example.test", "work_dir": tmp_path}
|
||||
)
|
||||
assert len(config.workload_allowlist) == 1
|
||||
assert config.workload_allowlist[0].workload.name == "descriptor-batch"
|
||||
|
||||
monkeypatch.setenv("SCIMESH_WORKLOAD_ALLOWLIST", "not-json")
|
||||
with pytest.raises(ValueError, match="valid JSON"):
|
||||
WorkerConfig.from_environment(
|
||||
{"coordinator_url": "https://example.test", "work_dir": tmp_path}
|
||||
)
|
||||
Reference in New Issue
Block a user