69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
"""Worker registry abstraction and read helpers."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
|
|
from .capabilities import has_capabilities
|
|
from .errors import NotFoundError
|
|
from .models import WorkerRegistration
|
|
from .store import TaskPlaneStore
|
|
|
|
|
|
class WorkerRegistry:
|
|
def __init__(self, store: TaskPlaneStore):
|
|
self.store = store
|
|
|
|
def register(
|
|
self,
|
|
payload: dict,
|
|
*,
|
|
actor_identity: str,
|
|
is_admin: bool,
|
|
expected_revision: int | None,
|
|
) -> WorkerRegistration:
|
|
reg = WorkerRegistration.from_dict(payload)
|
|
return self.store.register_worker(
|
|
reg,
|
|
actor_identity=actor_identity,
|
|
is_admin=is_admin,
|
|
expected_revision=expected_revision,
|
|
)
|
|
|
|
def heartbeat(
|
|
self,
|
|
worker_id: str,
|
|
heartbeat_interval_seconds: int | None = None,
|
|
*,
|
|
actor_identity: str,
|
|
expected_revision: int | None,
|
|
) -> WorkerRegistration:
|
|
return self.store.heartbeat_worker(
|
|
worker_id,
|
|
heartbeat_interval_seconds=heartbeat_interval_seconds,
|
|
actor_identity=actor_identity,
|
|
expected_revision=expected_revision,
|
|
)
|
|
|
|
def list(self):
|
|
return self.store.list_workers()
|
|
|
|
def get(self, worker_id: str) -> WorkerRegistration:
|
|
try:
|
|
return self.store.get_worker(worker_id)
|
|
except NotFoundError:
|
|
raise
|
|
|
|
def can_execute(self, worker_id: str, required_capabilities: list[str]) -> bool:
|
|
worker = self.store.get_worker(worker_id)
|
|
return has_capabilities(worker.capabilities, required_capabilities)
|
|
|
|
def is_online(self, worker_id: str, *, now: float | None = None) -> bool:
|
|
worker = self.store.get_worker(worker_id)
|
|
deadline = worker.last_heartbeat + (worker.heartbeat_interval_seconds * 3)
|
|
return not worker.disabled and (time.time() if now is None else now) <= deadline
|
|
|
|
def stale_workers(self, *, now: float | None = None) -> list[WorkerRegistration]:
|
|
current = time.time() if now is None else now
|
|
return [worker for worker in self.list() if not self.is_online(worker.worker_id, now=current)]
|