"""One worker, several servers, one queue. A volunteer with a spare machine has no reason to offer it to exactly one server. They might run one for themselves, join a friend's, and join a public one, and none of those should require a second copy of the worker, a second model download or a second slice of the machine. **Why this is a transport and not a second loop.** The obvious design is one :class:`~distinct_agent.worker.WorkerLoop` per server, and it is wrong. Each loop would own its own queue, so a machine advertising "two free slots" to three servers would be advertising six, and three servers would each fill what they believed was an empty queue. Capacity is a property of the machine, so there is one queue, one loop and one capacity figure; what fans out is the conversation with each server. So this implements the same :class:`~distinct_agent.worker.AgentTransport` interface the loop already talks to, and multiplexes underneath. The loop is unchanged and does not know how many servers exist. **Routing.** A job arrives from one server and its result belongs to that server and to no other. Every offer accepted is recorded against the transport that supplied it, and status, results and acknowledgements follow that record. A result must never be published to a server that did not ask for it: it would be handing one server's user's answer to another server's operator. **Failure.** One unreachable server must not stop the others. A transport that raises is put aside for a back-off interval and retried; the worker keeps serving everyone else meanwhile, and says so once rather than on every poll. """ from __future__ import annotations import threading import time from collections.abc import Callable, Iterable, Sequence from dataclasses import dataclass, field from typing import Any from distinct_protocol import AgentSnapshot, JobResult, JobSpec, JobStatus #: How long a server that raised is left alone before it is tried again, and #: the ceiling that back-off climbs to. Short enough that a server restarting #: is picked up quickly, long enough that a server that is simply gone does #: not dominate every poll cycle. BACKOFF_START_SECONDS = 5.0 BACKOFF_CEILING_SECONDS = 120.0 @dataclass class ServerLink: """One paired server, and what the worker knows about its health.""" name: str url: str transport: Any failures: int = 0 resume_at: float = 0.0 last_error: str = "" #: Jobs currently owned by this server. Kept so a result can be routed #: back, and so a link's outstanding work can be described. job_ids: set[str] = field(default_factory=set) def healthy(self, now: float) -> bool: return now >= self.resume_at class MultiServerTransport: """Poll several servers in turn; route each job's traffic back to its own. Thread safety matters here and is not incidental: the poll thread calls :meth:`poll_offer` while worker threads call :meth:`publish_result` for jobs that arrived several cycles ago. The routing table is guarded, and the underlying transports are assumed to be individually safe in the same way the single-server case already assumes. """ def __init__( self, links: Sequence[ServerLink], *, monotonic: Callable[[], float] = time.monotonic, report: Callable[[str], None] | None = None, ) -> None: if not links: raise ValueError("a worker needs at least one server") self._links = list(links) self._monotonic = monotonic self._report = report or (lambda message: None) self._lock = threading.RLock() self._owner: dict[str, ServerLink] = {} self._cursor = 0 # Introspection ------------------------------------------------------ @property def links(self) -> tuple[ServerLink, ...]: with self._lock: return tuple(self._links) def describe(self) -> str: """One line per server, for the operator's start-up banner.""" now = self._monotonic() lines = [] for link in self.links: if link.healthy(now): state = f"{len(link.job_ids)} running" if link.job_ids else "idle" else: state = f"unreachable, retrying ({link.last_error[:60]})" lines.append(f" {link.name}: {link.url} [{state}]") return "\n".join(lines) # Health ------------------------------------------------------------- def _fail(self, link: ServerLink, exc: BaseException) -> None: with self._lock: first = link.failures == 0 link.failures += 1 delay = min( BACKOFF_CEILING_SECONDS, BACKOFF_START_SECONDS * (2 ** (link.failures - 1)) ) link.resume_at = self._monotonic() + delay link.last_error = f"{type(exc).__name__}: {exc}" if first: # Once, not on every cycle. A server that is down for an hour # should not produce an hour of identical lines. self._report( f"{link.name} is not responding ({link.last_error[:120]}). " f"Still serving the others; retrying in {delay:.0f}s." ) def _recover(self, link: ServerLink) -> None: with self._lock: if not link.failures: return link.failures = 0 link.resume_at = 0.0 link.last_error = "" self._report(f"{link.name} is responding again.") # AgentTransport ----------------------------------------------------- def poll_offer(self, snapshot: AgentSnapshot) -> JobSpec | None: """Ask each healthy server in turn, starting where the last poll left off. Round-robin rather than always-first-server, so a busy server cannot starve the others of this worker's attention. The first offer wins and the cursor advances past it, which is the fairness this needs: no server is asked twice before every other has been asked once. """ now = self._monotonic() with self._lock: order = self._links[self._cursor :] + self._links[: self._cursor] for index, link in enumerate(order): if not link.healthy(now): continue try: offer = link.transport.poll_offer(snapshot) except Exception as exc: # noqa: BLE001 - one server must not stop the rest self._fail(link, exc) continue self._recover(link) if offer is None: continue with self._lock: self._owner[offer.id] = link link.job_ids.add(offer.id) self._cursor = (self._cursor + index + 1) % len(self._links) return offer with self._lock: if self._links: self._cursor = (self._cursor + 1) % len(self._links) return None def poll_cancellations(self) -> Iterable[str]: """Every server's cancellations, together. A cancellation names a job id, and job ids are unique across servers because the server that issued one is the only one that knows it. So the union is safe: an id from server A cannot collide with an id from server B, and if it somehow did, the queue would cancel a job that server A had every right to cancel anyway. """ now = self._monotonic() cancelled: list[str] = [] for link in self.links: if not link.healthy(now): continue try: cancelled.extend(str(job_id) for job_id in link.transport.poll_cancellations()) except Exception as exc: # noqa: BLE001 self._fail(link, exc) return cancelled def _link_for(self, job_id: str) -> ServerLink | None: with self._lock: return self._owner.get(job_id) def acknowledge_offer(self, decision: Any) -> None: link = self._link_for(getattr(decision, "job_id", "")) if link is None: return try: link.transport.acknowledge_offer(decision) except Exception as exc: # noqa: BLE001 self._fail(link, exc) if not getattr(decision, "accepted", True): self._release(getattr(decision, "job_id", "")) def publish_status( self, job_id: str, status: JobStatus, reason: str | None = None ) -> None: link = self._link_for(job_id) if link is None: return try: link.transport.publish_status(job_id, status, reason) except Exception as exc: # noqa: BLE001 self._fail(link, exc) def publish_result(self, result: JobResult) -> None: """To the server that asked, and to no other. A missing owner is dropped rather than broadcast. Publishing an answer to every server would hand one server's user's text to operators who never asked for it, which is a worse outcome than losing a result the worker cannot attribute. """ link = self._link_for(result.job_id) if link is None: self._report( f"Result for {result.job_id} has no server to return to; discarding it " "rather than sending it to a server that did not ask." ) return try: link.transport.publish_result(result) except Exception as exc: # noqa: BLE001 self._fail(link, exc) finally: self._release(result.job_id) def _release(self, job_id: str) -> None: with self._lock: link = self._owner.pop(job_id, None) if link is not None: link.job_ids.discard(job_id) def close(self) -> None: for link in self.links: close = getattr(link.transport, "close", None) if callable(close): try: close() except Exception: # noqa: BLE001 - closing must not raise pass __all__ = [ "BACKOFF_CEILING_SECONDS", "BACKOFF_START_SECONDS", "MultiServerTransport", "ServerLink", ]