atlas-iac/services/hermes/scripts/execution_pool_client.py

#!/usr/bin/env python3
"""Ordinal-local signing boundary for a model-facing Hermes worker."""

from __future__ import annotations

import json
import os
import threading
import urllib.error
import urllib.request
from http.server import BaseHTTPRequestHandler
from pathlib import Path
from typing import Any

from execution_pool_protocol import (
    MAX_WIRE_BYTES,
    BoundedHTTPServer,
    ProtocolError,
    canonical_json,
    parse_wire,
    read_key,
    sign_envelope,
    verify_envelope,
)


KEY_PATH = Path(os.environ.get("HERMES_EXECUTION_POOL_KEY_FILE", "/pool-access/execution-pool-key"))
COORDINATOR = os.environ.get(
    "HERMES_EXECUTION_POOL_URL", "http://hermes-execution-pool.hermes.svc.cluster.local:9007"
).rstrip("/")
ORDINAL = int(os.environ.get("HERMES_WORKER_ORDINAL", "-1"))
PORT = int(os.environ.get("HERMES_EXECUTION_CLIENT_PORT", "9009"))


def _binding(value: dict[str, Any]) -> dict[str, Any]:
    return {name: value[name] for name in (
        "board", "task_id", "run_id", "worker_ordinal", "attempt"
    )}


class ClientBoundary:
    """Hold the master key and fence local requests to the current ordinal run."""

    def __init__(self, key: bytes):
        self.key = key
        self.current: dict[str, Any] | None = None
        self.lock = threading.RLock()

    def _post(self, path: str, envelope: dict[str, Any]) -> dict[str, Any]:
        request = urllib.request.Request(
            COORDINATOR + path, data=canonical_json(envelope), method="POST",
            headers={"Content-Type": "application/json", "Cache-Control": "no-store"},
        )
        with urllib.request.urlopen(request, timeout=60) as response:
            body = response.read(MAX_WIRE_BYTES + 1)
        if len(body) > MAX_WIRE_BYTES:
            raise ProtocolError("coordinator response exceeds the wire limit")
        value = json.loads(body)
        return verify_envelope(self.key, value)

    def poll(self) -> dict[str, Any]:
        poll_binding = {
            "board": "", "task_id": "", "run_id": "",
            "worker_ordinal": ORDINAL, "attempt": 0,
        }
        response = self._post(
            "/v1/poll", sign_envelope(self.key, "poll", poll_binding, {"ready": True})
        )
        with self.lock:
            if response["kind"] == "ack":
                self.current = None
                return {"assignment": None}
            if response["kind"] != "assignment" or response["worker_ordinal"] != ORDINAL:
                raise ProtocolError("coordinator returned a foreign assignment")
            self.current = _binding(response)
        return {"assignment": response}

    def forward(self, kind: str, request: dict[str, Any]) -> dict[str, Any]:
        supplied = request.get("binding")
        payload = request.get("payload")
        if not isinstance(supplied, dict) or not isinstance(payload, dict):
            raise ProtocolError("local request binding and payload must be objects")
        with self.lock:
            if self.current is None or supplied != self.current:
                raise ProtocolError("local request does not own the current assignment")
            binding = dict(self.current)
        response = self._post(
            f"/v1/{kind}", sign_envelope(self.key, kind, binding, payload)
        )
        if response["kind"] != "ack" or _binding(response) != binding:
            raise ProtocolError("coordinator acknowledgement binding changed")
        if kind == "result" and response["payload"].get("accepted"):
            with self.lock:
                self.current = None
        return {"ack": response["payload"]}


def handler_factory(boundary: ClientBoundary) -> type[BaseHTTPRequestHandler]:
    class Handler(BaseHTTPRequestHandler):
        server_version = "hermes-execution-client/1"

        def _reply(self, status: int, value: dict[str, Any]) -> None:
            body = canonical_json(value)
            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 do_GET(self) -> None:  # noqa: N802
            self._reply(200, {"ready": True}) if self.path == "/ready" else self._reply(404, {"error": "not found"})

        def do_POST(self) -> None:  # noqa: N802
            try:
                length = int(self.headers.get("Content-Length", "0"))
                request = parse_wire(self.rfile.read(length)) if 0 < length <= MAX_WIRE_BYTES else None
                if not request or set(request) - {"operation", "binding", "payload"}:
                    raise ProtocolError("invalid local client request")
                operation = str(request.get("operation") or "")
                if operation == "poll":
                    result = boundary.poll()
                elif operation in {"heartbeat", "result"}:
                    result = boundary.forward(operation, request)
                else:
                    raise ProtocolError("unsupported local client operation")
                self._reply(200, result)
            except (ProtocolError, OSError, ValueError, urllib.error.URLError) as error:
                self._reply(409, {"error": str(error)[:2000]})

        def log_message(self, _format: str, *_arguments: Any) -> None:
            return

    return Handler


def main() -> int:
    if ORDINAL not in range(3):
        raise SystemExit("HERMES_WORKER_ORDINAL must be 0, 1, or 2")
    key = read_key(KEY_PATH)
    BoundedHTTPServer(
        ("0.0.0.0", PORT), handler_factory(ClientBoundary(key)), max_workers=4
    ).serve_forever()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())