| 1 | """Echo generator admission / traffic control. |
| 2 | |
| 3 | Before submitting generation jobs, poll ``{echoGenerator.baseUrl}/health`` and |
| 4 | reject new work when the algorithm side is overloaded: |
| 5 | |
| 6 | workers = scheduler.queues.inference.workers |
| 7 | works = scheduler.gpu_busy + scheduler.total_queued |
| 8 | capacity = workers |
| 9 | reject when works >= capacity |
| 10 | """ |
| 11 | |
| 12 | from __future__ import annotations |
| 13 | |
| 14 | import errno |
| 15 | import json |
| 16 | from dataclasses import dataclass |
| 17 | from typing import Any |
| 18 | from urllib import error as urllib_error |
| 19 | from urllib import request as urllib_request |
| 20 | |
| 21 | from loguru import logger |
| 22 | |
| 23 | from nanobot.config.schema import EchoGeneratorConfig |
| 24 | |
| 25 | BUSY_MESSAGE = "The video generation service is busy. Please try again shortly." |
| 26 | UNAVAILABLE_MESSAGE = "The video generation service is temporarily unavailable. Please try again shortly." |
| 27 | _GENERATION_OPERATIONS = frozenset({"generate_echo_shot", "r2v_generate"}) |
| 28 | |
| 29 | |
| 30 | class EchoGeneratorBusyError(RuntimeError): |
| 31 | """Raised when the Echo generator should not accept more generation work.""" |
| 32 | |
| 33 | def __init__(self, message: str = BUSY_MESSAGE, *, snapshot: "AdmissionSnapshot | None" = None): |
| 34 | super().__init__(message) |
| 35 | self.snapshot = snapshot |
| 36 | |
| 37 | |
| 38 | class EchoGeneratorUnavailableError(RuntimeError): |
| 39 | """Raised when the Echo generator host cannot be reached.""" |
| 40 | |
| 41 | def __init__(self, message: str = UNAVAILABLE_MESSAGE): |
| 42 | super().__init__(message) |
| 43 | |
| 44 | |
| 45 | def is_connection_refused(exc: BaseException) -> bool: |
| 46 | current: BaseException | None = exc |
| 47 | for _ in range(5): |
| 48 | if current is None: |
| 49 | break |
| 50 | if isinstance(current, ConnectionRefusedError): |
| 51 | return True |
| 52 | if getattr(current, "errno", None) == errno.ECONNREFUSED: |
| 53 | return True |
| 54 | reason = getattr(current, "reason", None) |
| 55 | if isinstance(reason, BaseException) and reason is not current: |
| 56 | current = reason |
| 57 | continue |
| 58 | current = current.__cause__ or current.__context__ |
| 59 | text = str(exc).lower() |
| 60 | return "connection refused" in text or "errno 111" in text |
| 61 | |
| 62 | |
| 63 | @dataclass(frozen=True) |
| 64 | class AdmissionSnapshot: |
| 65 | """Parsed health signals used for the admission decision.""" |
| 66 | |
| 67 | workers: float | None |
| 68 | capacity: float | None |
| 69 | works: float | None |
| 70 | busy: bool |
| 71 | reason: str |
| 72 | raw: dict[str, Any] | None = None |
| 73 | |
| 74 | |
| 75 | class EchoAdmissionController: |
| 76 | """Traffic gate for Echo / JoyEcho generation endpoints.""" |
| 77 | |
| 78 | def __init__( |
| 79 | self, |
| 80 | config: EchoGeneratorConfig | None = None, |
| 81 | *, |
| 82 | base_url: str | None = None, |
| 83 | timeout_sec: float | None = None, |
| 84 | fail_open: bool = True, |
| 85 | ) -> None: |
| 86 | cfg = config or EchoGeneratorConfig() |
| 87 | resolved_base = (base_url if base_url is not None else cfg.base_url) or "" |
| 88 | self.base_url = str(resolved_base).strip().rstrip("/") |
| 89 | raw_timeout = timeout_sec if timeout_sec is not None else cfg.http_timeout_sec |
| 90 | try: |
| 91 | self.timeout_sec = max(1.0, float(raw_timeout)) |
| 92 | except (TypeError, ValueError): |
| 93 | self.timeout_sec = 30.0 |
| 94 | self.fail_open = fail_open |
| 95 | |
| 96 | @classmethod |
| 97 | def from_tools_config(cls, tools_config: Any, **kwargs: Any) -> "EchoAdmissionController": |
| 98 | return cls( |
| 99 | getattr(tools_config, "echo_generator", None) or EchoGeneratorConfig(), |
| 100 | **kwargs, |
| 101 | ) |
| 102 | |
| 103 | def applies_to_operation(self, operation: str | None) -> bool: |
| 104 | return (operation or "") in _GENERATION_OPERATIONS |
| 105 | |
| 106 | def fetch_health(self) -> dict[str, Any]: |
| 107 | if not self.base_url: |
| 108 | raise RuntimeError("echoGenerator.baseUrl is not configured") |
| 109 | url = f"{self.base_url}/health" |
| 110 | req = urllib_request.Request(url, method="GET") |
| 111 | with urllib_request.urlopen(req, timeout=self.timeout_sec) as resp: |
| 112 | raw = resp.read().decode("utf-8") |
| 113 | data = json.loads(raw) if raw else {} |
| 114 | if not isinstance(data, dict): |
| 115 | raise RuntimeError(f"unexpected /health payload type: {type(data).__name__}") |
| 116 | return data |
| 117 | |
| 118 | @staticmethod |
| 119 | def _as_number(value: Any) -> float | None: |
| 120 | if isinstance(value, bool) or value is None: |
| 121 | return None |
| 122 | if isinstance(value, (int, float)): |
| 123 | return float(value) |
| 124 | if isinstance(value, str) and value.strip(): |
| 125 | try: |
| 126 | return float(value.strip()) |
| 127 | except ValueError: |
| 128 | return None |
| 129 | return None |
| 130 | |
| 131 | @classmethod |
| 132 | def _dig(cls, data: dict[str, Any], *path: str) -> Any: |
| 133 | cur: Any = data |
| 134 | for key in path: |
| 135 | if not isinstance(cur, dict): |
| 136 | return None |
| 137 | cur = cur.get(key) |
| 138 | return cur |
| 139 | |
| 140 | @classmethod |
| 141 | def extract_signals(cls, health: dict[str, Any]) -> tuple[float | None, float | None]: |
| 142 | """Resolve workers / works from /health. |
| 143 | |
| 144 | - workers = ``scheduler.queues.inference.workers`` |
| 145 | - works = ``scheduler.gpu_busy`` + ``scheduler.total_queued`` |
| 146 | """ |
| 147 | workers = cls._as_number( |
| 148 | cls._dig(health, "scheduler", "queues", "inference", "workers") |
| 149 | ) |
| 150 | gpu_busy = cls._as_number(cls._dig(health, "scheduler", "gpu_busy")) |
| 151 | total_queued = cls._as_number(cls._dig(health, "scheduler", "total_queued")) |
| 152 | if gpu_busy is None or total_queued is None: |
| 153 | works = None |
| 154 | else: |
| 155 | works = gpu_busy + total_queued |
| 156 | return workers, works |
| 157 | |
| 158 | def evaluate(self, health: dict[str, Any]) -> AdmissionSnapshot: |
| 159 | if str(health.get("status") or "").strip().lower() not in {"", "ok"}: |
| 160 | return AdmissionSnapshot( |
| 161 | workers=None, |
| 162 | capacity=None, |
| 163 | works=None, |
| 164 | busy=True, |
| 165 | reason="status_unhealthy", |
| 166 | raw=health, |
| 167 | ) |
| 168 | |
| 169 | workers, works = self.extract_signals(health) |
| 170 | if workers is None or works is None: |
| 171 | return AdmissionSnapshot( |
| 172 | workers=workers, |
| 173 | capacity=None, |
| 174 | works=works, |
| 175 | busy=False, |
| 176 | reason="missing_workers_or_works", |
| 177 | raw=health, |
| 178 | ) |
| 179 | |
| 180 | capacity = workers |
| 181 | busy = works >= capacity |
| 182 | return AdmissionSnapshot( |
| 183 | workers=workers, |
| 184 | capacity=capacity, |
| 185 | works=works, |
| 186 | busy=busy, |
| 187 | reason="at_or_over_capacity" if busy else "within_capacity", |
| 188 | raw=health, |
| 189 | ) |
| 190 | |
| 191 | def check(self) -> AdmissionSnapshot: |
| 192 | """Fetch health and decide. Raises ``EchoGeneratorBusyError`` when overloaded.""" |
| 193 | if not self.base_url: |
| 194 | return AdmissionSnapshot( |
| 195 | workers=None, |
| 196 | capacity=None, |
| 197 | works=None, |
| 198 | busy=False, |
| 199 | reason="base_url_not_configured", |
| 200 | ) |
| 201 | try: |
| 202 | health = self.fetch_health() |
| 203 | except ( |
| 204 | urllib_error.URLError, |
| 205 | urllib_error.HTTPError, |
| 206 | TimeoutError, |
| 207 | OSError, |
| 208 | json.JSONDecodeError, |
| 209 | RuntimeError, |
| 210 | ) as exc: |
| 211 | logger.warning( |
| 212 | "Echo admission health check failed (fail_open={}): {}", |
| 213 | self.fail_open, |
| 214 | exc, |
| 215 | ) |
| 216 | if is_connection_refused(exc): |
| 217 | raise EchoGeneratorUnavailableError(UNAVAILABLE_MESSAGE) from exc |
| 218 | if self.fail_open: |
| 219 | return AdmissionSnapshot( |
| 220 | workers=None, |
| 221 | capacity=None, |
| 222 | works=None, |
| 223 | busy=False, |
| 224 | reason=f"health_check_failed:{exc}", |
| 225 | ) |
| 226 | raise EchoGeneratorBusyError(BUSY_MESSAGE) from exc |
| 227 | |
| 228 | snapshot = self.evaluate(health) |
| 229 | logger.info( |
| 230 | "Echo admission check workers={} capacity={} works={} busy={} reason={}", |
| 231 | snapshot.workers, |
| 232 | snapshot.capacity, |
| 233 | snapshot.works, |
| 234 | snapshot.busy, |
| 235 | snapshot.reason, |
| 236 | ) |
| 237 | if snapshot.busy: |
| 238 | raise EchoGeneratorBusyError(BUSY_MESSAGE, snapshot=snapshot) |
| 239 | return snapshot |
| 240 | |
| 241 | def ensure_allowed(self, *, operation: str | None = None) -> AdmissionSnapshot: |
| 242 | """Gate generation operations; no-op for unrelated operations.""" |
| 243 | if operation is not None and not self.applies_to_operation(operation): |
| 244 | return AdmissionSnapshot( |
| 245 | workers=None, |
| 246 | capacity=None, |
| 247 | works=None, |
| 248 | busy=False, |
| 249 | reason=f"skip_operation:{operation}", |
| 250 | ) |
| 251 | return self.check() |
| 252 |