Replace Python worker daemon with the Go worker agent
This commit is contained in:
Emil
@@ -1,5 +1,8 @@
|
||||
"""Worker daemon for executing coordinator-assigned SciMesh workloads."""
|
||||
"""Per-task execution for the Go worker agent.
|
||||
|
||||
from .daemon import WorkerDaemon
|
||||
|
||||
__all__ = ["WorkerDaemon"]
|
||||
This package contains the Python side of a claimed task: the wire value
|
||||
objects (``models``), the SDK execution bridge (``runners``), and the
|
||||
command-line task entry (``task``). The agent lifecycle itself lives in the
|
||||
Go worker agent (``coordinator/internal/agent``); this package is only ever
|
||||
invoked by it, one process per task.
|
||||
"""
|
||||
|
||||
@@ -1,141 +0,0 @@
|
||||
"""Input/output artifact transport kept separate from the daemon state machine."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import http.client
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Protocol
|
||||
from urllib.error import HTTPError
|
||||
from urllib.parse import quote, urljoin, urlsplit
|
||||
from urllib.request import Request, build_opener
|
||||
|
||||
from .auth import StaticTokenProvider, TokenProvider
|
||||
from .coordinator import CoordinatorConflictError
|
||||
from .models import ClaimedTask, ProducedArtifact, UploadedArtifact
|
||||
from .transport import SameOriginAuthRedirectHandler, origin
|
||||
|
||||
# Compatibility aliases for focused transport tests.
|
||||
_SameOriginAuthRedirectHandler = SameOriginAuthRedirectHandler
|
||||
_origin = origin
|
||||
|
||||
class ArtifactClient(Protocol):
|
||||
def download(self, uri: str, destination: Path) -> None: ...
|
||||
|
||||
def upload(
|
||||
self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact
|
||||
) -> UploadedArtifact: ...
|
||||
|
||||
|
||||
class HttpArtifactClient:
|
||||
"""Transfers artifacts through the coordinator without leaking credentials."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
coordinator_url: str,
|
||||
timeout: float,
|
||||
bearer_token: str | None = None,
|
||||
*,
|
||||
token_provider: TokenProvider | None = None,
|
||||
) -> None:
|
||||
self.coordinator_url = coordinator_url.rstrip("/")
|
||||
self.timeout = timeout
|
||||
self._tokens: TokenProvider = token_provider or StaticTokenProvider(bearer_token)
|
||||
self.coordinator_origin = origin(coordinator_url)
|
||||
self._opener = build_opener(SameOriginAuthRedirectHandler(self.coordinator_origin))
|
||||
|
||||
@property
|
||||
def bearer_token(self) -> str | None:
|
||||
return self._tokens.token()
|
||||
|
||||
def download(self, uri: str, destination: Path) -> None:
|
||||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||||
resolved_uri = urljoin(f"{self.coordinator_url}/", uri)
|
||||
self._download_once(resolved_uri, destination, allow_refresh=True)
|
||||
|
||||
def _download_once(self, resolved_uri: str, destination: Path, *, allow_refresh: bool) -> None:
|
||||
request = Request(resolved_uri, headers=self._auth_headers_for(resolved_uri))
|
||||
try:
|
||||
with self._opener.open(request, timeout=self.timeout) as response, destination.open("wb") as target:
|
||||
while chunk := response.read(1024 * 1024):
|
||||
target.write(chunk)
|
||||
except HTTPError as error:
|
||||
# Refresh an expired token and retry once, mirroring the coordinator
|
||||
# client, so a token that lapses mid-task does not fail the download.
|
||||
if error.code == 401 and allow_refresh:
|
||||
self._tokens.refresh()
|
||||
self._download_once(resolved_uri, destination, allow_refresh=False)
|
||||
return
|
||||
raise
|
||||
|
||||
def upload(
|
||||
self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact
|
||||
) -> UploadedArtifact:
|
||||
return self._upload_once(task, worker_id, artifact, allow_refresh=True)
|
||||
|
||||
def _upload_once(
|
||||
self, task: ClaimedTask, worker_id: str, artifact: ProducedArtifact, *, allow_refresh: bool
|
||||
) -> UploadedArtifact:
|
||||
"""Stream an artifact and require durable coordinator-owned metadata."""
|
||||
url = (
|
||||
f"{self.coordinator_url}/tasks/{quote(task.task_id, safe='')}/artifacts/"
|
||||
f"{quote(artifact.path.name, safe='')}"
|
||||
)
|
||||
parsed = urlsplit(url)
|
||||
if parsed.scheme not in {"http", "https"} or not parsed.hostname:
|
||||
raise ValueError("coordinator URL must be an absolute HTTP(S) URL")
|
||||
connection_class = (
|
||||
http.client.HTTPSConnection if parsed.scheme == "https" else http.client.HTTPConnection
|
||||
)
|
||||
connection = connection_class(parsed.hostname, parsed.port, timeout=self.timeout)
|
||||
local_size = artifact.path.stat().st_size
|
||||
local_sha256 = sha256_file(artifact.path)
|
||||
try:
|
||||
path = parsed.path + (f"?{parsed.query}" if parsed.query else "")
|
||||
connection.putrequest("PUT", path)
|
||||
connection.putheader("Content-Type", artifact.content_type)
|
||||
connection.putheader("Content-Length", str(local_size))
|
||||
connection.putheader("X-Worker-ID", worker_id)
|
||||
connection.putheader("X-Task-Attempt", str(task.attempt))
|
||||
for name, value in self._auth_headers_for(url).items():
|
||||
connection.putheader(name, value)
|
||||
connection.endheaders()
|
||||
with artifact.path.open("rb") as source:
|
||||
while chunk := source.read(1024 * 1024):
|
||||
connection.send(chunk)
|
||||
response = connection.getresponse()
|
||||
body = response.read()
|
||||
if response.status == 401 and allow_refresh:
|
||||
# Token lapsed mid-task: refresh and retry the upload once.
|
||||
self._tokens.refresh()
|
||||
connection.close()
|
||||
return self._upload_once(task, worker_id, artifact, allow_refresh=False)
|
||||
if response.status == 409:
|
||||
raise CoordinatorConflictError("artifact upload rejected because the task lease was lost")
|
||||
if response.status != 200:
|
||||
raise RuntimeError(f"artifact upload rejected with status {response.status}")
|
||||
try:
|
||||
response_data = json.loads(body)
|
||||
uploaded = UploadedArtifact.from_json(response_data)
|
||||
except (ValueError, json.JSONDecodeError) as error:
|
||||
raise RuntimeError("artifact upload returned invalid metadata") from error
|
||||
if uploaded.sha256 != local_sha256 or uploaded.size_bytes != local_size:
|
||||
raise RuntimeError("artifact upload metadata does not match local artifact")
|
||||
return uploaded
|
||||
finally:
|
||||
connection.close()
|
||||
|
||||
def _auth_headers_for(self, uri: str) -> dict[str, str]:
|
||||
"""Only coordinator-owned URLs receive the coordinator bearer token."""
|
||||
token = self._tokens.token()
|
||||
if token and origin(uri) == self.coordinator_origin:
|
||||
return {"Authorization": f"Bearer {token}"}
|
||||
return {}
|
||||
|
||||
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()
|
||||
@@ -1,128 +0,0 @@
|
||||
"""Bearer-token strategies for the worker's coordinator calls.
|
||||
|
||||
A worker authenticates in one of two ways:
|
||||
|
||||
* a *static* token — the shared service token or a directly supplied JWT, fixed
|
||||
for the life of the process; or
|
||||
* a *worker key* — a long-lived per-user credential the worker trades for a
|
||||
short-lived JWT at the userservice, refreshing before that JWT expires.
|
||||
|
||||
Both are exposed through the small ``TokenProvider`` protocol so the HTTP
|
||||
clients neither know nor care which one is in play.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from typing import Callable, Protocol
|
||||
from urllib.error import HTTPError, URLError
|
||||
from urllib.request import Request, build_opener
|
||||
|
||||
from .transport import NoRedirectHandler
|
||||
|
||||
|
||||
class TokenExchangeError(RuntimeError):
|
||||
"""The userservice refused or failed to exchange a worker key."""
|
||||
|
||||
|
||||
class TokenProvider(Protocol):
|
||||
def token(self) -> str | None:
|
||||
"""Return the current bearer token, refreshing it if necessary."""
|
||||
|
||||
def refresh(self) -> None:
|
||||
"""Force the next token to be re-fetched (e.g. after a 401)."""
|
||||
|
||||
|
||||
class StaticTokenProvider:
|
||||
"""Serves a fixed token forever. ``None`` means "send no Authorization"."""
|
||||
|
||||
def __init__(self, token: str | None) -> None:
|
||||
self._token = token
|
||||
|
||||
def token(self) -> str | None:
|
||||
return self._token
|
||||
|
||||
def refresh(self) -> None: # noqa: D401 - nothing to refresh
|
||||
return None
|
||||
|
||||
|
||||
class WorkerKeyTokenProvider:
|
||||
"""Exchanges a long-lived worker key for short-lived JWTs and refreshes them.
|
||||
|
||||
The token is cached until roughly ``1 - refresh_leeway`` of its lifetime has
|
||||
elapsed, so the worker renews ahead of expiry instead of waiting for a 401.
|
||||
A monotonic clock is injectable to keep tests deterministic.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
userservice_url: str,
|
||||
worker_key: str,
|
||||
timeout: float,
|
||||
*,
|
||||
refresh_leeway: float = 0.2,
|
||||
now: Callable[[], float] = time.monotonic,
|
||||
) -> None:
|
||||
self._url = userservice_url.rstrip("/")
|
||||
self._key = worker_key
|
||||
self._timeout = timeout
|
||||
self._leeway = refresh_leeway
|
||||
self._now = now
|
||||
self._token: str | None = None
|
||||
self._refresh_at: float = 0.0
|
||||
self._opener = build_opener(NoRedirectHandler())
|
||||
|
||||
def token(self) -> str:
|
||||
if self._token is None or self._now() >= self._refresh_at:
|
||||
self._exchange()
|
||||
assert self._token is not None # _exchange sets it or raises
|
||||
return self._token
|
||||
|
||||
def refresh(self) -> None:
|
||||
self._exchange()
|
||||
|
||||
def _exchange(self) -> None:
|
||||
request = Request(
|
||||
f"{self._url}/worker-tokens/exchange",
|
||||
data=json.dumps({"key": self._key}).encode(),
|
||||
method="POST",
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
try:
|
||||
with self._opener.open(request, timeout=self._timeout) as response:
|
||||
raw = response.read()
|
||||
data = json.loads(raw) if raw else {}
|
||||
except HTTPError as error:
|
||||
# A revoked or unknown key is a permanent 401; there is nothing the
|
||||
# worker can do but stop, so surface it rather than retry forever.
|
||||
raise TokenExchangeError(
|
||||
f"worker key exchange rejected with status {error.code}"
|
||||
) from error
|
||||
except (URLError, TimeoutError, json.JSONDecodeError) as error:
|
||||
raise TokenExchangeError("worker key exchange request failed") from error
|
||||
|
||||
token = data.get("token")
|
||||
if not isinstance(token, str) or not token:
|
||||
raise TokenExchangeError("worker key exchange response is missing a token")
|
||||
|
||||
expires_in = data.get("expires_in")
|
||||
ttl = float(expires_in) if isinstance(expires_in, (int, float)) and expires_in > 0 else 0.0
|
||||
self._token = token
|
||||
# Renew once ~(1 - leeway) of the lifetime is gone. An unknown TTL falls
|
||||
# back to re-exchanging on the next call — correct, just chattier.
|
||||
self._refresh_at = self._now() + ttl * (1.0 - self._leeway)
|
||||
|
||||
|
||||
def provider_from_config(
|
||||
*,
|
||||
worker_key: str | None,
|
||||
userservice_url: str | None,
|
||||
bearer_token: str | None,
|
||||
request_timeout: float,
|
||||
) -> TokenProvider:
|
||||
"""Pick the token strategy: a worker key (exchange mode) wins over a static
|
||||
bearer token, which in turn wins over no credential at all."""
|
||||
if worker_key and userservice_url:
|
||||
return WorkerKeyTokenProvider(userservice_url, worker_key, request_timeout)
|
||||
return StaticTokenProvider(bearer_token)
|
||||
@@ -1,106 +0,0 @@
|
||||
"""Console entry point for ``scimesh-worker``."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
from pathlib import Path
|
||||
|
||||
from .artifacts import HttpArtifactClient
|
||||
from .auth import provider_from_config
|
||||
from .config import WorkerConfig
|
||||
from .coordinator import HttpCoordinatorClient
|
||||
from .daemon import WorkerDaemon
|
||||
from .runners import SciMeshRunner
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
"""Build the worker CLI parser for command-line use and focused tests."""
|
||||
parser = argparse.ArgumentParser(
|
||||
prog="scimesh-worker",
|
||||
epilog=(
|
||||
"Environment: SCIMESH_COORDINATOR_URL, SCIMESH_WORK_DIR, "
|
||||
"SCIMESH_WORKER_NAME, SCIMESH_CPU_COUNT, SCIMESH_MEMORY_MB, "
|
||||
"SCIMESH_POLL_INTERVAL, SCIMESH_REQUEST_TIMEOUT, "
|
||||
"SCIMESH_HEARTBEAT_INTERVAL, SCIMESH_CLEANUP_AFTER_SECONDS, "
|
||||
"SCIMESH_MAX_TASKS, SCIMESH_BEARER_TOKEN, SCIMESH_WORKER_KEY, and "
|
||||
"SCIMESH_USERSERVICE_URL. SCIMESH_WORKER_ID is a legacy/test override."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--coordinator-url")
|
||||
parser.add_argument("--worker-id")
|
||||
parser.add_argument("--work-dir")
|
||||
parser.add_argument("--worker-name")
|
||||
parser.add_argument(
|
||||
"--worker-key",
|
||||
help="Long-lived worker key from the web UI; the worker exchanges it for "
|
||||
"short-lived tokens, binding it to your account. Requires --userservice-url.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--userservice-url",
|
||||
help="Base URL of the userservice that issues tokens for --worker-key",
|
||||
)
|
||||
parser.add_argument("--cpu-count", type=int)
|
||||
parser.add_argument("--memory-mb", type=int)
|
||||
parser.add_argument("--poll-interval", type=float)
|
||||
parser.add_argument("--request-timeout", type=float)
|
||||
parser.add_argument("--heartbeat-interval", type=float)
|
||||
parser.add_argument("--cleanup-after-seconds", type=float)
|
||||
lifecycle = parser.add_mutually_exclusive_group()
|
||||
lifecycle.add_argument(
|
||||
"--once",
|
||||
action="store_true",
|
||||
help="Claim at most one task, then exit; exit immediately when the queue is empty",
|
||||
)
|
||||
lifecycle.add_argument(
|
||||
"--max-tasks",
|
||||
type=int,
|
||||
help="Process this many claimed tasks, then exit",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args(argv)
|
||||
overrides = {
|
||||
key: value
|
||||
for key, value in vars(args).items()
|
||||
if value is not None and key != "once"
|
||||
}
|
||||
if args.once:
|
||||
overrides["max_tasks"] = 1
|
||||
overrides["exit_when_idle"] = True
|
||||
if "work_dir" in overrides:
|
||||
overrides["work_dir"] = Path(overrides["work_dir"])
|
||||
try:
|
||||
config = WorkerConfig.from_environment(overrides)
|
||||
except (TypeError, ValueError) as error:
|
||||
parser.error(str(error))
|
||||
logging.basicConfig(
|
||||
level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
|
||||
)
|
||||
# One shared token strategy backs both clients: a worker key (exchanged and
|
||||
# refreshed) or a static bearer token, decided by what the config carries.
|
||||
tokens = provider_from_config(
|
||||
worker_key=config.worker_key,
|
||||
userservice_url=config.userservice_url,
|
||||
bearer_token=config.bearer_token,
|
||||
request_timeout=config.request_timeout,
|
||||
)
|
||||
client = HttpCoordinatorClient(
|
||||
config.coordinator_url, config.request_timeout, token_provider=tokens
|
||||
)
|
||||
completed_without_interruption = WorkerDaemon(
|
||||
config,
|
||||
client,
|
||||
HttpArtifactClient(
|
||||
config.coordinator_url, config.request_timeout, token_provider=tokens
|
||||
),
|
||||
SciMeshRunner.for_worker(config),
|
||||
).run_forever()
|
||||
return 0 if completed_without_interruption else 130
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -1,221 +0,0 @@
|
||||
"""Configuration parsing for the worker command."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from math import isfinite
|
||||
from pathlib import Path
|
||||
import os
|
||||
import socket
|
||||
from typing import Mapping
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
from scimesh.sdk.registry import AllowedPackage, workload_allowlist_from_json
|
||||
|
||||
|
||||
def _clean_url(value: object | None) -> str | None:
|
||||
"""Normalise an optional URL: drop a blank one, strip a trailing slash."""
|
||||
if value is None:
|
||||
return None
|
||||
text = str(value).strip()
|
||||
return text.rstrip("/") or None
|
||||
|
||||
|
||||
def _int_value(value: object | None, name: str) -> int | None:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float, str)):
|
||||
raise ValueError(f"{name} must be a number")
|
||||
return int(value)
|
||||
|
||||
|
||||
def _float_value(value: object | None, name: str) -> float | None:
|
||||
if value is None:
|
||||
return None
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float, str)):
|
||||
raise ValueError(f"{name} must be a number")
|
||||
return float(value)
|
||||
|
||||
|
||||
def _positive_number(value: object, name: str, *, allow_zero: bool = False) -> None:
|
||||
if (
|
||||
isinstance(value, bool)
|
||||
or not isinstance(value, (int, float))
|
||||
or not isfinite(value)
|
||||
or value < 0
|
||||
or (not allow_zero and value == 0)
|
||||
):
|
||||
qualifier = "non-negative" if allow_zero else "positive"
|
||||
raise ValueError(f"{name} must be {qualifier}")
|
||||
|
||||
|
||||
def _capabilities(value: object) -> tuple[str, ...]:
|
||||
"""Parse a comma-separated capability list into unique non-empty names."""
|
||||
if value is None:
|
||||
return ("similarity-search", "similarity_search")
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
raise ValueError("capabilities must be a comma-separated list")
|
||||
names = tuple(
|
||||
dict.fromkeys(item.strip() for item in value.split(",") if item.strip())
|
||||
)
|
||||
if not names:
|
||||
raise ValueError("capabilities cannot be empty")
|
||||
return names
|
||||
|
||||
|
||||
def _workload_allowlist(value: object) -> tuple[AllowedPackage, ...]:
|
||||
return workload_allowlist_from_json(value)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class WorkerConfig:
|
||||
coordinator_url: str
|
||||
worker_id: str | None
|
||||
work_dir: Path
|
||||
worker_name: str = "scimesh-worker"
|
||||
cpu_count: int = 1
|
||||
memory_mb: int | None = None
|
||||
poll_interval: float = 2.0
|
||||
request_timeout: float = 30.0
|
||||
heartbeat_interval: float = 15.0
|
||||
bearer_token: str | None = None
|
||||
# A long-lived per-user credential. When set (with userservice_url), the
|
||||
# worker exchanges it for short-lived JWTs instead of using bearer_token,
|
||||
# binding the worker to that user's account.
|
||||
worker_key: str | None = None
|
||||
userservice_url: str | None = None
|
||||
cleanup_after_seconds: float | None = None
|
||||
max_tasks: int | None = None
|
||||
exit_when_idle: bool = False
|
||||
# Distributed similarity-graph requires triangular block-pair planning and
|
||||
# is deliberately not advertised until CTX-10. A normal worker must never
|
||||
# make a multi-shard graph job appear scientifically complete.
|
||||
# The local CLI uses hyphens; the first coordinator contract used
|
||||
# underscores, so retain the search alias during migration.
|
||||
capabilities: tuple[str, ...] = (
|
||||
"similarity-search",
|
||||
"similarity_search",
|
||||
)
|
||||
# Optional allowlist of installed SDK workload packages to execute. When
|
||||
# empty, the worker runs the built-in similarity-search only. Entries are
|
||||
# ``{distribution, name, version, digest}`` JSON objects matching the
|
||||
# installed ``scimesh.workloads`` entry points.
|
||||
workload_allowlist: tuple[AllowedPackage, ...] = ()
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
parsed = urlsplit(self.coordinator_url)
|
||||
if parsed.scheme not in {"http", "https"} or not parsed.hostname:
|
||||
raise ValueError("coordinator_url must be an absolute HTTP(S) URL")
|
||||
if not isinstance(self.worker_name, str) or not self.worker_name.strip():
|
||||
raise ValueError("worker_name must be non-empty")
|
||||
if self.userservice_url is not None:
|
||||
us = urlsplit(self.userservice_url)
|
||||
if us.scheme not in {"http", "https"} or not us.hostname:
|
||||
raise ValueError("userservice_url must be an absolute HTTP(S) URL")
|
||||
if self.worker_key is not None and not self.userservice_url:
|
||||
raise ValueError(
|
||||
"worker_key requires userservice_url (SCIMESH_USERSERVICE_URL)"
|
||||
)
|
||||
if (
|
||||
isinstance(self.cpu_count, bool)
|
||||
or not isinstance(self.cpu_count, int)
|
||||
or self.cpu_count < 1
|
||||
):
|
||||
raise ValueError("cpu_count must be positive")
|
||||
if self.worker_id is not None and not isinstance(self.worker_id, str):
|
||||
raise ValueError("worker_id must be a string when set")
|
||||
if self.memory_mb is not None and (
|
||||
isinstance(self.memory_mb, bool)
|
||||
or not isinstance(self.memory_mb, int)
|
||||
or self.memory_mb < 1
|
||||
):
|
||||
raise ValueError("memory_mb must be positive when set")
|
||||
_positive_number(self.poll_interval, "poll_interval")
|
||||
_positive_number(self.request_timeout, "request_timeout")
|
||||
_positive_number(self.heartbeat_interval, "heartbeat_interval")
|
||||
if self.cleanup_after_seconds is not None:
|
||||
_positive_number(
|
||||
self.cleanup_after_seconds, "cleanup_after_seconds", allow_zero=True
|
||||
)
|
||||
if self.max_tasks is not None:
|
||||
if (
|
||||
isinstance(self.max_tasks, bool)
|
||||
or not isinstance(self.max_tasks, int)
|
||||
or self.max_tasks < 1
|
||||
):
|
||||
raise ValueError("max_tasks must be positive when set")
|
||||
if not isinstance(self.exit_when_idle, bool):
|
||||
raise ValueError("exit_when_idle must be a boolean")
|
||||
if not self.capabilities:
|
||||
raise ValueError("capabilities cannot be empty")
|
||||
if any(
|
||||
not isinstance(capability, str) or not capability.strip()
|
||||
for capability in self.capabilities
|
||||
):
|
||||
raise ValueError("capabilities must contain non-empty names")
|
||||
if len(self.capabilities) != len(set(self.capabilities)):
|
||||
raise ValueError("capabilities must be unique")
|
||||
if any(
|
||||
not isinstance(package, AllowedPackage)
|
||||
for package in self.workload_allowlist
|
||||
):
|
||||
raise ValueError("workload_allowlist must contain AllowedPackage values")
|
||||
# Runner subprocesses use a task directory as their cwd. Keep the
|
||||
# configured root absolute so input/output paths remain valid there
|
||||
# even when the CLI received a convenient relative --work-dir value.
|
||||
object.__setattr__(self, "work_dir", self.work_dir.expanduser().resolve())
|
||||
|
||||
@classmethod
|
||||
def from_environment(
|
||||
cls, overrides: Mapping[str, object] | None = None
|
||||
) -> "WorkerConfig":
|
||||
"""Build config from environment, allowing typed CLI values to override it."""
|
||||
values = overrides or {}
|
||||
|
||||
def value(
|
||||
name: str, environment: str, default: object | None = None
|
||||
) -> object | None:
|
||||
override = values.get(name)
|
||||
return override if override is not None else os.getenv(environment, default)
|
||||
|
||||
url = value("coordinator_url", "SCIMESH_COORDINATOR_URL")
|
||||
if not isinstance(url, str) or not url:
|
||||
raise ValueError("SCIMESH_COORDINATOR_URL or --coordinator-url is required")
|
||||
cleanup = value("cleanup_after_seconds", "SCIMESH_CLEANUP_AFTER_SECONDS")
|
||||
cpu_count = value("cpu_count", "SCIMESH_CPU_COUNT", os.cpu_count() or 1)
|
||||
memory_mb = value("memory_mb", "SCIMESH_MEMORY_MB")
|
||||
max_tasks = value("max_tasks", "SCIMESH_MAX_TASKS")
|
||||
capabilities = value("capabilities", "SCIMESH_CAPABILITIES")
|
||||
allowlist = value("workload_allowlist", "SCIMESH_WORKLOAD_ALLOWLIST")
|
||||
worker_id = value("worker_id", "SCIMESH_WORKER_ID")
|
||||
work_dir = value("work_dir", "SCIMESH_WORK_DIR", "./scimesh-worker-data")
|
||||
worker_name = value("worker_name", "SCIMESH_WORKER_NAME", socket.gethostname())
|
||||
poll_interval = value("poll_interval", "SCIMESH_POLL_INTERVAL", "2")
|
||||
request_timeout = value("request_timeout", "SCIMESH_REQUEST_TIMEOUT", "30")
|
||||
heartbeat_interval = value(
|
||||
"heartbeat_interval", "SCIMESH_HEARTBEAT_INTERVAL", "15"
|
||||
)
|
||||
bearer_token = value("bearer_token", "SCIMESH_BEARER_TOKEN")
|
||||
worker_key = value("worker_key", "SCIMESH_WORKER_KEY")
|
||||
userservice_url = _clean_url(
|
||||
value("userservice_url", "SCIMESH_USERSERVICE_URL")
|
||||
)
|
||||
return cls(
|
||||
coordinator_url=url.rstrip("/"),
|
||||
worker_id=str(worker_id) if worker_id is not None else None,
|
||||
work_dir=Path(str(work_dir)),
|
||||
worker_name=str(worker_name),
|
||||
cpu_count=_int_value(cpu_count, "cpu_count") or 1,
|
||||
memory_mb=_int_value(memory_mb, "memory_mb"),
|
||||
poll_interval=_float_value(poll_interval, "poll_interval") or 2.0,
|
||||
request_timeout=_float_value(request_timeout, "request_timeout") or 30.0,
|
||||
heartbeat_interval=_float_value(heartbeat_interval, "heartbeat_interval") or 15.0,
|
||||
bearer_token=str(bearer_token) if bearer_token is not None else None,
|
||||
worker_key=str(worker_key) if worker_key is not None else None,
|
||||
userservice_url=userservice_url,
|
||||
cleanup_after_seconds=_float_value(cleanup, "cleanup_after_seconds"),
|
||||
max_tasks=_int_value(max_tasks, "max_tasks"),
|
||||
exit_when_idle=bool(values.get("exit_when_idle", False)),
|
||||
capabilities=_capabilities(capabilities),
|
||||
workload_allowlist=_workload_allowlist(allowlist),
|
||||
)
|
||||
@@ -1,149 +0,0 @@
|
||||
"""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, build_opener
|
||||
|
||||
from .auth import StaticTokenProvider, TokenProvider
|
||||
from .models import ClaimedTask, RegisteredWorker
|
||||
from .transport import NoRedirectHandler
|
||||
|
||||
|
||||
class CoordinatorError(RuntimeError):
|
||||
"""A non-retriable coordinator response."""
|
||||
|
||||
|
||||
class CoordinatorTransientError(CoordinatorError):
|
||||
"""A timeout, connection error, or 5xx coordinator response."""
|
||||
|
||||
|
||||
class CoordinatorConflictError(CoordinatorError):
|
||||
"""The worker no longer owns the task lease or attempted a conflicting mutation."""
|
||||
|
||||
|
||||
class CoordinatorClient(Protocol):
|
||||
def register(
|
||||
self, name: str, capabilities: tuple[str, ...], cpu_count: int, memory_mb: int | None
|
||||
) -> RegisteredWorker: ...
|
||||
|
||||
def claim(self, worker_id: str, capabilities: tuple[str, ...]) -> ClaimedTask | None: ...
|
||||
|
||||
def submit(self, task: ClaimedTask, payload: dict[str, Any]) -> None: ...
|
||||
|
||||
def fail(self, task: ClaimedTask, payload: dict[str, Any]) -> None: ...
|
||||
|
||||
def heartbeat(self, task: ClaimedTask, worker_id: str) -> str: ...
|
||||
|
||||
|
||||
class HttpCoordinatorClient:
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
timeout: float,
|
||||
bearer_token: str | None = None,
|
||||
*,
|
||||
token_provider: TokenProvider | None = None,
|
||||
) -> None:
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.timeout = timeout
|
||||
# A bearer_token argument keeps older call sites working; internally
|
||||
# everything goes through a provider so refresh is uniform.
|
||||
self._tokens: TokenProvider = token_provider or StaticTokenProvider(bearer_token)
|
||||
self._opener = build_opener(NoRedirectHandler())
|
||||
|
||||
@property
|
||||
def bearer_token(self) -> str | None:
|
||||
return self._tokens.token()
|
||||
|
||||
def register(
|
||||
self, name: str, capabilities: tuple[str, ...], cpu_count: int, memory_mb: int | None
|
||||
) -> RegisteredWorker:
|
||||
payload: dict[str, Any] = {
|
||||
"name": name,
|
||||
"capabilities": list(capabilities),
|
||||
"cpu_count": cpu_count,
|
||||
}
|
||||
if memory_mb is not None:
|
||||
payload["memory_mb"] = memory_mb
|
||||
status, body = self._request("POST", "/workers/register", payload)
|
||||
if status != 201:
|
||||
raise CoordinatorError(f"worker registration rejected with status {status}")
|
||||
try:
|
||||
return RegisteredWorker.from_json(body)
|
||||
except ValueError as error:
|
||||
raise CoordinatorError("invalid worker registration response") from error
|
||||
|
||||
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):
|
||||
if status == 409:
|
||||
raise CoordinatorConflictError("result rejected because the task lease was lost")
|
||||
raise CoordinatorError(f"result rejected with status {status}")
|
||||
|
||||
def fail(self, task: ClaimedTask, payload: dict[str, Any]) -> None:
|
||||
status, _ = self._request("POST", f"/tasks/{task.task_id}/failure", payload)
|
||||
if status not in (200, 201, 202):
|
||||
if status == 409:
|
||||
raise CoordinatorConflictError("failure rejected because the task lease was lost")
|
||||
raise CoordinatorError(f"failure report rejected with status {status}")
|
||||
|
||||
def heartbeat(self, task: ClaimedTask, worker_id: str) -> str:
|
||||
status, body = self._request(
|
||||
"POST", f"/tasks/{task.task_id}/heartbeat",
|
||||
{"worker_id": worker_id, "attempt": task.attempt},
|
||||
)
|
||||
if status != 200:
|
||||
if status == 409:
|
||||
raise CoordinatorConflictError("heartbeat rejected because the task lease was lost")
|
||||
raise CoordinatorError(f"heartbeat rejected with status {status}")
|
||||
lease_expires_at = body.get("lease_expires_at")
|
||||
if not isinstance(lease_expires_at, str):
|
||||
raise CoordinatorError("heartbeat response is missing lease_expires_at")
|
||||
return lease_expires_at
|
||||
|
||||
def _request(self, method: str, path: str, payload: dict[str, Any]) -> tuple[int, dict[str, Any]]:
|
||||
return self._request_once(method, path, payload, allow_refresh=True)
|
||||
|
||||
def _request_once(
|
||||
self, method: str, path: str, payload: dict[str, Any], *, allow_refresh: bool
|
||||
) -> 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 self._opener.open(request, timeout=self.timeout) as response:
|
||||
raw = response.read()
|
||||
try:
|
||||
return response.status, json.loads(raw) if raw else {}
|
||||
except json.JSONDecodeError as error:
|
||||
raise CoordinatorError("coordinator returned invalid JSON") from error
|
||||
except HTTPError as error:
|
||||
# A 401 usually means the short-lived JWT expired; mint a fresh one
|
||||
# and retry exactly once so an in-flight worker rides over the gap.
|
||||
if error.code == 401 and allow_refresh:
|
||||
self._tokens.refresh()
|
||||
return self._request_once(method, path, payload, allow_refresh=False)
|
||||
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]:
|
||||
token = self._tokens.token()
|
||||
return {"Authorization": f"Bearer {token}"} if token else {}
|
||||
@@ -1,331 +0,0 @@
|
||||
"""The worker state machine and its safe failure handling."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import replace
|
||||
from dataclasses import dataclass
|
||||
import random
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import threading
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from .artifacts import ArtifactClient, sha256_file
|
||||
from .config import WorkerConfig
|
||||
from .coordinator import (
|
||||
CoordinatorClient,
|
||||
CoordinatorConflictError,
|
||||
CoordinatorTransientError,
|
||||
)
|
||||
from .models import ClaimedTask, UploadedArtifact
|
||||
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._worker_id = config.worker_id or ""
|
||||
self._stop = threading.Event()
|
||||
self._error: Exception | None = None
|
||||
self._thread: threading.Thread | None = None
|
||||
self._lease_expires_at = task.lease_expires_at
|
||||
|
||||
def start(self) -> None:
|
||||
# Verify ownership before expensive download or calculation begins.
|
||||
self._lease_expires_at = self.coordinator.heartbeat(self.task, self._worker_id)
|
||||
self._next_delay()
|
||||
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 = self._next_delay()
|
||||
while not self._stop.wait(max(delay, 0.01)):
|
||||
try:
|
||||
self._lease_expires_at = self.coordinator.heartbeat(
|
||||
self.task, self._worker_id
|
||||
)
|
||||
delay = self._next_delay()
|
||||
except (
|
||||
Exception
|
||||
) as error: # Surface the lease loss in the main state machine.
|
||||
self._error = error
|
||||
return
|
||||
|
||||
def _next_delay(self) -> float:
|
||||
return min(self.config.heartbeat_interval, self._seconds_until_expiry() / 2)
|
||||
|
||||
def _seconds_until_expiry(self) -> float:
|
||||
try:
|
||||
expiry = datetime.fromisoformat(
|
||||
self._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
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RunOnceOutcome:
|
||||
"""Whether a claim was made and whether that claimed task completed."""
|
||||
|
||||
claimed: bool
|
||||
completed: bool
|
||||
|
||||
|
||||
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.worker_id = config.worker_id
|
||||
self._registered = False
|
||||
self.log = logging.getLogger("scimesh.worker")
|
||||
|
||||
def run_forever(self) -> bool:
|
||||
"""Run until stopped; return false only when interrupted by the operator."""
|
||||
failures = 0
|
||||
completed_tasks = 0
|
||||
self._log(
|
||||
"started",
|
||||
max_tasks=self.config.max_tasks,
|
||||
exit_when_idle=self.config.exit_when_idle,
|
||||
)
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
if not self._registered:
|
||||
self._register_worker()
|
||||
self._cleanup_expired_directories()
|
||||
outcome = self.run_once()
|
||||
failures = 0
|
||||
if outcome.claimed:
|
||||
if outcome.completed:
|
||||
completed_tasks += 1
|
||||
if self.config.exit_when_idle:
|
||||
self._log(
|
||||
"stopped",
|
||||
reason="one_claim_processed",
|
||||
completed_tasks=completed_tasks,
|
||||
)
|
||||
return True
|
||||
if (
|
||||
outcome.completed
|
||||
and self.config.max_tasks is not None
|
||||
and completed_tasks >= self.config.max_tasks
|
||||
):
|
||||
self._log(
|
||||
"stopped",
|
||||
reason="max_tasks_reached",
|
||||
completed_tasks=completed_tasks,
|
||||
)
|
||||
return True
|
||||
elif self.config.exit_when_idle:
|
||||
self._log(
|
||||
"stopped",
|
||||
reason="queue_empty",
|
||||
completed_tasks=completed_tasks,
|
||||
)
|
||||
return True
|
||||
else:
|
||||
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)
|
||||
)
|
||||
except KeyboardInterrupt:
|
||||
self._log("stopped", reason="interrupted", completed_tasks=completed_tasks)
|
||||
return False
|
||||
|
||||
def run_once(self) -> RunOnceOutcome:
|
||||
worker_id = self._worker_id()
|
||||
self._log("claiming", log_level=logging.DEBUG)
|
||||
task = self.coordinator.claim(worker_id, self.config.capabilities)
|
||||
if task is None:
|
||||
self._log("idle", log_level=logging.DEBUG)
|
||||
return RunOnceOutcome(claimed=False, completed=False)
|
||||
started = time.monotonic()
|
||||
task_dir = self.config.work_dir / task.task_id / str(task.attempt)
|
||||
heartbeat = LeaseHeartbeat(task, self.coordinator, self.config)
|
||||
completed = False
|
||||
try:
|
||||
task_dir.mkdir(parents=True, exist_ok=False)
|
||||
heartbeat.start()
|
||||
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)
|
||||
heartbeat.raise_if_failed()
|
||||
if len(result.artifacts) != 1:
|
||||
raise ValueError("runner must produce exactly one result artifact")
|
||||
artifact = result.artifacts[0]
|
||||
uploaded = self.artifacts.upload(task, worker_id, artifact)
|
||||
manifest = self._result_manifest(uploaded)
|
||||
self._log("submitting", task)
|
||||
heartbeat.raise_if_failed()
|
||||
self.coordinator.submit(
|
||||
task,
|
||||
{
|
||||
"worker_id": worker_id,
|
||||
"attempt": task.attempt,
|
||||
"result": manifest,
|
||||
"metrics": {
|
||||
**result.metrics,
|
||||
"elapsed_seconds": round(time.monotonic() - started, 3),
|
||||
},
|
||||
},
|
||||
)
|
||||
completed = True
|
||||
self._log(
|
||||
"completed", task, elapsed_seconds=round(time.monotonic() - started, 3)
|
||||
)
|
||||
except KeyboardInterrupt:
|
||||
self._log("interrupted", task)
|
||||
try:
|
||||
self._report_failure(
|
||||
task, InterruptedError("worker interrupted by operator")
|
||||
)
|
||||
except CoordinatorTransientError:
|
||||
self._log("failed", task, error_type="FailureReportError")
|
||||
raise
|
||||
except CoordinatorConflictError as error:
|
||||
self._log("lease_lost", task, error_type=type(error).__name__)
|
||||
except Exception as error:
|
||||
self._log("failed", task, error_type=type(error).__name__)
|
||||
self._report_failure(task, error)
|
||||
finally:
|
||||
heartbeat.stop()
|
||||
return RunOnceOutcome(claimed=True, completed=completed)
|
||||
|
||||
def _report_failure(self, task: ClaimedTask, error: Exception) -> None:
|
||||
message = self._sanitize_error_message(error)
|
||||
try:
|
||||
self.coordinator.fail(
|
||||
task,
|
||||
{
|
||||
"worker_id": self._worker_id(),
|
||||
"attempt": task.attempt,
|
||||
"error_code": type(error).__name__,
|
||||
"error_message": message,
|
||||
"retryable": self._is_retryable(error),
|
||||
},
|
||||
)
|
||||
except CoordinatorTransientError:
|
||||
raise
|
||||
except Exception:
|
||||
self._log("failed", task, error_type="FailureReportError")
|
||||
|
||||
@staticmethod
|
||||
def _is_retryable(error: Exception) -> bool:
|
||||
"""Retry transient worker/transport failures, never invalid scientific input."""
|
||||
return not isinstance(
|
||||
error, (ValueError, FileNotFoundError, subprocess.CalledProcessError)
|
||||
)
|
||||
|
||||
def _sanitize_error_message(self, error: Exception) -> str:
|
||||
"""Keep coordinator-visible failures useful without exposing local paths."""
|
||||
message = str(error).replace(str(self.config.work_dir), "<worker-dir>")
|
||||
# CalledProcessError includes the complete argv, including sys.executable
|
||||
# outside work_dir. Replace POSIX and Windows absolute paths before the
|
||||
# message reaches the coordinator database or operator UI.
|
||||
message = re.sub(r"(?<![\w:])[A-Za-z]:\\[^\s'\"\],)]+", "<path>", message)
|
||||
message = re.sub(r"(?<![\w:])/(?:[^\s'\"\],)]+)", "<path>", message)
|
||||
return message[:300]
|
||||
|
||||
def _register_worker(self) -> None:
|
||||
registered = self.coordinator.register(
|
||||
self.config.worker_name,
|
||||
self.config.capabilities,
|
||||
self.config.cpu_count,
|
||||
self.config.memory_mb,
|
||||
)
|
||||
self.worker_id = registered.worker_id
|
||||
self.config = replace(
|
||||
self.config,
|
||||
worker_id=registered.worker_id,
|
||||
heartbeat_interval=registered.heartbeat_interval_seconds,
|
||||
)
|
||||
self._registered = True
|
||||
self._log("registered")
|
||||
|
||||
def _worker_id(self) -> str:
|
||||
if not self.worker_id:
|
||||
raise ValueError("worker is not registered")
|
||||
return self.worker_id
|
||||
|
||||
@staticmethod
|
||||
def _result_manifest(uploaded: UploadedArtifact) -> dict[str, object]:
|
||||
"""Keep completion payload exact: coordinator owns all artifact metadata."""
|
||||
return {"artifact_id": uploaded.artifact_id}
|
||||
|
||||
def _log(
|
||||
self,
|
||||
state: str,
|
||||
task: ClaimedTask | None = None,
|
||||
*,
|
||||
log_level: int = logging.INFO,
|
||||
**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.log(log_level, "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))
|
||||
+20
-61
@@ -1,10 +1,13 @@
|
||||
"""Value objects shared by the worker daemon components."""
|
||||
"""Value objects for one claimed task and its run result.
|
||||
|
||||
These are the wire types the Go worker agent hands to the per-task Python
|
||||
entry point; the daemon lifecycle itself lives in the Go agent.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from math import isfinite
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urlsplit
|
||||
@@ -34,7 +37,9 @@ def _coordinator_uri(value: object, field: str) -> str:
|
||||
|
||||
def _sha256(value: object, field: str) -> str:
|
||||
digest = _required_string(value, field).lower()
|
||||
if len(digest) != 64 or any(character not in "0123456789abcdef" for character in digest):
|
||||
if len(digest) != 64 or any(
|
||||
character not in "0123456789abcdef" for character in digest
|
||||
):
|
||||
raise ValueError(f"{field} must be a SHA-256 hex digest")
|
||||
return digest
|
||||
|
||||
@@ -61,11 +66,20 @@ class ClaimedTask:
|
||||
if not isinstance(input_data, dict):
|
||||
raise ValueError("input must be an object")
|
||||
raw_attempt = data["attempt"]
|
||||
if isinstance(raw_attempt, bool) or not isinstance(raw_attempt, int) or raw_attempt < 1:
|
||||
if (
|
||||
isinstance(raw_attempt, bool)
|
||||
or not isinstance(raw_attempt, int)
|
||||
or raw_attempt < 1
|
||||
):
|
||||
raise ValueError("attempt must be a positive integer")
|
||||
task_id = str(UUID(_required_string(data["task_id"], "task_id")))
|
||||
lease_expires_at = _required_string(data["lease_expires_at"], "lease_expires_at")
|
||||
if datetime.fromisoformat(lease_expires_at.replace("Z", "+00:00")).tzinfo is None:
|
||||
lease_expires_at = _required_string(
|
||||
data["lease_expires_at"], "lease_expires_at"
|
||||
)
|
||||
if (
|
||||
datetime.fromisoformat(lease_expires_at.replace("Z", "+00:00")).tzinfo
|
||||
is None
|
||||
):
|
||||
raise ValueError("lease_expires_at must include a timezone")
|
||||
parameters = data.get("parameters", {})
|
||||
if not isinstance(parameters, dict):
|
||||
@@ -91,61 +105,6 @@ class ProducedArtifact:
|
||||
content_type: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class UploadedArtifact:
|
||||
"""Coordinator-owned artifact metadata returned after a successful upload."""
|
||||
|
||||
artifact_id: str
|
||||
uri: str
|
||||
sha256: str
|
||||
size_bytes: int
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, data: object) -> "UploadedArtifact":
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("artifact upload response must be an object")
|
||||
raw_size = data.get("size_bytes")
|
||||
if isinstance(raw_size, bool) or not isinstance(raw_size, int) or raw_size < 0:
|
||||
raise ValueError("artifact size_bytes must be a non-negative integer")
|
||||
try:
|
||||
return cls(
|
||||
artifact_id=str(UUID(_required_string(data.get("artifact_id"), "artifact_id"))),
|
||||
uri=_coordinator_uri(data.get("uri"), "uri"),
|
||||
sha256=_sha256(data.get("sha256"), "sha256"),
|
||||
size_bytes=raw_size,
|
||||
)
|
||||
except ValueError as error:
|
||||
raise ValueError("invalid artifact upload response") from error
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RegisteredWorker:
|
||||
"""Identity and heartbeat policy returned by worker registration."""
|
||||
|
||||
worker_id: str
|
||||
heartbeat_interval_seconds: float
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, data: object) -> "RegisteredWorker":
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("worker registration response must be an object")
|
||||
raw_interval = data.get("heartbeat_interval_seconds")
|
||||
if (
|
||||
isinstance(raw_interval, bool)
|
||||
or not isinstance(raw_interval, (int, float))
|
||||
or not isfinite(raw_interval)
|
||||
or raw_interval <= 0
|
||||
):
|
||||
raise ValueError("heartbeat_interval_seconds must be positive")
|
||||
try:
|
||||
return cls(
|
||||
worker_id=str(UUID(_required_string(data.get("worker_id"), "worker_id"))),
|
||||
heartbeat_interval_seconds=float(raw_interval),
|
||||
)
|
||||
except ValueError as error:
|
||||
raise ValueError("invalid worker registration response") from error
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RunResult:
|
||||
artifacts: tuple[ProducedArtifact, ...]
|
||||
|
||||
+42
-38
@@ -46,9 +46,9 @@ from scimesh.sdk.runtime import (
|
||||
)
|
||||
from scimesh.sdk.workflow import StageKind
|
||||
|
||||
from .config import WorkerConfig
|
||||
from .models import ClaimedTask, ProducedArtifact, RunResult
|
||||
|
||||
|
||||
def _utc_now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
|
||||
|
||||
@@ -57,6 +57,38 @@ class Runner(Protocol):
|
||||
def run(self, task: ClaimedTask, task_dir: Path) -> RunResult: ...
|
||||
|
||||
|
||||
def _definitions_from_environment() -> dict[str, WorkloadDefinition]:
|
||||
"""Load definitions from ``SCIMESH_WORKLOAD_ALLOWLIST`` or the built-ins.
|
||||
|
||||
The environment-driven discovery mirrors the former worker configuration;
|
||||
the Go agent passes the allowlist through unchanged.
|
||||
"""
|
||||
import os
|
||||
|
||||
from scimesh.sdk.registry import workload_allowlist_from_json
|
||||
|
||||
allowlist = workload_allowlist_from_json(os.getenv("SCIMESH_WORKLOAD_ALLOWLIST"))
|
||||
if allowlist:
|
||||
registry = WorkloadRegistry()
|
||||
registry.discover_installed(allowlist)
|
||||
definitions: dict[str, WorkloadDefinition] = {}
|
||||
for description in registry.descriptions():
|
||||
definition, _ = registry.require(
|
||||
description.workload.name,
|
||||
description.workload.version,
|
||||
description.package_digest,
|
||||
)
|
||||
definitions[description.workload.name] = definition
|
||||
if not definitions:
|
||||
raise ValueError("workload_allowlist discovered no workloads")
|
||||
return definitions
|
||||
from scimesh.workloads.search import similarity_search_sdk_definition
|
||||
|
||||
return {
|
||||
"similarity-search": similarity_search_sdk_definition().definition(),
|
||||
}
|
||||
|
||||
|
||||
def _inventory_for(
|
||||
definitions: Mapping[str, WorkloadDefinition],
|
||||
*,
|
||||
@@ -91,7 +123,12 @@ def _runtime_for(
|
||||
|
||||
|
||||
class SciMeshRunner:
|
||||
"""Execute claimed coordinator tasks through SDK-built workloads."""
|
||||
"""Execute claimed coordinator tasks through SDK-built workloads.
|
||||
|
||||
Definitions come from ``SCIMESH_WORKLOAD_ALLOWLIST`` when set, otherwise
|
||||
the built-in ``similarity-search``; callers may supply an explicit mapping
|
||||
for tests.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -100,13 +137,9 @@ class SciMeshRunner:
|
||||
inventory: ResourceInventory | None = None,
|
||||
runtime: RuntimeCapabilities | None = None,
|
||||
) -> None:
|
||||
self._definitions = dict(definitions or {})
|
||||
if "similarity-search" not in self._definitions:
|
||||
from scimesh.workloads.search import similarity_search_sdk_definition
|
||||
|
||||
self._definitions["similarity-search"] = (
|
||||
similarity_search_sdk_definition().definition()
|
||||
)
|
||||
if definitions is None:
|
||||
definitions = _definitions_from_environment()
|
||||
self._definitions = dict(definitions)
|
||||
self._inventory = inventory or _inventory_for(
|
||||
self._definitions,
|
||||
cpu_cores=1,
|
||||
@@ -115,35 +148,6 @@ class SciMeshRunner:
|
||||
self._runtime = runtime or _runtime_for(self._definitions, self._inventory)
|
||||
self._pool = ResourcePool(self._runtime.inventory, max_concurrency=1)
|
||||
|
||||
@classmethod
|
||||
def for_worker(cls, config: WorkerConfig) -> "SciMeshRunner":
|
||||
"""Build a runner for one worker: discover allowlisted workloads or use built-ins."""
|
||||
definitions: dict[str, WorkloadDefinition] = {}
|
||||
if config.workload_allowlist:
|
||||
registry = WorkloadRegistry()
|
||||
registry.discover_installed(config.workload_allowlist)
|
||||
for description in registry.descriptions():
|
||||
definition, _ = registry.require(
|
||||
description.workload.name,
|
||||
description.workload.version,
|
||||
description.package_digest,
|
||||
)
|
||||
definitions[description.workload.name] = definition
|
||||
if not definitions:
|
||||
raise ValueError("workload_allowlist discovered no workloads")
|
||||
else:
|
||||
from scimesh.workloads.search import similarity_search_sdk_definition
|
||||
|
||||
definitions["similarity-search"] = (
|
||||
similarity_search_sdk_definition().definition()
|
||||
)
|
||||
inventory = _inventory_for(
|
||||
definitions,
|
||||
cpu_cores=config.cpu_count,
|
||||
memory_mb=config.memory_mb or 1024,
|
||||
)
|
||||
return cls(definitions=definitions, inventory=inventory)
|
||||
|
||||
def run(self, task: ClaimedTask, task_dir: Path) -> RunResult:
|
||||
task_dir = task_dir.resolve()
|
||||
workload = task.workload.replace("_", "-")
|
||||
|
||||
@@ -1,53 +0,0 @@
|
||||
"""Small HTTP transport helpers shared by coordinator and artifact clients."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from http.client import HTTPMessage
|
||||
from typing import IO
|
||||
from urllib.request import HTTPRedirectHandler, Request
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
|
||||
def origin(uri: str) -> tuple[str, str, int | None]:
|
||||
"""Return a normalized HTTP origin for authorization decisions."""
|
||||
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):
|
||||
"""Strip coordinator authorization when an artifact redirect 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: IO[bytes],
|
||||
code: int,
|
||||
msg: str,
|
||||
headers: HTTPMessage,
|
||||
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 NoRedirectHandler(HTTPRedirectHandler):
|
||||
"""Reject redirects for mutating coordinator API calls."""
|
||||
|
||||
def redirect_request(
|
||||
self,
|
||||
req: Request,
|
||||
fp: object,
|
||||
code: int,
|
||||
msg: str,
|
||||
headers: object,
|
||||
newurl: str,
|
||||
) -> Request | None:
|
||||
return None
|
||||
Reference in New Issue
Block a user