atlas-iac/services/hermes/model-gate-configmap.yaml

# services/hermes/model-gate-configmap.yaml
apiVersion: v1
kind: ConfigMap
metadata:
  name: hermes-model-gate
  namespace: hermes
data:
  model_gate.py: |
    #!/usr/bin/env python3
    """Normalize Jetson text requests and coordinate the titan-24 image handoff."""

    from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
    import hmac
    import ipaddress
    import json
    import os
    from pathlib import Path
    import ssl
    import threading
    import time
    import lan_generate
    from urllib.error import HTTPError, URLError
    from urllib.request import Request, urlopen


    LISTEN_HOST = os.environ.get("LISTEN_HOST", "0.0.0.0")
    LISTEN_PORT = int(os.environ.get("LISTEN_PORT", "8080"))
    LAN_LISTEN_PORT = int(os.environ.get("LAN_LISTEN_PORT", "8082"))
    HANDOFF_PORT = int(os.environ.get("HANDOFF_PORT", "8081"))
    UPSTREAM_URL = os.environ.get("UPSTREAM_URL", "http://ollama.ai.svc.cluster.local:11434").rstrip("/")
    LOCAL_IMAGE_URL = os.environ.get("LOCAL_IMAGE_URL", "http://hermes-local-image.hermes.svc.cluster.local:9004").rstrip("/")
    HANDOFF_TIMEOUT_SEC = float(os.environ.get("HANDOFF_TIMEOUT_SEC", "1200"))
    IMAGE_NAMESPACE = os.environ.get("IMAGE_NAMESPACE", "hermes")
    IMAGE_DEPLOYMENT = os.environ.get("IMAGE_DEPLOYMENT", "hermes-local-image")
    LAN_API_TOKEN_FILE = Path(os.environ.get("LAN_API_TOKEN_FILE", "/vault/secrets/lan-api-token"))
    LAN_MODEL = lan_generate.MODEL
    LAN_NETWORK = ipaddress.ip_network("192.168.22.0/24")
    LAN_MAX_BODY = lan_generate.MAX_BODY
    KUBE_HOST = os.environ.get("KUBERNETES_SERVICE_HOST", "kubernetes.default.svc")
    KUBE_PORT = os.environ.get("KUBERNETES_SERVICE_PORT_HTTPS", "443")
    KUBE_TOKEN = Path("/var/run/secrets/kubernetes.io/serviceaccount/token")
    KUBE_CA = Path("/var/run/secrets/kubernetes.io/serviceaccount/ca.crt")
    _handoff_lock = threading.Lock()
    _image_lane_warm = True


    def _image_deployment_ready() -> bool:
      """Fail closed on transient network loss, but do not wait on a dead pod."""

      token = KUBE_TOKEN.read_text(encoding="utf-8").strip()
      url = (
        f"https://{KUBE_HOST}:{KUBE_PORT}/apis/apps/v1/namespaces/"
        f"{IMAGE_NAMESPACE}/deployments/{IMAGE_DEPLOYMENT}"
      )
      request = Request(url, headers={"Authorization": f"Bearer {token}"})
      context = ssl.create_default_context(cafile=str(KUBE_CA))
      with urlopen(request, timeout=5, context=context) as response:
        payload = json.load(response)
      return int((payload.get("status") or {}).get("readyReplicas") or 0) > 0


    def _normalize_reasoning(body: bytes | None) -> bytes | None:
      """Translate routed aliases and clamp effort to the local server tier."""

      if not body:
        return body
      try:
        payload = json.loads(body)
      except (TypeError, ValueError, json.JSONDecodeError):
        return body
      if not isinstance(payload, dict):
        return body

      changed = False
      model = payload.get("model")
      if isinstance(model, str) and model.startswith("route/local/qwen2.5-14b/"):
        payload["model"] = "qwen2.5:14b-instruct-q4_0"
        changed = True
      elif not isinstance(model, str) or not model.strip():
        # This gate serves one text model, so an omitted routed target has one
        # unambiguous and safe default.
        payload["model"] = "qwen2.5:14b-instruct-q4_0"
        changed = True
      # This non-thinking Qwen build rejects hosted-provider reasoning fields.
      # Effort remains a Switchyard routing concern, not an Ollama request.
      for key in ("reasoning_effort", "reasoning", "thinking", "output_config"):
        value = payload.get(key)
        if value is not None:
          payload.pop(key, None)
          changed = True
      if not changed:
        return body
      return json.dumps(payload, separators=(",", ":")).encode("utf-8")


    def _wait_for_image_idle() -> tuple[bool, str]:
      """Keep Wolf handoff blocked until an in-flight FLUX render is released."""

      deadline = time.monotonic() + HANDOFF_TIMEOUT_SEC
      last_error = "local image service did not report idle"
      while time.monotonic() < deadline:
        try:
          with urlopen(f"{LOCAL_IMAGE_URL}/health", timeout=5) as response:
            payload = json.load(response)
          if not bool(payload.get("busy")):
            return True, ""
          last_error = f"local image renderer is {payload.get('phase', 'busy')}"
        except Exception as exc:
          last_error = f"local image health unavailable: {exc}"
          try:
            if not _image_deployment_ready():
              # A stopped/crashed process has no CUDA context left to drain.
              return True, ""
          except Exception as kube_exc:
            last_error += f"; deployment state unavailable: {kube_exc}"
        time.sleep(1)
      return False, last_error


    def _is_lan_client(client_ip: str) -> bool:
      """Validate the address appended by the trusted ingress proxy."""
      try:
        return ipaddress.ip_address(client_ip) in LAN_NETWORK
      except ValueError:
        return False


    class LanHandler(BaseHTTPRequestHandler):
      """Narrow authenticated LAN API: health plus one pinned Ollama operation."""

      protocol_version = "HTTP/1.1"

      def setup(self) -> None:
        """Bound the time a client can hold a request body open."""
        super().setup()
        self.connection.settimeout(10)

      def _json(self, status: int, payload: dict) -> None:
        """End each request so rejected bodies cannot become another request."""
        body = json.dumps(payload, separators=(",", ":")).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.send_header("Connection", "close")
        self.close_connection = True
        self.end_headers()
        self.wfile.write(body)

      def _authorized(self) -> bool:
        """Require both the ingress-observed LAN address and the Vault token."""
        forwarded_for = self.headers.get("X-Forwarded-For", "").split(",")[-1].strip()
        if not _is_lan_client(forwarded_for):
          self._json(403, {"error": "LAN source required"})
          return False
        try:
          expected = LAN_API_TOKEN_FILE.read_text(encoding="utf-8").strip()
        except OSError:
          self._json(503, {"error": "LAN API credential unavailable"})
          return False
        supplied = self.headers.get("Authorization", "")
        if not expected or not supplied.startswith("Bearer ") or not hmac.compare_digest(supplied[7:].encode(), expected.encode()):
          self._json(401, {"error": "Bearer authentication required"})
          return False
        return True

      def do_GET(self) -> None:
        """Expose authenticated readiness without model or session metadata."""
        if self.path != "/healthz":
          self._json(404, {"error": "not found"})
          return
        if self._authorized():
          self._json(*lan_generate.health())

      def do_POST(self) -> None:
        """Validate and serialize bounded inference against local Ollama only."""
        if self.path != "/api/generate":
          self._json(404, {"error": "not found"})
          return
        if not self._authorized():
          return
        try:
          if self.headers.get("Transfer-Encoding"):
            raise ValueError("Transfer-Encoding is unsupported")
          length = int(self.headers.get("Content-Length", "0"))
          maximum = LAN_MAX_BODY
          if not 0 < length <= maximum:
            self._json(413, {"error": f"body must be between 1 and {maximum} bytes"})
            return
          body = self.rfile.read(length)
          if len(body) != length:
            raise ValueError("incomplete body")
        except ValueError:
          self._json(400, {"error": "invalid request framing"})
          return
        except TimeoutError:
          self._json(408, {"error": "request body timed out"})
          return
        self._json(*lan_generate.generate(body))

      def send_error(self, code, message=None, explain=None) -> None:
        """Do not reflect malformed methods, request targets, or parser input."""
        self._json(code, {"error": "unsupported or malformed HTTP request"})

      def log_request(self, code="-", size="-") -> None:
        """Only fixed route names and status codes enter container logs."""
        route = {"/healthz": "health", "/api/generate": "generate"}.get(getattr(self, "path", ""), "rejected")
        print(f"model-gate-lan route={route} status={code}", flush=True)

      def log_message(self, format_string: str, *args) -> None:
        # Parser errors may contain attacker-controlled text; omit it entirely.
        return


    class Handler(BaseHTTPRequestHandler):
      """Proxy the non-preemptible titan-20 text fallback."""

      protocol_version = "HTTP/1.1"

      def _json(self, status: int, payload: dict) -> None:
        body = json.dumps(payload, separators=(",", ":")).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(body)

      def _proxy(self) -> None:
        length = int(self.headers.get("Content-Length", "0") or 0)
        body = self.rfile.read(length) if length else None
        body = _normalize_reasoning(body)
        headers = {"Content-Type": self.headers.get("Content-Type", "application/json")}
        if self.headers.get("Accept"):
          headers["Accept"] = self.headers["Accept"]
        request = Request(f"{UPSTREAM_URL}{self.path}", data=body, headers=headers, method=self.command)
        try:
          response = urlopen(request, timeout=1800)
        except HTTPError as exc:
          response = exc
        except (TimeoutError, URLError) as exc:
          self._json(503, {"error": {"message": f"local model upstream unavailable: {exc}", "type": "server_error"}})
          return

        self.send_response(response.status)
        content_type = response.headers.get("Content-Type")
        if content_type:
          self.send_header("Content-Type", content_type)
        content_length = response.headers.get("Content-Length")
        if content_length:
          self.send_header("Content-Length", content_length)
        else:
          self.send_header("Connection", "close")
          self.close_connection = True
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        while True:
          chunk = response.read(65536)
          if not chunk:
            break
          self.wfile.write(chunk)
          self.wfile.flush()
        response.close()

      def do_GET(self) -> None:
        if self.path == "/healthz":
          self._json(200, {"status": "ok"})
          return
        if self.path == "/gate/status":
          self._json(200, {"local_inference_allowed": True, "node": "titan-20"})
          return
        self._proxy()

      def do_POST(self) -> None:
        self._proxy()

      def log_message(self, format_string: str, *args) -> None:
        print(f"model-gate {self.address_string()} {format_string % args}", flush=True)


    class HandoffHandler(BaseHTTPRequestHandler):
      """Ollama-compatible adapter for Ariadne's image/Wolf handoff contract."""

      protocol_version = "HTTP/1.1"

      def _json(self, status: int, payload: dict) -> None:
        body = json.dumps(payload, separators=(",", ":")).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(body)

      def _proxy(self) -> None:
        global _image_lane_warm
        if self.command == "GET" and self.path == "/api/ps":
          models = [{"name": "flux-2-klein-4b-local"}] if _image_lane_warm else []
          self._json(200, {"models": models})
          return
        if self.command != "POST" or self.path != "/api/generate":
          self._json(404, {"error": "unsupported handoff operation"})
          return
        length = int(self.headers.get("Content-Length", "0") or 0)
        try:
          payload = json.loads(self.rfile.read(length) or b"{}")
        except (ValueError, json.JSONDecodeError):
          self._json(400, {"error": "invalid JSON"})
          return
        keep_alive = payload.get("keep_alive")
        with _handoff_lock:
          idle, error = _wait_for_image_idle()
          if not idle:
            self._json(503, {"error": error, "gpu_handoff_blocked": True})
            return
          if keep_alive == 0:
            _image_lane_warm = False
            self._json(200, {"response": "", "done": True})
            return
          _image_lane_warm = True
          response = "READY" if str(payload.get("prompt") or "").strip() else ""
          self._json(200, {"response": response, "done": True})

      def do_GET(self) -> None:
        if self.path == "/healthz":
          self._json(200, {"status": "ok"})
          return
        self._proxy()

      def do_POST(self) -> None:
        self._proxy()

      def log_message(self, format_string: str, *args) -> None:
        print(f"gpu-handoff {self.address_string()} {format_string % args}", flush=True)


    if __name__ == "__main__":
      handoff = ThreadingHTTPServer((LISTEN_HOST, HANDOFF_PORT), HandoffHandler)
      threading.Thread(target=handoff.serve_forever, name="gpu-handoff", daemon=True).start()
      lan_api = ThreadingHTTPServer((LISTEN_HOST, LAN_LISTEN_PORT), LanHandler)
      threading.Thread(target=lan_api.serve_forever, name="lan-model-api", daemon=True).start()
      ThreadingHTTPServer((LISTEN_HOST, LISTEN_PORT), Handler).serve_forever()