Spaces:
Paused
Paused
Atomic node repair 2026-10-22: controller_token config + registration recovery/readiness gating (candidate ee4db2ec)
748e442 verified | """The node agent: serve a model, join the mesh, hold the lease, leave cleanly. | |
| Lifecycle | |
| --------- | |
| 1. Resolve config and (optionally) a device identity. | |
| 2. Load the backend and start the OpenAI-compatible server on a background thread. | |
| 3. Determine the public URL — a Space origin, a tunnel, or an override. | |
| 4. Register with the mesh, retrying transport failures with backoff. | |
| 5. Heartbeat on a cadence strictly shorter than the TTL, refreshing telemetry. | |
| 6. On shutdown, deregister so the endpoint leaves routing immediately. | |
| Ephemerality is the design constraint that shapes all of it. A Colab session can | |
| vanish without warning — the VM is reclaimed, the browser tab closes, the | |
| 12-hour cap hits — and there is no shutdown hook that reliably fires. Clean | |
| deregistration is therefore treated as an optimisation, not a guarantee, and | |
| correctness rests on the lease in ``capabilities.thox_lease``: the node keeps | |
| pushing an ``expires_at`` a TTL into the future, and any consumer treats a lapsed | |
| lease as ineligible. That way a node that dies mid-sentence stops receiving | |
| traffic within one TTL whether or not it ever got to say goodbye. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| import os | |
| import signal | |
| import threading | |
| import time | |
| from datetime import datetime | |
| from typing import Any | |
| import uvicorn | |
| from .backends import Backend, build_backend | |
| from .config import NodeConfig, detect_node_kind, load_config, space_public_url | |
| from .errors import ConfigError, MeshRejectedError, MeshTransportError, PublicUrlRejectedError | |
| from .mesh_client import ControllerClient, MeshClient, Registration, build_client | |
| from .server import create_app | |
| from .signing import DeviceIdentity, load_identity_from_env | |
| from .telemetry import TelemetryRecorder, build_capabilities | |
| logger = logging.getLogger("thox.agent") | |
| class NodeAgent: | |
| """Owns the server, the mesh registration and the heartbeat loop.""" | |
| def __init__( | |
| self, | |
| config: NodeConfig, | |
| *, | |
| backend: Backend, | |
| identity: DeviceIdentity | None = None, | |
| client: MeshClient | None = None, | |
| ) -> None: | |
| self.config = config | |
| self.backend = backend | |
| self.recorder = TelemetryRecorder() | |
| self.state: dict[str, Any] = { | |
| "started_at": time.time(), | |
| "node_kind": config.node_kind, | |
| "transport": None, | |
| "endpoint_id": None, | |
| "public_url": config.public_url or None, | |
| "last_heartbeat_at": None, | |
| "last_error": None, | |
| "stopping": False, | |
| } | |
| self._identity = identity | |
| self._client = client or build_client(config, identity) | |
| self.state.update( | |
| mesh_required=self._client.name != "none", | |
| registration_status="pending" if self._client.name != "none" else "disabled", | |
| lease_expires_at=None, | |
| ) | |
| self._registration: Registration | None = None | |
| # SIGTERM handlers can call shutdown on this same thread while a | |
| # registration result is being adopted. | |
| self._membership_lock = threading.RLock() | |
| self._stop = threading.Event() | |
| self._heartbeat_thread: threading.Thread | None = None | |
| self._server: uvicorn.Server | None = None | |
| self._server_thread: threading.Thread | None = None | |
| self.app = create_app( | |
| backend=backend, | |
| recorder=self.recorder, | |
| model_id=config.model_id, | |
| context_window=config.context_window, | |
| max_tokens_cap=config.max_tokens_cap, | |
| node_state=self.state, | |
| ) | |
| # ── Capabilities ───────────────────────────────────────────────────── | |
| def _capabilities(self) -> dict[str, Any]: | |
| """Current capability + telemetry payload for register/heartbeat.""" | |
| return build_capabilities( | |
| base=self.config.capabilities, | |
| recorder=self.recorder, | |
| node_kind=self.config.node_kind, | |
| ephemeral=self.config.ephemeral, | |
| heartbeat_seconds=self.config.heartbeat_seconds, | |
| ttl_seconds=self.config.ttl_seconds, | |
| ) | |
| # ── Server ─────────────────────────────────────────────────────────── | |
| def start_server(self) -> None: | |
| """Run uvicorn on a daemon thread and wait until it accepts requests.""" | |
| uvicorn_config = uvicorn.Config( | |
| self.app, | |
| host=self.config.host, | |
| port=self.config.port, | |
| log_level="info", | |
| access_log=False, | |
| ) | |
| self._server = uvicorn.Server(uvicorn_config) | |
| self._server_thread = threading.Thread(target=self._server.run, daemon=True, name="thox-node-server") | |
| self._server_thread.start() | |
| deadline = time.time() + 60 | |
| while time.time() < deadline: | |
| if getattr(self._server, "started", False): | |
| logger.info("serving on %s:%d", self.config.host, self.config.port) | |
| return | |
| if not self._server_thread.is_alive(): | |
| raise ConfigError("the HTTP server thread exited during startup") | |
| time.sleep(0.25) | |
| raise ConfigError("the HTTP server did not start within 60s") | |
| # ── Registration ───────────────────────────────────────────────────── | |
| def register(self, public_url: str, *, max_attempts: int = 5) -> Registration: | |
| """Publish this node, retrying transport errors with linear backoff. | |
| A ``PublicUrlRejectedError`` on the canonical transport is not retried: | |
| it means the mesh's private-address policy refused this URL, and no | |
| number of retries changes that. When a controller URL is configured the | |
| agent degrades to it, which is what keeps an ephemeral public node usable | |
| while the policy question is resolved with the Mesh backend lane. | |
| """ | |
| self.state["public_url"] = public_url | |
| last_error: Exception | None = None | |
| for attempt in range(1, max_attempts + 1): | |
| try: | |
| if self._stop.is_set(): | |
| raise MeshTransportError("registration stopped during shutdown") | |
| capabilities = self._capabilities() | |
| lease_expires_at = datetime.fromisoformat( | |
| capabilities["thox_lease"]["expires_at"].replace("Z", "+00:00") | |
| ).timestamp() | |
| registration = self._client.register( | |
| self.config, public_url, capabilities, self.recorder.status() | |
| ) | |
| with self._membership_lock: | |
| if self._stop.is_set(): | |
| # A response can arrive after SIGTERM. Retire that late | |
| # admission rather than advertise a server being stopped. | |
| self._client.deregister(registration) | |
| raise MeshTransportError("registration stopped during shutdown") | |
| self._registration = registration | |
| self.state.update( | |
| transport=registration.transport, | |
| endpoint_id=registration.endpoint_id, | |
| last_error=None, | |
| registration_status="registered" if registration.transport != "none" else "disabled", | |
| lease_expires_at=lease_expires_at, | |
| ) | |
| if self._stop.is_set(): | |
| self._client.deregister(registration) | |
| self._registration = None | |
| self.state.update(endpoint_id=None, registration_status="stopped") | |
| raise MeshTransportError("registration stopped during shutdown") | |
| return registration | |
| except PublicUrlRejectedError as exc: | |
| logger.error("mesh refused the public URL: %s", exc) | |
| self.state["last_error"] = str(exc) | |
| self.state["registration_status"] = "rejected" | |
| if self.config.controller_url and not isinstance(self._client, ControllerClient): | |
| logger.warning( | |
| "falling back to controller transport at %s because the v2 control " | |
| "plane only accepts private addresses", | |
| self.config.controller_url, | |
| ) | |
| self._client = ControllerClient( | |
| self.config.controller_url, token=self.config.controller_token or None | |
| ) | |
| continue | |
| raise | |
| except MeshTransportError as exc: | |
| if self._stop.is_set(): | |
| self.state["registration_status"] = "stopped" | |
| raise | |
| last_error = exc | |
| self.state["last_error"] = str(exc) | |
| self.state["registration_status"] = "pending" | |
| wait = self.config.register_retry_seconds * attempt | |
| logger.warning("register attempt %d/%d failed: %s (retrying in %ds)", attempt, max_attempts, exc, wait) | |
| if attempt < max_attempts: | |
| if self._stop.wait(wait): | |
| raise MeshTransportError("registration stopped during shutdown") from exc | |
| except MeshRejectedError as exc: | |
| self.state["last_error"] = str(exc) | |
| self.state["registration_status"] = "rejected" | |
| logger.error("mesh rejected registration: %s", exc) | |
| raise | |
| raise MeshTransportError(f"could not register after {max_attempts} attempts: {last_error}") | |
| # ── Heartbeat ──────────────────────────────────────────────────────── | |
| def _heartbeat_loop(self) -> None: | |
| """Renew the lease until stopped, tolerating transient mesh outages.""" | |
| consecutive_failures = 0 | |
| while not self._stop.wait(self.config.heartbeat_seconds): | |
| registration = self._registration | |
| if registration is None: | |
| # Startup admission can fail while the model is already serving. | |
| # Retry once per cadence, without restarting or claiming membership. | |
| try: | |
| self.register(self.state["public_url"], max_attempts=1) | |
| except MeshTransportError: | |
| continue | |
| except MeshRejectedError as exc: | |
| self.state["last_error"] = str(exc) | |
| self.state["registration_status"] = "rejected" | |
| logger.error("registration rejected; operator action required: %s", exc) | |
| return | |
| continue | |
| try: | |
| capabilities = self._capabilities() | |
| lease_expires_at = datetime.fromisoformat( | |
| capabilities["thox_lease"]["expires_at"].replace("Z", "+00:00") | |
| ).timestamp() | |
| self._client.heartbeat( | |
| self.config, registration, capabilities, self.recorder.status() | |
| ) | |
| if self._stop.is_set(): | |
| # Controller heartbeats are registration upserts. A late | |
| # reply must not recreate membership after shutdown. | |
| self._client.deregister(registration) | |
| return | |
| with self._membership_lock: | |
| if self._stop.is_set(): | |
| return | |
| self.state["last_heartbeat_at"] = time.time() | |
| self.state["last_error"] = None | |
| self.state["registration_status"] = "registered" if registration.transport != "none" else "disabled" | |
| self.state["lease_expires_at"] = lease_expires_at | |
| if self._stop.is_set(): | |
| self.state["registration_status"] = "stopped" | |
| consecutive_failures = 0 | |
| except PublicUrlRejectedError as exc: | |
| # Policy, not an outage; stop hammering and let the lease lapse. | |
| logger.error("heartbeat refused by URL policy, stopping renewal: %s", exc) | |
| self.state["last_error"] = str(exc) | |
| self.state["registration_status"] = "rejected" | |
| return | |
| except MeshRejectedError as exc: | |
| self.state["last_error"] = str(exc) | |
| self.state["registration_status"] = "rejected" | |
| logger.error("heartbeat rejected; operator action required: %s", exc) | |
| return | |
| except MeshTransportError as exc: | |
| consecutive_failures += 1 | |
| self.state["last_error"] = str(exc) | |
| self.state["registration_status"] = "degraded" | |
| logger.warning("heartbeat failed (%d in a row): %s", consecutive_failures, exc) | |
| if consecutive_failures >= 3: | |
| logger.error( | |
| "3 consecutive heartbeat failures; the lease will lapse in ~%ds " | |
| "and consumers will stop routing here", | |
| self.config.ttl_seconds, | |
| ) | |
| def start_heartbeat(self) -> None: | |
| if self._heartbeat_thread is not None: | |
| return | |
| self._heartbeat_thread = threading.Thread(target=self._heartbeat_loop, daemon=True, name="thox-node-heartbeat") | |
| self._heartbeat_thread.start() | |
| logger.info( | |
| "heartbeat every %ds, lease TTL %ds", self.config.heartbeat_seconds, self.config.ttl_seconds | |
| ) | |
| # ── Shutdown ───────────────────────────────────────────────────────── | |
| def shutdown(self) -> None: | |
| """Leave the mesh and stop serving. Safe to call more than once.""" | |
| if self._stop.is_set(): | |
| return | |
| self._stop.set() | |
| self.state["stopping"] = True | |
| logger.info("shutting down") | |
| with self._membership_lock: | |
| self.state["registration_status"] = "stopped" | |
| if self._registration is not None: | |
| self._client.deregister(self._registration) | |
| self._registration = None | |
| self.state["endpoint_id"] = None | |
| if self._server is not None: | |
| self._server.should_exit = True | |
| try: | |
| self.backend.close() | |
| except Exception: # noqa: BLE001 | |
| logger.debug("backend close failed", exc_info=True) | |
| def install_signal_handlers(self) -> None: | |
| """Deregister on SIGINT/SIGTERM where the platform supports it. | |
| Colab interrupts land as SIGINT and Spaces stop containers with SIGTERM, | |
| so both give us a chance to leave cleanly. A hard VM reclaim gives no | |
| signal at all, which is exactly why the lease exists. | |
| """ | |
| def _handler(signum, _frame): # type: ignore[no-untyped-def] | |
| logger.info("signal %s received", signum) | |
| self.shutdown() | |
| for sig in (signal.SIGINT, signal.SIGTERM): | |
| try: | |
| signal.signal(sig, _handler) | |
| except (ValueError, OSError): | |
| # Not the main thread (common in notebooks); the lease still covers us. | |
| logger.debug("could not install handler for %s", sig) | |
| def wait(self) -> None: | |
| """Block until shutdown. Used by the CLI entry point.""" | |
| try: | |
| while not self._stop.is_set(): | |
| time.sleep(1) | |
| except KeyboardInterrupt: | |
| self.shutdown() | |
| def resolve_public_url(config: NodeConfig, *, tunnel_provider: str = "cloudflared") -> tuple[str, Any]: | |
| """Determine how the mesh will reach this node. | |
| Order: explicit override, then a Space's own origin, then a tunnel. Returns | |
| the URL and the tunnel handle (``None`` when no tunnel was started) so the | |
| caller can close it. | |
| """ | |
| if config.public_url: | |
| logger.info("using configured public URL %s", config.public_url) | |
| return config.public_url, None | |
| if config.node_kind == "hf-space": | |
| url = space_public_url() | |
| if not url: | |
| raise ConfigError( | |
| "running on a Space but neither SPACE_HOST nor SPACE_ID is set; " | |
| "set THOX_PUBLIC_URL explicitly rather than registering an unreachable address" | |
| ) | |
| return url, None | |
| from .tunnel import start_tunnel | |
| tunnel = start_tunnel(tunnel_provider, config.port) | |
| return tunnel.public_url, tunnel | |
| def run( | |
| *, | |
| env: dict[str, str] | None = None, | |
| echo: bool = False, | |
| tunnel_provider: str = "cloudflared", | |
| block: bool = True, | |
| ) -> NodeAgent: | |
| """Start a fully wired node. The single entry point notebooks and Spaces call.""" | |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(name)s] %(levelname)s %(message)s") | |
| config = load_config(env) | |
| if config.node_kind == "unknown": | |
| config = type(config)(**{**config.__dict__, "node_kind": detect_node_kind(env)}) | |
| identity = load_identity_from_env(env) | |
| backend = build_backend( | |
| model_id=config.model_id, | |
| repo_id=config.model_repo, | |
| filename=config.model_file, | |
| n_ctx=config.context_window, | |
| n_threads=config.n_threads, | |
| n_gpu_layers=config.n_gpu_layers, | |
| token=(env or os.environ).get("HF_TOKEN"), | |
| echo=echo, | |
| ) | |
| agent = NodeAgent(config, backend=backend, identity=identity) | |
| agent.install_signal_handlers() | |
| agent.start_server() | |
| tunnel = None | |
| try: | |
| public_url, tunnel = resolve_public_url(config, tunnel_provider=tunnel_provider) | |
| try: | |
| agent.register(public_url) | |
| except MeshTransportError as exc: | |
| if agent._stop.is_set(): | |
| raise | |
| logger.warning("model loaded; inference awaits mesh registration: %s", exc) | |
| except Exception: | |
| agent.shutdown() | |
| if tunnel is not None: | |
| tunnel.close() | |
| raise | |
| agent.start_heartbeat() | |
| logger.info("node serving: %s -> %s (mesh=%s)", config.model_id, public_url, agent.state["registration_status"]) | |
| if block: | |
| try: | |
| agent.wait() | |
| finally: | |
| agent.shutdown() | |
| if tunnel is not None: | |
| tunnel.close() | |
| return agent | |