353 lines
16 KiB
Python
353 lines
16 KiB
Python
#!/usr/bin/env python3
|
|
"""Ordinal-local signing and SCM mediation boundary for a model 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 collections.abc import Callable
|
|
from typing import Any
|
|
|
|
import cli_lane_goal
|
|
from execution_pool_protocol import (
|
|
MAX_WIRE_BYTES,
|
|
PROTOCOL_VERSION,
|
|
BoundedHTTPServer,
|
|
ProtocolError,
|
|
canonical_json,
|
|
parse_wire,
|
|
read_key,
|
|
sign_envelope,
|
|
verify_envelope,
|
|
)
|
|
from execution_pool_scm import Boundary as SCMBoundary
|
|
|
|
|
|
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"))
|
|
RESULT_FIELDS = frozenset(
|
|
{"status", "summary", "changed_files", "tests_run", "artifacts", "findings", "blockers"}
|
|
)
|
|
|
|
|
|
def _binding(value: dict[str, Any]) -> dict[str, Any]:
|
|
return {
|
|
name: value[name]
|
|
for name in ("board", "task_id", "run_id", "worker_ordinal", "attempt")
|
|
}
|
|
|
|
|
|
def _validate_result(payload: Any) -> dict[str, Any]:
|
|
if not isinstance(payload, dict):
|
|
raise ProtocolError("terminal result payload must be an object")
|
|
structured = payload.get("structured")
|
|
if not isinstance(structured, dict) or set(structured) != RESULT_FIELDS:
|
|
raise ProtocolError("terminal result fields do not match the reviewed schema")
|
|
if structured.get("status") not in cli_lane_goal.RESULT_STATUSES:
|
|
raise ProtocolError("terminal result status is invalid")
|
|
if not isinstance(structured.get("summary"), str) or not structured["summary"].strip():
|
|
raise ProtocolError("terminal result summary is required")
|
|
for name in RESULT_FIELDS - {"status", "summary"}:
|
|
value = structured.get(name)
|
|
if not isinstance(value, list) or any(not isinstance(item, str) for item in value):
|
|
raise ProtocolError(f"terminal result {name} must be a text list")
|
|
return payload
|
|
|
|
|
|
class ClientBoundary:
|
|
"""Keep HMAC and SCM authority outside the model-facing container."""
|
|
|
|
def __init__(self, key: bytes, scm: SCMBoundary | None = None):
|
|
self.key = key
|
|
self.scm = scm or SCMBoundary(key)
|
|
self.current: dict[str, Any] | None = None
|
|
self.resumed: dict[str, str] | None = None
|
|
self.resumed_evidence: dict[str, Any] | None = None
|
|
self.resume_transient = False
|
|
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")
|
|
return verify_envelope(self.key, json.loads(body))
|
|
|
|
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")
|
|
publication_only = isinstance(response.get("payload", {}).get("scm_resume"), dict)
|
|
checkout = {"workspace": "", "baseline_sha": ""} if publication_only else self.scm.checkout(response)
|
|
self.current = response
|
|
self.resumed = None
|
|
self.resumed_evidence = None
|
|
self.resume_transient = False
|
|
assignment = {
|
|
**_binding(response),
|
|
"payload": response["payload"],
|
|
"workspace": checkout["workspace"],
|
|
"baseline_sha": checkout["baseline_sha"],
|
|
"protocol_version": PROTOCOL_VERSION,
|
|
}
|
|
return {"assignment": assignment}
|
|
|
|
def _current_for(self, supplied: Any) -> tuple[dict[str, Any], dict[str, Any]]:
|
|
if not isinstance(supplied, dict):
|
|
raise ProtocolError("local request binding must be an object")
|
|
if self.current is None or supplied != _binding(self.current):
|
|
raise ProtocolError("local request does not own the current assignment")
|
|
return self.current, dict(supplied)
|
|
|
|
def _publication_lease(self, binding: dict[str, Any]) -> tuple[Callable[[], None], Callable[[], None]]:
|
|
"""Keep a resume lease alive outside the SCM lock and fail before another step."""
|
|
stop, failures = threading.Event(), []
|
|
|
|
def renew() -> None:
|
|
try:
|
|
response = self._post("/v1/heartbeat", sign_envelope(self.key, "heartbeat", binding, {"note": "SCM publication in progress"}))
|
|
if response["kind"] != "ack" or _binding(response) != binding or not response["payload"].get("accepted"):
|
|
raise ProtocolError("coordinator rejected SCM publication lease")
|
|
except (OSError, ProtocolError, ValueError, urllib.error.URLError) as error:
|
|
failures.append(error)
|
|
|
|
renew()
|
|
def keepalive() -> None:
|
|
while not stop.wait(20):
|
|
renew()
|
|
|
|
thread = threading.Thread(target=keepalive, daemon=True)
|
|
thread.start()
|
|
|
|
def checkpoint() -> None:
|
|
if failures:
|
|
raise ProtocolError("SCM publication lease was lost") from failures[0]
|
|
|
|
return checkpoint, lambda: (stop.set(), thread.join(timeout=1))
|
|
|
|
def heartbeat(self, request: dict[str, Any]) -> dict[str, Any]:
|
|
payload = request.get("payload")
|
|
if not isinstance(payload, dict):
|
|
raise ProtocolError("heartbeat payload must be an object")
|
|
with self.lock:
|
|
_assignment, binding = self._current_for(request.get("binding"))
|
|
response = self._post(
|
|
"/v1/heartbeat",
|
|
sign_envelope(self.key, "heartbeat", binding, payload),
|
|
)
|
|
if response["kind"] != "ack" or _binding(response) != binding:
|
|
raise ProtocolError("coordinator acknowledgement binding changed")
|
|
return {"ack": response["payload"]}
|
|
|
|
def _submit(
|
|
self, assignment: dict[str, Any], request: dict[str, Any],
|
|
structured: dict[str, Any],
|
|
) -> dict[str, str] | None:
|
|
"""Publish the run's work, downgrading the result rather than losing it.
|
|
|
|
A refused push or draft used to unwind the whole worker, so the run's
|
|
actual output was discarded and the task was blocked with only the SCM
|
|
error. The commits live on the durable workspace volume, so the result is
|
|
still reported here -- as blocked, carrying the reason -- which keeps the
|
|
prior work recoverable and reviewable by a human.
|
|
"""
|
|
try:
|
|
if isinstance(self.scm, SCMBoundary):
|
|
checkpoint, close = self._publication_lease(_binding(assignment))
|
|
try:
|
|
submission = self.scm.submit(assignment, request, checkpoint=checkpoint)
|
|
checkpoint()
|
|
finally:
|
|
close()
|
|
else:
|
|
submission = self.scm.submit(assignment, request)
|
|
except (OSError, ProtocolError, RuntimeError, ValueError) as error:
|
|
payload = assignment.get("payload")
|
|
resume = None
|
|
if isinstance(payload, dict) and payload.get("continuation_kind") == "repair":
|
|
try:
|
|
resume = self.scm.resume_artifact(assignment, request, structured)
|
|
except (OSError, ProtocolError, RuntimeError, ValueError):
|
|
resume = None
|
|
structured["status"] = "blocked"
|
|
reason = (
|
|
"Distributed SCM submission failed; the commits remain on this "
|
|
f"ordinal's workspace: {type(error).__name__}: {error}"
|
|
)
|
|
if reason not in structured["blockers"]:
|
|
structured["blockers"].append(reason)
|
|
return {"resume": resume} if resume is not None else None
|
|
pull = str(submission.get("pull_request") or "")
|
|
branch = str(submission.get("branch") or "")
|
|
for artifact in (pull, f"branch:{branch}" if branch else ""):
|
|
if artifact and artifact not in structured["artifacts"]:
|
|
structured["artifacts"].append(artifact)
|
|
head = str(submission.get("head") or "")
|
|
return {"branch": branch, "pull_request": pull, "head": head}
|
|
|
|
def resume(self, request: dict[str, Any]) -> dict[str, Any]:
|
|
"""Publish the coordinator-carried artifact without exposing it to the worker."""
|
|
with self.lock:
|
|
assignment, _binding = self._current_for(request.get("binding"))
|
|
artifact = assignment.get("payload", {}).get("scm_resume")
|
|
if not isinstance(artifact, dict):
|
|
raise ProtocolError("assignment has no SCM resume artifact")
|
|
evidence = _validate_result({"structured": artifact.get("structured")})["structured"]
|
|
if evidence["status"] != "completed":
|
|
raise ProtocolError("SCM resume evidence is not completed")
|
|
checkpoint, close = self._publication_lease(_binding)
|
|
try:
|
|
self.resume_transient = False
|
|
try:
|
|
self.resumed = self.scm.resume(assignment, artifact, checkpoint=checkpoint)
|
|
checkpoint()
|
|
except (OSError, RuntimeError, urllib.error.URLError) as error:
|
|
detail = str(error).lower()
|
|
if any(marker in detail for marker in ("timed out", "connection", "could not resolve", "http 502", "http 503", "http 504")):
|
|
self.resume_transient = True
|
|
return {"publication_retry_transient": True}
|
|
raise ProtocolError("SCM resume publication was rejected") from error
|
|
finally:
|
|
close()
|
|
self.resumed_evidence = json.loads(canonical_json(evidence))
|
|
return {"scm_submission": self.resumed}
|
|
|
|
def finish(self, request: dict[str, Any]) -> dict[str, Any]:
|
|
payload = _validate_result(request.get("payload"))
|
|
# The model-facing caller cannot classify a retry as transient. Only a
|
|
# preceding mediator resume may attach this coordinator control signal.
|
|
payload.pop("publication_retry_transient", None)
|
|
with self.lock:
|
|
assignment, binding = self._current_for(request.get("binding"))
|
|
structured = payload["structured"]
|
|
if self.resume_transient:
|
|
payload["publication_retry_transient"] = True
|
|
if structured["status"] == "completed" and int(payload.get("returncode", 1)) == 0:
|
|
if self.resumed is not None:
|
|
if self.resumed_evidence is None:
|
|
raise ProtocolError("SCM resume evidence is unavailable")
|
|
structured = json.loads(canonical_json(self.resumed_evidence))
|
|
payload["structured"] = structured
|
|
payload["returncode"] = 0
|
|
payload["scm_submission"] = self.resumed
|
|
else:
|
|
submission = self._submit(assignment, request, structured)
|
|
if submission is not None and "resume" not in submission:
|
|
# The mediator, not the model-facing request, records the
|
|
# broker-confirmed PR/branch in the signed terminal wire.
|
|
payload["scm_submission"] = submission
|
|
elif submission is not None:
|
|
# The exact clean commit remains on the durable worker volume.
|
|
# A broker refusal is execution infrastructure, not a provider
|
|
# capability verdict, so let the coordinator retry it after the
|
|
# scoped SCM condition is corrected.
|
|
payload["capacity_failure"] = True
|
|
payload["scm_resume"] = submission["resume"]
|
|
else:
|
|
payload["capacity_failure"] = True
|
|
response = self._post(
|
|
"/v1/result", sign_envelope(self.key, "result", binding, payload)
|
|
)
|
|
if response["kind"] != "ack" or _binding(response) != binding:
|
|
raise ProtocolError("coordinator acknowledgement binding changed")
|
|
if response["payload"].get("accepted"):
|
|
self.current = None
|
|
return {"ack": response["payload"], "structured": structured}
|
|
|
|
|
|
def handler_factory(boundary: ClientBoundary) -> type[BaseHTTPRequestHandler]:
|
|
class Handler(BaseHTTPRequestHandler):
|
|
server_version = f"hermes-execution-mediator/{PROTOCOL_VERSION}"
|
|
|
|
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
|
|
value = {"ready": True, "protocol_version": PROTOCOL_VERSION}
|
|
self._reply(200, value) 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
|
|
)
|
|
allowed = {"operation", "binding", "payload", "title", "body"}
|
|
if not request or set(request) - allowed:
|
|
raise ProtocolError("invalid local mediator request")
|
|
operation = str(request.get("operation") or "")
|
|
routes = {
|
|
"poll": boundary.poll,
|
|
"heartbeat": lambda: boundary.heartbeat(request),
|
|
"resume": lambda: boundary.resume(request),
|
|
"finish": lambda: boundary.finish(request),
|
|
}
|
|
if operation not in routes:
|
|
raise ProtocolError("unsupported local mediator operation")
|
|
self._reply(200, routes[operation]())
|
|
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__": # pragma: no cover - exercised by the container entrypoint
|
|
raise SystemExit(main())
|