From 55d7ae3f3fc5677e4c86aba6ecb5c6db4e897d1b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 9 Oct 2026 00:03:10 +0400 Subject: [PATCH 1/6] Add durable graph publication to the existing build service --- packages/microcosm-build/pyproject.toml | 2 + .../microcosm/build/graph_publication_cli.py | 60 +++ .../src/microcosm/build/telemetry_emitter.py | 36 +- .../v20261008_02_graph_publication.py | 27 ++ .../build/telemetry_emitter_service/auth.py | 161 ++++++++ .../telemetry_emitter_service/collector.py | 156 ++------ .../graph_publication.py | 375 ++++++++++++++++++ .../build/telemetry_emitter_service/main.py | 23 +- .../telemetry_emitter_service/migrations.py | 11 +- .../build/telemetry_emitter_service/models.py | 15 + .../telemetry_emitter_service/runtime.py | 62 ++- .../build/us_runtime/worker_identity.py | 2 +- .../shared/test_graph_publication.py | 245 ++++++++++++ .../shared/test_publication_workers.py | 105 +++++ .../shared/test_telemetry_emitter.py | 5 +- uv.lock | 2 + 16 files changed, 1129 insertions(+), 158 deletions(-) create mode 100644 packages/microcosm-build/src/microcosm/build/graph_publication_cli.py create mode 100644 packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/alembic/versions/v20261008_02_graph_publication.py create mode 100644 packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/auth.py create mode 100644 packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/graph_publication.py create mode 100644 packages/microcosm-build/tests/engine_free/shared/test_graph_publication.py create mode 100644 packages/microcosm-build/tests/engine_free/shared/test_publication_workers.py diff --git a/packages/microcosm-build/pyproject.toml b/packages/microcosm-build/pyproject.toml index c18475838..ed07ea12d 100644 --- a/packages/microcosm-build/pyproject.toml +++ b/packages/microcosm-build/pyproject.toml @@ -12,6 +12,7 @@ dependencies = [ "pandas>=2", "scipy>=1.13", "huggingface-hub>=0.20", + "httpx>=0.27,<1", "microcosm-frame>=0.1,<0.2", "microcosm-graph>=0.1,<0.2", "microcosm-fit>=0.1,<0.2", @@ -50,6 +51,7 @@ us = [ uk = ["policyengine-uk>=2.123.0", "microcosm-data>=0.1,<0.2", "h5py>=3", "tables>=3", "openpyxl>=3.1"] [project.scripts] +microcosm-publish-graph = "microcosm.build.graph_publication_cli:main" microcosm-build-uk = "microcosm.build.uk_runtime.full_build_cli:main" microcosm-export-us-l0-refit-h5 = "microcosm.build.us_runtime.l0_refit_export:main" diff --git a/packages/microcosm-build/src/microcosm/build/graph_publication_cli.py b/packages/microcosm-build/src/microcosm/build/graph_publication_cli.py new file mode 100644 index 000000000..95f6d2e20 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/graph_publication_cli.py @@ -0,0 +1,60 @@ +"""Retry preserved graph uploads without registering a telemetry run.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +from microcosm.build.telemetry_emitter import _cache_dir +from microcosm.build.telemetry_emitter_constants import TELEMETRY_SPOOL_FILENAME +from microcosm.build.telemetry_emitter_service.auth import ( + CollectorSession, + _collector_origin, +) +from microcosm.build.telemetry_emitter_service.constants import PRODUCTION_COLLECTOR_URL +from microcosm.build.telemetry_emitter_service.graph_publication import ( + DEFAULT_GRAPH_PUBLICATION_ORIGIN, + GraphPublicationDelivery, + GraphPublicationQueue, +) + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--spool", type=Path, default=_cache_dir() / TELEMETRY_SPOOL_FILENAME + ) + parser.add_argument("--publication-id", required=True) + parser.add_argument( + "--directory", + type=Path, + help="Optional preserved bundle containing orrery.upload.json; re-enqueues after moving machines.", + ) + parser.add_argument("--origin", default=DEFAULT_GRAPH_PUBLICATION_ORIGIN) + args = parser.parse_args(argv) + queue = GraphPublicationQueue(args.spool) + if args.directory is not None: + inventory = json.loads((args.directory / "orrery.upload.json").read_text()) + if inventory["publication_id"] != args.publication_id: + parser.error("The preserved inventory has a different publication ID.") + queue.enqueue(args.directory, inventory) + if queue.receipt(args.publication_id) is None: + parser.error("Publication job not found; supply its preserved directory.") + queue.retry(args.publication_id) + session = CollectorSession(_collector_origin(PRODUCTION_COLLECTOR_URL)) + delivery = GraphPublicationDelivery( + queue, + credential=session.credential, + origin=args.origin, + invalidate_credential=session.invalidate, + ) + # Explicit retries process this job only, not unrelated pending jobs. + delivery.flush_once(publication_id=args.publication_id) + receipt = queue.receipt(args.publication_id) + print(json.dumps(receipt, sort_keys=True)) + return 0 if receipt["status"] == "published" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter.py index c9fb5aaf7..160f33b09 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter.py @@ -122,12 +122,14 @@ def __init__( socket_path: Path | None, runtime_dir: Path | None, send_timeout_seconds: float = DEFAULT_SEND_TIMEOUT_SECONDS, + spool_path: Path | None = None, ) -> None: self.run = run self._process = process self._socket_path = socket_path self._runtime_dir = runtime_dir self._send_timeout_seconds = send_timeout_seconds + self._spool_path = spool_path self._closed = False self._warned = False self._transition_stage: str | None = None @@ -215,6 +217,7 @@ def start( process=process, socket_path=socket_path, runtime_dir=runtime_dir, + spool_path=queue_path, ) emitter.emit( event_type=EVENT_TYPE_RUN, @@ -246,7 +249,13 @@ def start( runtime_dir.rmdir() except OSError: pass - return cls(run=run, process=None, socket_path=None, runtime_dir=runtime_dir) + return cls( + run=run, + process=None, + socket_path=None, + runtime_dir=runtime_dir, + spool_path=Path(spool_path) if spool_path else None, + ) @property def available(self) -> bool: @@ -450,6 +459,31 @@ def close(self) -> None: self._send({"action": ACTION_CLOSE}) self._closed = True + def publish_graph( + self, directory: Path, inventory: dict, *, wait_seconds: float = 30.0 + ) -> dict: + """Preserve a job first, then wait a bounded time for its first result. + + Unlike telemetry, missing credentials never discard this job. Enqueue + directly into the service's shared SQLite database so a socket timeout + or unavailable service cannot lose the publication request. + """ + from microcosm.build.telemetry_emitter_service.graph_publication import ( + GraphPublicationQueue, + ) + + queue = GraphPublicationQueue( + self._spool_path or _cache_dir() / TELEMETRY_SPOOL_FILENAME + ) + queue.enqueue(directory, inventory) + deadline = time.monotonic() + max(0, wait_seconds) + while self.available and time.monotonic() < deadline: + receipt = queue.receipt(inventory["publication_id"]) + if receipt.get("attempts", 0) or receipt["status"] == "published": + return receipt + time.sleep(0.1) + return queue.receipt(inventory["publication_id"]) + def _close_transition_stage( self, *, diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/alembic/versions/v20261008_02_graph_publication.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/alembic/versions/v20261008_02_graph_publication.py new file mode 100644 index 000000000..14ae0318a --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/alembic/versions/v20261008_02_graph_publication.py @@ -0,0 +1,27 @@ +"""Add durable graph publication jobs without changing event retention.""" + +import sqlalchemy as sa +from alembic import op + +revision = "20261008_02" +down_revision = "20261007_01" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.create_table( + "graph_publication_jobs", + sa.Column("publication_id", sa.Text(), primary_key=True), + sa.Column("inventory", sa.Text(), nullable=False), + sa.Column("directory", sa.Text(), nullable=False), + sa.Column("status", sa.Text(), nullable=False), + sa.Column("attempts", sa.Integer(), nullable=False), + sa.Column("next_attempt_at", sa.Float(), nullable=False), + sa.Column("lease_until", sa.Float(), nullable=False), + sa.Column("receipt", sa.Text(), nullable=False), + ) + + +def downgrade() -> None: + op.drop_table("graph_publication_jobs") diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/auth.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/auth.py new file mode 100644 index 000000000..a8e0726f7 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/auth.py @@ -0,0 +1,161 @@ +"""HF-to-collector session shared by telemetry and graph publication.""" + +from __future__ import annotations + +import json +import os +import time +import urllib.error +import urllib.request +from collections.abc import Mapping +from http import HTTPStatus +from typing import Any +from urllib.parse import urlsplit + +from huggingface_hub import get_token + +from microcosm.build.telemetry_emitter_service.constants import ( + COLLECTOR_RESPONSE_TOO_LARGE_ERROR, + COLLECTOR_URL_HTTPS_ERROR, + COLLECTOR_URL_ORIGIN_ERROR, + DEFAULT_TOKEN_LIFETIME_SECONDS, + DEVELOPMENT_COLLECTOR_LOOPBACK_ERROR, + HTTP_TIMEOUT_SECONDS, + HTTP_USER_AGENT, + LOOPBACK_HOSTS, + MAX_HTTP_RESPONSE_BYTES, + MINIMUM_TOKEN_LIFETIME_SECONDS, + TOKEN_EXCHANGE_PATH, + TOKEN_REFRESH_MARGIN_SECONDS, +) + + +class _NoRedirectHandler(urllib.request.HTTPRedirectHandler): + """Keep bearer credentials on the explicitly configured origin.""" + + def redirect_request(self, req, fp, code, msg, headers, newurl): + return None + + +def _collector_origin(value: str, *, allow_loopback_http: bool = False) -> str: + parsed = urlsplit(value) + loopback = parsed.hostname in LOOPBACK_HOSTS + valid_scheme = parsed.scheme == "https" or ( + allow_loopback_http and loopback and parsed.scheme == "http" + ) + if not parsed.hostname or not valid_scheme: + raise ValueError(COLLECTOR_URL_HTTPS_ERROR) + if ( + parsed.username + or parsed.password + or parsed.query + or parsed.fragment + or parsed.path not in {"", "/"} + ): + raise ValueError(COLLECTOR_URL_ORIGIN_ERROR) + return value.rstrip("/") + + +def _development_collector_url(value: str) -> str: + parsed = urlsplit(value) + if parsed.hostname not in LOOPBACK_HOSTS: + raise ValueError(DEVELOPMENT_COLLECTOR_LOOPBACK_ERROR) + return _collector_origin(value, allow_loopback_http=True) + + +def _decode_response(body: bytes) -> dict[str, Any]: + try: + response = json.loads(body) if body else {} + except (json.JSONDecodeError, UnicodeDecodeError): + return {} + return response if isinstance(response, dict) else {} + + +def _http_post( + url: str, + payload: Mapping[str, Any], + bearer_token: str, + *, + timeout: float = HTTP_TIMEOUT_SECONDS, +) -> tuple[int, dict[str, Any]]: + request = urllib.request.Request( + url, + data=json.dumps(payload, separators=(",", ":")).encode(), + headers={ + "Authorization": f"Bearer {bearer_token}", + "Content-Type": "application/json", + "User-Agent": HTTP_USER_AGENT, + }, + method="POST", + ) + try: + opener = urllib.request.build_opener(_NoRedirectHandler) + with opener.open(request, timeout=timeout) as response: + body = response.read(MAX_HTTP_RESPONSE_BYTES + 1) + if len(body) > MAX_HTTP_RESPONSE_BYTES: + raise OSError(COLLECTOR_RESPONSE_TOO_LARGE_ERROR) + return response.status, _decode_response(body) + except urllib.error.HTTPError as error: + body = error.read(MAX_HTTP_RESPONSE_BYTES + 1) + if len(body) > MAX_HTTP_RESPONSE_BYTES: + return error.code, {} + return error.code, _decode_response(body) + + +def _huggingface_token() -> str | None: + return ( + os.environ.get("HF_TOKEN", "").strip() + or os.environ.get("HUGGINGFACE_TOKEN", "").strip() + or get_token() + ) + + +class CollectorSession: + """One thread-safe short-lived credential shared by independent publishers.""" + + def __init__(self, origin: str, *, post=None, credential=None): + import threading + + self.origin = origin + self._post = post or _http_post + self._credential = credential or _huggingface_token + self._session_token: tuple[str, float] | None = None + self._lock = threading.Lock() + + def invalidate(self) -> None: + with self._lock: + self._session_token = None + + def credential(self) -> tuple[str, str | None]: + with self._lock: + if ( + self._session_token is not None + and self._session_token[1] + > time.monotonic() + TOKEN_REFRESH_MARGIN_SECONDS + ): + return "ok", self._session_token[0] + hf_token = self._credential() + if not hf_token: + return "missing", None + try: + status, response = self._post( + self.origin + TOKEN_EXCHANGE_PATH, {}, hf_token + ) + except (OSError, TimeoutError): + return "unavailable", None + if status == HTTPStatus.OK and isinstance( + response.get("access_token"), str + ): + try: + expires_in = max( + MINIMUM_TOKEN_LIFETIME_SECONDS, + int(response.get("expires_in", DEFAULT_TOKEN_LIFETIME_SECONDS)), + ) + except (ValueError, TypeError): + return "unavailable", None + token = response["access_token"] + self._session_token = (token, time.monotonic() + expires_in) + return "ok", token + if status in {HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN}: + return "rejected", None + return "unavailable", None diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/collector.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/collector.py index 0d7e4d294..1b96dfab7 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/collector.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/collector.py @@ -1,130 +1,38 @@ -"""Collector authentication and best-effort event delivery.""" +"""Best-effort telemetry delivery; shared authentication lives outside this component.""" from __future__ import annotations -import json -import os import sys import time -import urllib.error -import urllib.request from collections.abc import Mapping from http import HTTPStatus from typing import Any -from urllib.parse import urlsplit - -from huggingface_hub import get_token +from microcosm.build.telemetry_emitter_service.auth import ( + CollectorSession, + _collector_origin, + _development_collector_url, + _http_post, + _huggingface_token, +) from microcosm.build.telemetry_emitter_service.constants import ( - COLLECTOR_RESPONSE_TOO_LARGE_ERROR, - COLLECTOR_URL_HTTPS_ERROR, - COLLECTOR_URL_ORIGIN_ERROR, - DEFAULT_TOKEN_LIFETIME_SECONDS, - DEVELOPMENT_COLLECTOR_LOOPBACK_ERROR, - HTTP_TIMEOUT_SECONDS, - HTTP_USER_AGENT, INITIAL_RETRY_SECONDS, LOCAL_ONLY_MISSING_CREDENTIAL, LOCAL_ONLY_REJECTED_COLLECTOR_AUTHORIZATION, LOCAL_ONLY_REJECTED_CREDENTIAL, LOCAL_ONLY_REJECTED_EVENTS, LOCAL_ONLY_REJECTED_REGISTRATION, - LOOPBACK_HOSTS, - MAX_HTTP_RESPONSE_BYTES, MAX_RETRY_SECONDS, - MINIMUM_TOKEN_LIFETIME_SECONDS, NO_CREDENTIAL_MESSAGE, PRODUCTION_COLLECTOR_URL, REJECTED_CREDENTIAL_MESSAGE, REJECTED_EVENTS_MESSAGE, RUN_EVENTS_PATH_TEMPLATE, RUN_REGISTRATION_PATH, - TOKEN_EXCHANGE_PATH, - TOKEN_REFRESH_MARGIN_SECONDS, ) from microcosm.build.telemetry_emitter_service.spool import EventSpool -class _NoRedirectHandler(urllib.request.HTTPRedirectHandler): - """Keep bearer credentials on the explicitly configured origin.""" - - def redirect_request(self, req, fp, code, msg, headers, newurl): - return None - - -def _collector_origin(value: str, *, allow_loopback_http: bool = False) -> str: - parsed = urlsplit(value) - loopback = parsed.hostname in LOOPBACK_HOSTS - valid_scheme = parsed.scheme == "https" or ( - allow_loopback_http and loopback and parsed.scheme == "http" - ) - if not parsed.hostname or not valid_scheme: - raise ValueError(COLLECTOR_URL_HTTPS_ERROR) - if ( - parsed.username - or parsed.password - or parsed.query - or parsed.fragment - or parsed.path not in {"", "/"} - ): - raise ValueError(COLLECTOR_URL_ORIGIN_ERROR) - return value.rstrip("/") - - -def _development_collector_url(value: str) -> str: - parsed = urlsplit(value) - if parsed.hostname not in LOOPBACK_HOSTS: - raise ValueError(DEVELOPMENT_COLLECTOR_LOOPBACK_ERROR) - return _collector_origin(value, allow_loopback_http=True) - - -def _decode_response(body: bytes) -> dict[str, Any]: - try: - response = json.loads(body) if body else {} - except (json.JSONDecodeError, UnicodeDecodeError): - return {} - return response if isinstance(response, dict) else {} - - -def _http_post( - url: str, - payload: Mapping[str, Any], - bearer_token: str, - *, - timeout: float = HTTP_TIMEOUT_SECONDS, -) -> tuple[int, dict[str, Any]]: - request = urllib.request.Request( - url, - data=json.dumps(payload, separators=(",", ":")).encode(), - headers={ - "Authorization": f"Bearer {bearer_token}", - "Content-Type": "application/json", - "User-Agent": HTTP_USER_AGENT, - }, - method="POST", - ) - try: - opener = urllib.request.build_opener(_NoRedirectHandler) - with opener.open(request, timeout=timeout) as response: - body = response.read(MAX_HTTP_RESPONSE_BYTES + 1) - if len(body) > MAX_HTTP_RESPONSE_BYTES: - raise OSError(COLLECTOR_RESPONSE_TOO_LARGE_ERROR) - return response.status, _decode_response(body) - except urllib.error.HTTPError as error: - body = error.read(MAX_HTTP_RESPONSE_BYTES + 1) - if len(body) > MAX_HTTP_RESPONSE_BYTES: - return error.code, {} - return error.code, _decode_response(body) - - -def _huggingface_token() -> str | None: - return ( - os.environ.get("HF_TOKEN", "").strip() - or os.environ.get("HUGGINGFACE_TOKEN", "").strip() - or get_token() - ) - - #: Client errors worth retrying: a timeout and a rate limit pass with time. #: Every other 4xx is the collector's settled answer to the same request, so #: retrying it only wedges the queue (a payload shape it does not accept, say). @@ -149,6 +57,7 @@ def __init__( spool: EventSpool, *, development_collector_url: str | None = None, + session: CollectorSession | None = None, ) -> None: self.collector_url = ( _development_collector_url(development_collector_url) @@ -156,7 +65,11 @@ def __init__( else _collector_origin(PRODUCTION_COLLECTOR_URL) ) self.spool = spool - self._session_token: tuple[str, float] | None = None + self.session = session or CollectorSession( + self.collector_url, + post=lambda *args: _http_post(*args), + credential=lambda: _huggingface_token(), + ) self._registered: set[str] = set() self._warned_no_token = False self._warned_denied: set[str] = set() @@ -195,7 +108,7 @@ def flush_once(self) -> bool: made_progress = True self._retry_seconds = INITIAL_RETRY_SECONDS elif status == HTTPStatus.UNAUTHORIZED: - self._session_token = None + self.session.invalidate() elif status == HTTPStatus.FORBIDDEN: self._make_local_only( registration, @@ -212,40 +125,21 @@ def flush_once(self) -> bool: return made_progress def _collector_token(self, registration: Mapping[str, Any]) -> str | None: - if ( - self._session_token is not None - and self._session_token[1] > time.monotonic() + TOKEN_REFRESH_MARGIN_SECONDS - ): - return self._session_token[0] - hf_token = _huggingface_token() - if not hf_token: + result, token = self.session_credential() + if result == "missing": if not self._warned_no_token: print(NO_CREDENTIAL_MESSAGE, file=sys.stderr, flush=True) self._warned_no_token = True self._make_local_only(registration, LOCAL_ONLY_MISSING_CREDENTIAL) - return None - try: - status, response = _http_post( - self.collector_url + TOKEN_EXCHANGE_PATH, - {}, - hf_token, - ) - except (OSError, TimeoutError): - self._defer_retry() - return None - if status == HTTPStatus.OK and isinstance(response.get("access_token"), str): - expires_in = max( - MINIMUM_TOKEN_LIFETIME_SECONDS, - int(response.get("expires_in", DEFAULT_TOKEN_LIFETIME_SECONDS)), - ) - token = response["access_token"] - self._session_token = (token, time.monotonic() + expires_in) - return token - if status in {HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN}: + elif result == "rejected": self._make_local_only(registration, LOCAL_ONLY_REJECTED_CREDENTIAL) - else: + elif result != "ok": self._defer_retry() - return None + return token + + def session_credential(self) -> tuple[str, str | None]: + """Compatibility accessor; authentication has no telemetry queue policy.""" + return self.session.credential() def _register(self, registration: Mapping[str, Any], token: str) -> bool: registration_key = _registration_key(registration) @@ -263,7 +157,7 @@ def _register(self, registration: Mapping[str, Any], token: str) -> bool: self._retry_seconds = INITIAL_RETRY_SECONDS return True if status == HTTPStatus.UNAUTHORIZED: - self._session_token = None + self.session.invalidate() elif status in {HTTPStatus.FORBIDDEN, HTTPStatus.CONFLICT} or ( _permanent_client_error(status) ): diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/graph_publication.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/graph_publication.py new file mode 100644 index 000000000..17072ed49 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/graph_publication.py @@ -0,0 +1,375 @@ +"""Durable, retryable immutable graph uploads owned by the emitter service. + +Graph bytes stay in preserved files, never in event details. Pending graph jobs +are not subject to the telemetry event queue's time or size pruning. +""" + +from __future__ import annotations + +import hashlib +import json +import re +import time +import uuid +from collections.abc import Callable, Mapping +from pathlib import Path +from urllib.parse import urlsplit + +import httpx +from sqlalchemy import select, update + +from microcosm.build.telemetry_emitter_service.database import ( + create_spool_engine, + create_spool_session_factory, +) +from microcosm.build.telemetry_emitter_service.migrations import upgrade_spool_database +from microcosm.build.telemetry_emitter_service.models import GraphPublicationJob + +MAX_GRAPH_FILE_BYTES = 64 * 1024 * 1024 +MAX_GRAPH_TOTAL_BYTES = 512 * 1024 * 1024 +GRAPH_PUBLICATION_ACTION = "graph_publication" +GRAPH_PUBLICATION_WAIT_SECONDS = 30.0 +GRAPH_PUBLICATION_LEASE_SECONDS = 300.0 +GRAPH_PUBLICATION_STATUS_FILENAME = "publication.status.json" +DEFAULT_GRAPH_PUBLICATION_ORIGIN = "https://microcosm-runs.vercel.app" +GRAPH_PUBLICATION_ROLES = { + "graph": r"graph\.orrery\.json", + "schema": r"graph\.schema\.json", + "index": r"execution\.evidence\.json", + "declaration": r"evidence-[A-Za-z0-9_-]+-graph\.json", + "manifest": r"evidence-[A-Za-z0-9_-]+-manifest\.json", + "binding": r"evidence-[A-Za-z0-9_-]+-binding\.json", + "summaries": r"evidence-[A-Za-z0-9_-]+-summaries\.json", + "upstream": r"upstream-[a-f0-9]{64}\.json", +} + + +def _digest(payload: bytes) -> str: + return hashlib.sha256(payload).hexdigest() + + +def inventory_digest(inventory: Mapping) -> str: + """Digest the versioned field order shared with the publication API.""" + inventory = { + "version": inventory["version"], + "publication_id": inventory["publication_id"], + "files": [ + {key: item[key] for key in ("name", "role", "bytes", "sha256")} + for item in sorted(inventory["files"], key=lambda file: file["name"]) + ], + } + return _digest( + json.dumps(inventory, separators=(",", ":"), ensure_ascii=True).encode() + ) + + +def publication_inventory( + directory: Path, roles: Mapping[str, str], *, publication_id: str | None = None +) -> dict: + """Hash an explicit allowlist of files; publication identity is not a run ID.""" + files = [] + for name, role in sorted(roles.items()): + if role not in GRAPH_PUBLICATION_ROLES or not re.fullmatch( + GRAPH_PUBLICATION_ROLES[role], name + ): + raise ValueError("Ineligible graph publication filename or role.") + path = Path(directory) / name + if path.is_symlink() or not path.is_file(): + raise ValueError( + "Graph publication files must be regular, preserved files." + ) + size = path.stat().st_size + if not 0 < size <= MAX_GRAPH_FILE_BYTES: + raise ValueError("Graph publication file exceeds its size limit.") + files.append( + { + "name": name, + "role": role, + "bytes": size, + "sha256": _digest(path.read_bytes()), + } + ) + if ( + not files + or len(files) > 256 + or sum(f["bytes"] for f in files) > MAX_GRAPH_TOTAL_BYTES + or sum(f["role"] == "graph" for f in files) != 1 + ): + raise ValueError("Invalid graph publication inventory.") + id_ = publication_id or uuid.uuid4().hex + if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_-]{0,127}", id_): + raise ValueError("Invalid graph publication ID.") + return {"version": 1, "publication_id": id_, "files": files} + + +class GraphPublicationQueue: + """SQLite job records with leases; retained until publication or operator action.""" + + def __init__(self, path: Path): + self.path = Path(path) + self.path.parent.mkdir(parents=True, exist_ok=True) + self._engine = create_spool_engine(self.path) + upgrade_spool_database(self._engine) + self._sessions = create_spool_session_factory(self._engine) + + def enqueue(self, directory: Path, inventory: dict) -> None: + """Durably register exact bytes without a telemetry run or credential.""" + directory = Path(directory).resolve() + expected = publication_inventory( + directory, + {f["name"]: f["role"] for f in inventory["files"]}, + publication_id=inventory["publication_id"], + ) + if expected != inventory: + raise ValueError("Graph publication digest mismatch.") + id_ = inventory["publication_id"] + with self._sessions.begin() as session: + existing = session.get(GraphPublicationJob, id_) + if existing is not None: + if existing.inventory != inventory or existing.directory != str( + directory + ): + raise ValueError( + "Graph publication identity conflicts with a queued job." + ) + return + session.add( + GraphPublicationJob( + publication_id=id_, + inventory=inventory, + directory=str(directory), + status="pending", + attempts=0, + next_attempt_at=0, + lease_until=0, + receipt={ + "version": 1, + "publication_id": id_, + "status": "pending", + "error_code": None, + }, + ) + ) + + def claim( + self, + *, + lease_seconds: float = GRAPH_PUBLICATION_LEASE_SECONDS, + publication_id: str | None = None, + ) -> dict | None: + """Atomically lease one eligible job across concurrent emitter processes.""" + now = time.time() + with self._sessions.begin() as session: + statement = select(GraphPublicationJob.publication_id) + if publication_id is not None: + statement = statement.where( + GraphPublicationJob.publication_id == publication_id + ) + id_ = session.scalar( + statement.where( + GraphPublicationJob.status == "pending", + GraphPublicationJob.next_attempt_at <= now, + GraphPublicationJob.lease_until <= now, + ).limit(1) + ) + if id_ is None: + return None + result = session.execute( + update(GraphPublicationJob) + .where( + GraphPublicationJob.publication_id == id_, + GraphPublicationJob.status == "pending", + GraphPublicationJob.lease_until <= now, + ) + .values(lease_until=now + lease_seconds) + ) + if not result.rowcount: + return None + job = session.get(GraphPublicationJob, id_) + return { + "publication_id": id_, + "directory": job.directory, + "inventory": job.inventory, + "lease_until": job.lease_until, + } + + def receipt(self, id_: str) -> dict | None: + """Read the latest local result; old HF receipt snapshots remain unchanged.""" + with self._sessions() as session: + job = session.get(GraphPublicationJob, id_) + return None if job is None else dict(job.receipt) + + def finish( + self, + id_: str, + *, + error_code: str | None = None, + url: str | None = None, + expected_lease: float | None = None, + ) -> None: + with self._sessions.begin() as session: + statement = update(GraphPublicationJob).where( + GraphPublicationJob.publication_id == id_, + GraphPublicationJob.status == "pending", + ) + if expected_lease is not None: + statement = statement.where( + GraphPublicationJob.lease_until == expected_lease + ) + job = session.scalar( + statement.values( + attempts=GraphPublicationJob.attempts + 1, + status="published" if url else "pending", + lease_until=0, + ).returning(GraphPublicationJob) + ) + # Completion is conditional in SQL, not a read-then-write check: + # expired workers cannot release another lease or replace success. + if job is None: + return + job.next_attempt_at = time.time() + min(300, 2 ** min(job.attempts, 8)) + job.receipt = { + "version": 1, + "publication_id": id_, + "status": job.status, + "error_code": error_code, + "url": url, + "attempts": job.attempts, + "inventory_sha256": inventory_digest(job.inventory), + } + # This is a separate latest-result file, not the build/HF receipt. + path = Path(job.directory) / GRAPH_PUBLICATION_STATUS_FILENAME + temporary = path.with_suffix(f".{uuid.uuid4().hex}.tmp") + temporary.write_text(json.dumps(job.receipt, sort_keys=True) + "\n") + temporary.replace(path) + + def retry(self, id_: str | None = None) -> None: + """Make preserved pending jobs eligible again, including missing credentials.""" + statement = update(GraphPublicationJob).where( + GraphPublicationJob.status == "pending" + ) + if id_ is not None: + statement = statement.where(GraphPublicationJob.publication_id == id_) + with self._sessions.begin() as session: + session.execute(statement.values(next_attempt_at=0)) + + +class GraphPublicationDelivery: + """Upload exact files with scoped signed URLs, then seal the publication.""" + + def __init__( + self, + queue: GraphPublicationQueue, + *, + credential: Callable[[], tuple[str, str | None]], + origin: str, + client: httpx.Client | None = None, + invalidate_credential: Callable[[], None] | None = None, + ): + parsed = urlsplit(origin) + if parsed.scheme != "https" and not ( + parsed.scheme == "http" + and parsed.hostname in {"localhost", "127.0.0.1", "::1"} + ): + raise ValueError( + "Graph publication origin must be HTTPS or local development." + ) + if ( + parsed.username + or parsed.password + or parsed.query + or parsed.fragment + or parsed.path not in {"", "/"} + ): + raise ValueError("Graph publication destination must be an origin.") + self.queue, self.credential, self.origin = queue, credential, origin.rstrip("/") + self.client = client or httpx.Client(timeout=60, follow_redirects=False) + self.invalidate_credential = invalidate_credential + + def flush_once(self, *, publication_id: str | None = None) -> bool: + """Try one job; failures retain it without changing telemetry eligibility.""" + job = self.queue.claim(publication_id=publication_id) + if job is None: + return False + id_, inventory, directory = ( + job["publication_id"], + job["inventory"], + Path(job["directory"]), + ) + error = "publication_unavailable" + lease = job["lease_until"] + try: + status, token = self.credential() + if token is None: + self.queue.finish( + id_, error_code=f"credential_{status}", expected_lease=lease + ) + return False + current = publication_inventory( + directory, + {f["name"]: f["role"] for f in inventory["files"]}, + publication_id=id_, + ) + if current != inventory: + raise ValueError("preserved_bytes_changed") + headers = {"Authorization": f"Bearer {token}"} + response = self.client.post( + self.origin + "/api/uploads/", json=inventory, headers=headers + ) + response.raise_for_status() + body = response.json() + uploads = body.get("uploads") + if not isinstance(uploads, list): + raise ValueError("invalid_upload_response") + files = {f["name"]: f for f in inventory["files"]} + for upload in uploads: + name = upload.get("name") + if name not in files: + raise ValueError("invalid_upload_response") + if upload.get("already_uploaded"): + continue + url = urlsplit(upload["url"]) + if ( + url.scheme != "https" + or not ( + (url.hostname == "vercel.com" and url.path == "/api/blob/") + or url.hostname == "blob.vercel-storage.com" + or (url.hostname or "").endswith(".blob.vercel-storage.com") + ) + or url.username + or url.password + ): + raise ValueError("invalid_upload_url") + with (directory / name).open("rb") as stream: + result = self.client.put( + upload["url"], + content=stream, + headers={ + "Content-Type": "application/json", + "Content-Length": str(files[name]["bytes"]), + }, + ) + result.raise_for_status() + response = self.client.post( + self.origin + "/api/publications/", json=inventory, headers=headers + ) + response.raise_for_status() + result = response.json() + if result.get("publication_id") != id_ or result.get( + "inventory_sha256" + ) != inventory_digest(inventory): + raise ValueError("invalid_publication_response") + self.queue.finish( + id_, url=f"{self.origin}/runs/{id_}/", expected_lease=lease + ) + return True + except httpx.HTTPStatusError as exception: + error = f"publication_http_{exception.response.status_code}" + if exception.response.status_code == 401 and self.invalidate_credential: + self.invalidate_credential() + except httpx.HTTPError: + error = "publication_network_error" + except (OSError, ValueError, KeyError, TypeError): + error = "publication_invalid_evidence_or_response" + self.queue.finish(id_, error_code=error, expected_lease=lease) + return False diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/main.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/main.py index 2872bf91a..51d18fa74 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/main.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/main.py @@ -10,6 +10,11 @@ from microcosm.build.telemetry_emitter_service.constants import ( DEFAULT_HEARTBEAT_SECONDS, ) +from microcosm.build.telemetry_emitter_service.graph_publication import ( + DEFAULT_GRAPH_PUBLICATION_ORIGIN, + GraphPublicationDelivery, + GraphPublicationQueue, +) from microcosm.build.telemetry_emitter_service.resources import ProcessTreeSampler from microcosm.build.telemetry_emitter_service.runtime import EmitterService from microcosm.build.telemetry_emitter_service.spool import EventSpool @@ -22,7 +27,7 @@ def build_parser() -> argparse.ArgumentParser: parser.add_argument("--socket", type=Path, required=True) parser.add_argument("--spool", type=Path, required=True) parser.add_argument("--development-collector-url") - parser.add_argument("--registration-json", required=True) + parser.add_argument("--registration-json", help="Omit for graph-only delivery.") parser.add_argument("--parent-pid", type=int, required=True) parser.add_argument( "--heartbeat-seconds", @@ -36,15 +41,23 @@ def main(argv: list[str] | None = None) -> int: """Run the telemetry emitter service until the build disconnects.""" args = build_parser().parse_args(argv) - registration = json.loads(args.registration_json) + registration = ( + json.loads(args.registration_json) if args.registration_json else None + ) spool = EventSpool(args.spool) + delivery = CollectorDelivery( + spool, development_collector_url=args.development_collector_url + ) service = EmitterService( socket_path=args.socket, registration=registration, spool=spool, - delivery=CollectorDelivery( - spool, - development_collector_url=args.development_collector_url, + delivery=delivery, + graph_delivery=GraphPublicationDelivery( + GraphPublicationQueue(args.spool), + credential=delivery.session.credential, + invalidate_credential=delivery.session.invalidate, + origin=DEFAULT_GRAPH_PUBLICATION_ORIGIN, ), sampler=ProcessTreeSampler(args.parent_pid), heartbeat_seconds=args.heartbeat_seconds, diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/migrations.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/migrations.py index c1c69c61c..8446a5a2f 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/migrations.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/migrations.py @@ -2,6 +2,7 @@ from __future__ import annotations +import fcntl from collections.abc import Iterator from contextlib import contextmanager from importlib.resources import as_file, files @@ -37,9 +38,13 @@ def alembic_config( def upgrade_spool_database(engine: Engine) -> None: """Apply every committed spool migration to an engine.""" - with engine.begin() as connection: - with alembic_config(connection=connection) as config: - command.upgrade(config, _MIGRATION_TARGET) + # Concurrent build processes share this spool and must serialize upgrades. + path = Path(engine.url.database) + with path.with_suffix(".migration.lock").open("a") as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + with engine.begin() as connection: + with alembic_config(connection=connection) as config: + command.upgrade(config, _MIGRATION_TARGET) def current_database_revision(path: Path | str) -> str | None: diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/models.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/models.py index 87f1c678d..d49bb8277 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/models.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/models.py @@ -6,6 +6,7 @@ from typing import Any from sqlalchemy import ( + Float, ForeignKeyConstraint, Index, Integer, @@ -128,3 +129,17 @@ class TelemetryEventRecord(SpoolModel): ) created_at: Mapped[str] = mapped_column(Text, nullable=False) run: Mapped[TelemetryRunRecord] = relationship(back_populates="events") + + +class GraphPublicationJob(SpoolModel): + """Durable artifact delivery, independent of event retention and run IDs.""" + + __tablename__ = "graph_publication_jobs" + publication_id: Mapped[str] = mapped_column(Text, primary_key=True) + inventory: Mapped[dict[str, Any]] = mapped_column(JsonObjectText, nullable=False) + directory: Mapped[str] = mapped_column(Text, nullable=False) + status: Mapped[str] = mapped_column(Text, nullable=False) + attempts: Mapped[int] = mapped_column(Integer, nullable=False) + next_attempt_at: Mapped[float] = mapped_column(Float, nullable=False) + lease_until: Mapped[float] = mapped_column(Float, nullable=False) + receipt: Mapped[dict[str, Any]] = mapped_column(JsonObjectText, nullable=False) diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/runtime.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/runtime.py index 0bb1d9c2a..3794d65a1 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/runtime.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/runtime.py @@ -1,4 +1,4 @@ -"""Unix-socket runtime for the telemetry emitter service.""" +"""Existing local service with independent telemetry and graph workers.""" from __future__ import annotations @@ -25,6 +25,10 @@ UNSUPPORTED_ACTION_ERROR, WORKER_INTERVAL_SECONDS, ) +from microcosm.build.telemetry_emitter_service.graph_publication import ( + GRAPH_PUBLICATION_ACTION, + GraphPublicationDelivery, +) from microcosm.build.telemetry_emitter_service.resources import ProcessTreeSampler from microcosm.build.telemetry_emitter_service.spool import EventSpool from microcosm.build.telemetry_emitter_service.timestamps import utc_now @@ -81,15 +85,17 @@ def __init__( self, *, socket_path: Path, - registration: Mapping[str, Any], + registration: Mapping[str, Any] | None, spool: EventSpool, delivery: CollectorDelivery, sampler: ProcessTreeSampler, heartbeat_seconds: float, drain_seconds: float = DEFAULT_DRAIN_SECONDS, + graph_delivery: GraphPublicationDelivery | None = None, ) -> None: self.socket_path = socket_path - self.registration = dict(registration) + self.registration = dict(registration) if registration is not None else None + self.graph_delivery = graph_delivery self.spool = spool self.delivery = delivery self.sampler = sampler @@ -104,7 +110,8 @@ def __init__( def run(self) -> None: """Serve local messages until the client closes or exits.""" - self.spool.register(self.registration) + if self.registration is not None: + self.spool.register(self.registration) self.socket_path.parent.mkdir(parents=True, exist_ok=True) try: self.socket_path.unlink(missing_ok=True) @@ -113,11 +120,23 @@ def run(self) -> None: os.chmod(self.socket_path, 0o600) server.listen(SOCKET_LISTEN_BACKLOG) server.settimeout(SOCKET_ACCEPT_TIMEOUT_SECONDS) - worker = threading.Thread(target=self._worker, daemon=True) - worker.start() + workers = [] + if self.registration is not None: + workers.append(threading.Thread(target=self._worker, daemon=True)) + if self.graph_delivery is not None: + workers.append( + threading.Thread(target=self._graph_worker, daemon=True) + ) + for worker in workers: + worker.start() self._serve(server) - worker.join(timeout=self.drain_seconds + WORKER_INTERVAL_SECONDS) + deadline = ( + time.monotonic() + self.drain_seconds + WORKER_INTERVAL_SECONDS + ) + for worker in workers: + worker.join(timeout=max(0, deadline - time.monotonic())) finally: + self._stop.set() self.socket_path.unlink(missing_ok=True) try: self.socket_path.parent.rmdir() @@ -126,6 +145,15 @@ def run(self) -> None: def _serve(self, server: socket.socket) -> None: while not self._stop.is_set(): + if not self.sampler.parent_alive(): + if self.registration is not None: + self.spool.append( + self.registration, + _unexpected_exit_event(self._last_stage), + resources=self.sampler.sample(), + ) + self._stop.set() + break try: connection, _ = server.accept() except TimeoutError: @@ -162,7 +190,11 @@ def _read_message(connection: socket.socket) -> dict[str, Any]: def _handle(self, message: Mapping[str, Any]) -> None: action = message.get("action") - if action == ACTION_EVENT: + if action == GRAPH_PUBLICATION_ACTION and self.graph_delivery is not None: + self.graph_delivery.queue.enqueue( + Path(message["directory"]), message["inventory"] + ) + elif action == ACTION_EVENT and self.registration is not None: event = message.get("event") if not isinstance(event, Mapping): raise ValueError(EVENT_OBJECT_ERROR) @@ -200,16 +232,14 @@ def _worker(self) -> None: ) next_heartbeat = now + self.heartbeat_seconds self.delivery.flush_once() - if not self.sampler.parent_alive(): - self.spool.append( - self.registration, - _unexpected_exit_event(self._last_stage), - resources=self.sampler.sample(), - ) - self._stop.set() - break self._drain() + def _graph_worker(self) -> None: + # Uploads may take much longer than a heartbeat. They never run on the + # telemetry worker or the local socket thread. Pending jobs survive exit. + while not self._stop.wait(WORKER_INTERVAL_SECONDS): + self.graph_delivery.flush_once() + def _drain(self) -> None: deadline = time.monotonic() + self.drain_seconds while self.spool.has_deliverable() and time.monotonic() < deadline: diff --git a/packages/microcosm-build/src/microcosm/build/us_runtime/worker_identity.py b/packages/microcosm-build/src/microcosm/build/us_runtime/worker_identity.py index 618330734..2d5aff5b0 100644 --- a/packages/microcosm-build/src/microcosm/build/us_runtime/worker_identity.py +++ b/packages/microcosm-build/src/microcosm/build/us_runtime/worker_identity.py @@ -28,7 +28,7 @@ PRIMARY_QRF_WORKER_MODULE = "microcosm.build.us_runtime.puf_qrf_worker" PRIMARY_QRF_INTERPRETER_PLACEHOLDER = "{python_interpreter}" APPROVED_UV_LOCK_SHA256 = ( - "cbbfe1c0bd9b09e4fd54b8461fe5b2935d8604008a03b6a2928c458e25cdeb00" + "d73966a4020231e1b6d52bb5d69d852bbfc95827519dd4c98fc5c71ac8b20e19" ) LEGACY_CAMPAIGN_UV_LOCK_SHA256 = ( "27f47e385cfa35e2644a37410d1804b361ad9aee123577551c8421547bda65ee" diff --git a/packages/microcosm-build/tests/engine_free/shared/test_graph_publication.py b/packages/microcosm-build/tests/engine_free/shared/test_graph_publication.py new file mode 100644 index 000000000..1696a56af --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/shared/test_graph_publication.py @@ -0,0 +1,245 @@ +"""Durable graph jobs use fake HTTP, independently of the visualization app.""" + +import json + +import httpx +import pytest + +from microcosm.build.telemetry_emitter_service.graph_publication import ( + GraphPublicationDelivery, + GraphPublicationQueue, + inventory_digest, + publication_inventory, +) + + +def source(tmp_path): + directory = tmp_path / "source" + directory.mkdir() + (directory / "graph.orrery.json").write_text('{"fixture":true}') + return directory, publication_inventory(directory, {"graph.orrery.json": "graph"}) + + +def test_jobs_survive_restart_without_run_registration_or_event_retention(tmp_path): + directory, inventory = source(tmp_path) + queue = GraphPublicationQueue(tmp_path / "jobs.sqlite3") + queue.enqueue(directory, inventory) + restarted = GraphPublicationQueue(queue.path) + job = restarted.claim() + assert job["inventory"] == inventory + assert "run_id" not in inventory + assert job["publication_id"] == inventory["publication_id"] + assert restarted.claim() is None + + +def test_queue_rejects_changed_bytes_and_conflicting_inventory(tmp_path): + directory, inventory = source(tmp_path) + queue = GraphPublicationQueue(tmp_path / "jobs.sqlite3") + (directory / "graph.orrery.json").write_text("changed") + with pytest.raises(ValueError, match="digest"): + queue.enqueue(directory, inventory) + with pytest.raises(ValueError): + publication_inventory(directory, {"private.h5": "graph"}) + + +def test_delivery_retries_missing_credentials_and_network_failure(tmp_path): + directory, inventory = source(tmp_path) + queue = GraphPublicationQueue(tmp_path / "jobs.sqlite3") + queue.enqueue(directory, inventory) + delivery = GraphPublicationDelivery( + queue, credential=lambda: ("missing", None), origin="https://runs.example" + ) + delivery.flush_once() + assert queue.receipt(inventory["publication_id"])["status"] == "pending" + queue.retry(inventory["publication_id"]) + client = httpx.Client( + transport=httpx.MockTransport( + lambda request: (_ for _ in ()).throw( + httpx.ConnectError("secret", request=request) + ) + ) + ) + GraphPublicationDelivery( + queue, + credential=lambda: ("ok", "session"), + origin="https://runs.example", + client=client, + ).flush_once() + receipt = queue.receipt(inventory["publication_id"]) + assert receipt["status"] == "pending" + assert "secret" not in json.dumps(receipt) + + +@pytest.mark.parametrize( + "upload_url", + [ + "https://store.blob.vercel-storage.com/upload?signed=credential", + "https://vercel.com/api/blob/?pathname=object&signed=credential", + ], +) +def test_delivery_uses_exact_bytes_and_does_not_forward_session_to_signed_upload( + tmp_path, + upload_url, +): + directory, inventory = source(tmp_path) + queue = GraphPublicationQueue(tmp_path / "jobs.sqlite3") + queue.enqueue(directory, inventory) + requests = [] + + def respond(request): + requests.append(request) + if request.url.path == "/api/uploads/": + return httpx.Response( + 200, + json={ + "uploads": [ + { + "name": "graph.orrery.json", + "url": upload_url, + "already_uploaded": False, + } + ] + }, + ) + if request.method == "PUT": + assert "authorization" not in request.headers + assert request.content == (directory / "graph.orrery.json").read_bytes() + return httpx.Response(200, json={}) + return httpx.Response( + 200, + json={ + "publication_id": inventory["publication_id"], + "inventory_sha256": inventory_digest(inventory), + }, + ) + + delivery = GraphPublicationDelivery( + queue, + credential=lambda: ("ok", "session"), + origin="https://runs.example", + client=httpx.Client(transport=httpx.MockTransport(respond)), + ) + delivery.flush_once() + receipt = queue.receipt(inventory["publication_id"]) + assert receipt["status"] == "published" + assert len(requests) == 3 + assert "signed=" not in json.dumps(receipt) + assert "session" not in json.dumps(receipt) + + +def test_queue_resumes_after_worker_exits_and_keeps_pending_jobs(tmp_path): + directory, inventory = source(tmp_path) + queue = GraphPublicationQueue(tmp_path / "jobs.sqlite3") + queue.enqueue(directory, inventory) + queue.claim(lease_seconds=0) + assert GraphPublicationQueue(queue.path).claim() is not None + + +def test_late_failed_attempt_cannot_replace_a_published_receipt(tmp_path): + directory, inventory = source(tmp_path) + queue = GraphPublicationQueue(tmp_path / "jobs.sqlite3") + queue.enqueue(directory, inventory) + queue.claim(lease_seconds=0) + replacement = GraphPublicationQueue(queue.path) + replacement.claim() + id_ = inventory["publication_id"] + replacement.finish(id_, url=f"https://runs.example/runs/{id_}/") + published = queue.receipt(id_) + queue.finish(id_, error_code="publication_network_error") + assert queue.receipt(id_) == published + assert json.loads((directory / "publication.status.json").read_bytes()) == published + + +def test_expired_worker_cannot_release_a_replacement_workers_lease(tmp_path): + directory, inventory = source(tmp_path) + queue = GraphPublicationQueue(tmp_path / "jobs.sqlite3") + queue.enqueue(directory, inventory) + expired = queue.claim(lease_seconds=0) + replacement = GraphPublicationQueue(queue.path) + current = replacement.claim() + id_ = inventory["publication_id"] + initial = queue.receipt(id_) + queue.finish( + id_, + error_code="publication_network_error", + expected_lease=expired["lease_until"], + ) + assert queue.receipt(id_) == initial + assert queue.claim() is None + replacement.finish( + id_, + url=f"https://runs.example/runs/{id_}/", + expected_lease=current["lease_until"], + ) + assert queue.receipt(id_)["status"] == "published" + + +def test_telemetry_pruning_cannot_remove_graph_jobs(tmp_path): + from microcosm.build.telemetry_emitter_service.spool import EventSpool + + directory, inventory = source(tmp_path) + queue = GraphPublicationQueue(tmp_path / "jobs.sqlite3") + queue.enqueue(directory, inventory) + EventSpool(queue.path).prune() + assert ( + GraphPublicationQueue(queue.path).receipt(inventory["publication_id"])["status"] + == "pending" + ) + + +def test_rejected_graph_session_is_refreshed_without_discarding_the_job(tmp_path): + directory, inventory = source(tmp_path) + queue = GraphPublicationQueue(tmp_path / "jobs.sqlite3") + queue.enqueue(directory, inventory) + invalidations = [] + delivery = GraphPublicationDelivery( + queue, + credential=lambda: ("ok", "expired"), + invalidate_credential=lambda: invalidations.append(True), + origin="https://runs.example", + client=httpx.Client( + transport=httpx.MockTransport(lambda request: httpx.Response(401)) + ), + ) + assert not delivery.flush_once() + assert invalidations == [True] + assert queue.receipt(inventory["publication_id"])["status"] == "pending" + + +def test_retry_cli_reuses_identity_and_leaves_hf_snapshot_unchanged( + tmp_path, monkeypatch +): + from microcosm.build import graph_publication_cli as cli + + directory, inventory = source(tmp_path) + (directory / "orrery.upload.json").write_text(json.dumps(inventory)) + snapshot = directory / "orrery.publication.json" + snapshot.write_text('{"status":"pending"}') + calls = [] + + class Delivery: + def __init__(self, queue, **kwargs): + self.queue = queue + + def flush_once(self, *, publication_id): + calls.append(publication_id) + self.queue.finish( + publication_id, url="https://runs.example/runs/published/" + ) + + monkeypatch.setattr(cli, "GraphPublicationDelivery", Delivery) + assert ( + cli.main( + [ + "--spool", + str(tmp_path / "jobs.sqlite3"), + "--publication-id", + inventory["publication_id"], + "--directory", + str(directory), + ] + ) + == 0 + ) + assert calls == [inventory["publication_id"]] + assert snapshot.read_text() == '{"status":"pending"}' diff --git a/packages/microcosm-build/tests/engine_free/shared/test_publication_workers.py b/packages/microcosm-build/tests/engine_free/shared/test_publication_workers.py new file mode 100644 index 000000000..0289c8012 --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/shared/test_publication_workers.py @@ -0,0 +1,105 @@ +"""Independent publication delivery in the existing local service.""" + +import socket +import tempfile +import threading +import time +from pathlib import Path + +from microcosm.build.telemetry_emitter_service.auth import CollectorSession +from microcosm.build.telemetry_emitter_service.main import build_parser +from microcosm.build.telemetry_emitter_service.runtime import EmitterService +from microcosm.build.telemetry_emitter_service.spool import EventSpool + + +def test_slow_graph_upload_does_not_block_telemetry_or_socket( + tmp_path, real_local_telemetry +): + entered, release, telemetry_worked = (threading.Event() for _ in range(3)) + + class GraphDelivery: + def flush_once(self): + entered.set() + release.wait(5) + + class TelemetryDelivery: + def flush_once(self): + telemetry_worked.set() + return False + + class Sampler: + def parent_alive(self): + return True + + def sample(self): + return {} + + path = Path(tempfile.mkdtemp(prefix="mc-pub-test-", dir="/tmp")) / "e.sock" + service = EmitterService( + socket_path=path, + registration={ + "run_id": "run", + "producer_id": "producer", + "country_code": "GB", + "pipeline": "test", + }, + spool=EventSpool(tmp_path / "events.sqlite3"), + delivery=TelemetryDelivery(), + graph_delivery=GraphDelivery(), + sampler=Sampler(), + heartbeat_seconds=60, + drain_seconds=0, + ) + thread = threading.Thread(target=service.run) + thread.start() + try: + assert entered.wait(3) + assert telemetry_worked.wait(3) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as peer: + peer.settimeout(0.5) + peer.connect(str(path)) + peer.sendall(b'{"action":"ping"}\n') + assert peer.recv(16) == b"ok\n" + finally: + service._stop.set() + release.set() + thread.join(3) + assert not thread.is_alive() + + +def test_graph_only_service_requires_no_run_registration(tmp_path): + args = build_parser().parse_args( + [ + "--socket", + str(tmp_path / "e.sock"), + "--spool", + str(tmp_path / "jobs.sqlite3"), + "--parent-pid", + "1", + ] + ) + assert args.registration_json is None + + +def test_both_workers_share_one_collector_exchange(): + calls = [] + + def exchange(*args): + calls.append(args) + time.sleep(0.01) + return 200, {"access_token": "synthetic-session", "expires_in": 3600} + + session = CollectorSession( + "https://collector.example", post=exchange, credential=lambda: "hf-synthetic" + ) + results = [] + threads = [ + threading.Thread(target=lambda: results.append(session.credential())) + for _ in range(2) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(2) + assert results == [("ok", "synthetic-session")] * 2 + assert len(calls) == 1 diff --git a/packages/microcosm-build/tests/engine_free/shared/test_telemetry_emitter.py b/packages/microcosm-build/tests/engine_free/shared/test_telemetry_emitter.py index cee28f434..a2cb50116 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_telemetry_emitter.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_telemetry_emitter.py @@ -308,7 +308,10 @@ def test_current_pre_alembic_spool_is_adopted_without_losing_events(tmp_path) -> path = tmp_path / "events.sqlite3" registration = _registration() engine = create_engine(sqlite_database_url(path)) - SpoolModel.metadata.create_all(engine) + SpoolModel.metadata.create_all( + engine, + tables=[TelemetryRunRecord.__table__, TelemetryEventRecord.__table__], + ) with Session(engine) as session, session.begin(): session.add( TelemetryRunRecord( diff --git a/uv.lock b/uv.lock index 9ea420ac3..8da37c6d3 100644 --- a/uv.lock +++ b/uv.lock @@ -801,6 +801,7 @@ version = "0.1.0" source = { editable = "packages/microcosm-build" } dependencies = [ { name = "alembic" }, + { name = "httpx" }, { name = "huggingface-hub" }, { name = "jsonschema" }, { name = "microcosm-calibrate" }, @@ -844,6 +845,7 @@ requires-dist = [ { name = "alembic", specifier = ">=1.13.3,<2" }, { name = "h5py", marker = "extra == 'uk'", specifier = ">=3" }, { name = "h5py", marker = "extra == 'us'", specifier = ">=3" }, + { name = "httpx", specifier = ">=0.27,<1" }, { name = "huggingface-hub", specifier = ">=0.20" }, { name = "jsonschema", specifier = ">=4.23,<5" }, { name = "microcosm-calibrate", editable = "packages/microcosm-calibrate" }, From a16fa2b372396277791e4fb14329eb8efdc9c8ef Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 9 Oct 2026 00:03:10 +0400 Subject: [PATCH 2/6] Export and publish UK execution evidence before artifact staging --- .../build/uk_runtime/full_build_cli.py | 109 ++++++++- .../build/uk_runtime/orrery_publication.py | 216 ++++++++++++++++++ .../engine_free/uk/test_uk_full_build_cli.py | 2 +- .../uk/test_uk_orrery_publication.py | 147 ++++++++++++ .../uk/test_uk_rowwise_national_role.py | 4 +- test_support/microcosm_build/telemetry.py | 17 ++ 6 files changed, 488 insertions(+), 7 deletions(-) create mode 100644 packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_publication.py create mode 100644 packages/microcosm-build/tests/engine_free/uk/test_uk_orrery_publication.py diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/full_build_cli.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/full_build_cli.py index c39f63901..70725f23a 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/full_build_cli.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/full_build_cli.py @@ -619,6 +619,20 @@ def parse_args(argv: list[str] | None = None) -> argparse.Namespace: default_upload_interval_seconds=STAGING_UPLOAD_INTERVAL_SECONDS, ) add_staged_dataset_arguments(parser, repository=UK_STAGED_DATASET_REPOSITORY) + graph_publication = parser.add_mutually_exclusive_group() + graph_publication.add_argument( + "--publish-orrery", + dest="publish_orrery", + action="store_true", + help="Publish the saved graph, including when dataset staging is local-only.", + ) + graph_publication.add_argument( + "--no-publish-orrery", + dest="publish_orrery", + action="store_false", + help="Export graph files locally without submitting a publication job.", + ) + parser.set_defaults(publish_orrery=None) args = parser.parse_args(argv) validate_staging_arguments(parser, args) validate_staged_dataset_arguments(parser, args) @@ -1210,6 +1224,69 @@ def _release_stem(posture) -> tuple[str, str]: ) +def _run_with_graph_publication(execute, prepared, args, *, telemetry, attempt, record): + """Export before temporary files disappear, even if a kernel raises.""" + from .orrery_publication import ( + EVIDENCE_MANIFEST_NAME, + finalize_graph, + stage_graph_evidence, + ) + + try: + status = execute( + prepared, args, telemetry=telemetry, attempt=attempt, record=record + ) + except BaseException: + finalize_graph(args, record) + # A failed build may have no dataset. Preserve only its explicit graph + # inventory, not incomplete H5 files or arbitrary temporary JSON. + from microcosm.build.artifact_files import publish_staged_bundle + + output = Path(args.published_out) + names = [*record.get("orrery_files", []), EVIDENCE_MANIFEST_NAME] + staged = { + name: args.out / name for name in names if (args.out / name).is_file() + } + marker = args.out / "orrery.failure.json" + materialize_bytes( + canonical_json( + { + "schema_version": 1, + "kind": "uk_build_failed", + "releasable": False, + "outputs": { + name: file_artifact(path) for name, path in staged.items() + }, + } + ), + marker, + ) + staged["manifest"] = marker + publish_staged_bundle( + staged, + {role: output / path.name for role, path in staged.items()}, + completion_role="manifest", + ) + try: + stage_graph_evidence( + args, + output, + run_id=( + telemetry.run_id + if telemetry is not None + else Path(args.attempt_evidence).name + ), + ) + except Exception as error: + print( + f"Graph evidence staging failed ({type(error).__name__}); files remain local.", + file=sys.stderr, + ) + raise + finalize_graph(args, record) + return status + + def execute_full_build( prepared: PreparedUKFullBuild, args: argparse.Namespace, @@ -1245,8 +1322,13 @@ def execute_full_build( staged_args.out = Path(temporary) staged_args.graph_store = graph_store staged_args.published_out = output - status = _execute_full_build( - prepared, staged_args, telemetry=telemetry, attempt=attempt, record=record + status = _run_with_graph_publication( + _execute_full_build, + prepared, + staged_args, + telemetry=telemetry, + attempt=attempt, + record=record, ) if not (staged_args.out / "build.json").exists(): materialize_bytes( @@ -1349,6 +1431,13 @@ def _close_attempt( output / f"{stem}.h5", repository_hint=REPOSITORY ) else: + from .orrery_publication import stage_graph_evidence + + stage_graph_evidence( + args, + output, + run_id=state.build_id if telemetry is None else telemetry.run_id, + ) _close_staging(args, telemetry, classified, record) spool_path = record_candidate_attempt( state=state, @@ -2180,8 +2269,13 @@ def execute_national_build( staged_args.out = Path(temporary) staged_args.graph_store = graph_store staged_args.published_out = output - status = _execute_national_build( - prepared, staged_args, telemetry=telemetry, attempt=attempt, record=record + status = _run_with_graph_publication( + _execute_national_build, + prepared, + staged_args, + telemetry=telemetry, + attempt=attempt, + record=record, ) if not (staged_args.out / "build.json").exists(): materialize_bytes( @@ -2723,6 +2817,13 @@ def _close_national_attempt( }, ) else: + from .orrery_publication import stage_graph_evidence + + stage_graph_evidence( + args, + output, + run_id=state.build_id if telemetry is None else telemetry.run_id, + ) _close_staging(args, telemetry, classified, record) spool_path = record_candidate_attempt( state=state, diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_publication.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_publication.py new file mode 100644 index 000000000..c2bd2047a --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_publication.py @@ -0,0 +1,216 @@ +"""Export saved UK execution evidence before staging the build's artifact bundle. + +Only the graph, declarations, digest-bound native receipts, and reviewed +aggregate summaries are eligible. Population files and content-store payloads +are never uploaded to the public visualization service. +""" + +from __future__ import annotations + +import json +import re +from pathlib import Path + +from microcosm.build.artifact_files import file_artifact, materialize_bytes +from microcosm.build.telemetry_emitter_service.graph_publication import ( + publication_inventory, +) +from microcosm.graph import ContentStore +from microcosm.graph.canonical import canonical_json +from microcosm.graph.evidence import collect_execution_evidence, load_run_evidence +from microcosm.graph.orrery import orrery_json_from_schema + +RECEIPT_NAME = "orrery.publication.json" +INVENTORY_NAME = "orrery.upload.json" +EVIDENCE_MANIFEST_NAME = "orrery.evidence-manifest.json" + + +def publication_enabled(args) -> bool: + explicit = getattr(args, "publish_orrery", None) + if explicit is not None: + return explicit + return not ( + getattr(args, "no_staging", False) + or getattr(args, "staging_local_only", False) + or getattr(args, "no_staged_dataset", False) + ) + + +def _urls(value): + if isinstance(value, dict): + for key, item in value.items(): + if key == "url" and isinstance(item, str): + yield item + else: + yield from _urls(item) + elif isinstance(value, list): + for item in value: + yield from _urls(item) + + +def _output(path, published_root): + info = file_artifact(path) + return { + "path": str(published_root / path.name), + "sha256": info["sha256"], + "bytes": info["size_bytes"], + } + + +def finalize_graph(args, record: dict) -> dict: + """Freeze bytes and a receipt snapshot; retry never rewrites that snapshot. + + This runs inside the temporary build directory, including on exceptions. + Graph publication failure does not replace a dataset's scientific verdict. + """ + from . import rowwise_staging + + output = Path(args.out) + published_root = Path(getattr(args, "published_out", output)) + schema_path = output / "graph.schema.json" + receipt = { + "version": 1, + "status": "no_evidence", + "publication_id": None, + "error_code": None, + "url": None, + } + names = [] + try: + if schema_path.is_file(): + schema = json.loads(schema_path.read_bytes()) + index_path = output / "execution.evidence.json" + execution = None + roles = {schema_path.name: "schema", "graph.orrery.json": "graph"} + if index_path.is_file(): + store = ContentStore(args.graph_store, create=False) + runs = load_run_evidence(index_path, store=store, reference_base=output) + execution = collect_execution_evidence(schema, runs=runs, store=store) + roles[index_path.name] = "index" + index = json.loads(index_path.read_bytes()) + for run in index["runs"]: + for field, role in ( + ("graph", "declaration"), + ("manifest", "manifest"), + ("binding", "binding"), + ("summaries", "summaries"), + ): + roles[run[field]["path"]] = role + payload = orrery_json_from_schema(schema, execution=execution).encode() + materialize_bytes(payload, output / "graph.orrery.json") + for url in _urls(json.loads(payload)): + name = url.split("#", 1)[0] + if re.fullmatch(r"upstream-[a-f0-9]{64}\.json", name): + if file_artifact(output / name)["sha256"] != name[9:-5]: + raise ValueError("Upstream graph evidence digest mismatch.") + roles[name] = "upstream" + # Keep exact sources outside the build's temporary directory and + # outside HF's immutable receipt snapshot, including when offline. + names = list(roles) + preserved = Path(args.attempt_evidence) / "orrery-publication" + for name in roles: + materialize_bytes((output / name).read_bytes(), preserved / name) + inventory = publication_inventory(preserved, roles) + materialize_bytes(canonical_json(inventory), preserved / INVENTORY_NAME) + materialize_bytes(canonical_json(inventory), output / INVENTORY_NAME) + names = [*roles, INVENTORY_NAME] + receipt = { + "version": 1, + "publication_id": inventory["publication_id"], + "status": "skipped", + "error_code": "local_only", + "url": None, + "preserved_directory": str(preserved), + "recorded_execution": execution is not None, + } + if publication_enabled(args): + emitter = rowwise_staging._ACTIVE_EMITTER + if emitter is None: + # Durable enqueue is still possible without a running + # emitter; the retry command requires no telemetry run. + from microcosm.build.telemetry_emitter import _cache_dir + from microcosm.build.telemetry_emitter_constants import ( + TELEMETRY_SPOOL_FILENAME, + ) + from microcosm.build.telemetry_emitter_service.graph_publication import ( + GraphPublicationQueue, + ) + + queue = GraphPublicationQueue( + _cache_dir() / TELEMETRY_SPOOL_FILENAME + ) + queue.enqueue(preserved, inventory) + receipt.update(queue.receipt(inventory["publication_id"])) + else: + receipt.update(emitter.publish_graph(preserved, inventory)) + except Exception: + # No credential, exception text, paths to raw input data, or signed URL + # is copied from an exception into public evidence or telemetry. + receipt.update(status="failed", error_code="graph_export_or_enqueue_failed") + materialize_bytes(canonical_json(receipt), output / RECEIPT_NAME) + names.append(RECEIPT_NAME) + record["orrery_files"] = names + record["orrery_publication"] = receipt + manifest = record.get("manifest") + if manifest is not None: + from .rowwise_cli import MANIFEST_FILENAME, json_text + + outputs = manifest.setdefault("outputs", {}) + for name in names: + outputs[f"orrery/{name}"] = _output(output / name, published_root) + manifest["orrery_publication"] = receipt + materialize_bytes(json_text(manifest).encode(), output / MANIFEST_FILENAME) + marker = output / "build.json" + if marker.is_file(): + completion = json.loads(marker.read_bytes()) + completion["rowwise_candidate_manifest"] = file_artifact( + output / MANIFEST_FILENAME + ) + materialize_bytes(canonical_json(completion), marker) + elif names: + # Failed/refused builds have no dataset manifest. Give their saved + # evidence its own inventory rather than pretending an H5 was built. + evidence = { + "schema_version": 1, + "build_kind": "uk_graph_evidence", + "releasable": False, + "orrery_publication": receipt, + "outputs": {name: _output(output / name, published_root) for name in names}, + } + materialize_bytes(canonical_json(evidence), output / EVIDENCE_MANIFEST_NAME) + return receipt + + +def stage_graph_evidence(args, output: Path, *, run_id: str) -> dict | None: + """Stage failure evidence through the existing generic HF bundle adapter.""" + from microcosm.build.staging_dataset import ( + StagedDatasetBundle, + disabled_staged_dataset, + local_only_staged_dataset, + stage_bundle, + write_sidecars, + ) + from microcosm.build.staging_storage import HuggingFaceDatasetStorage + + from . import rowwise_staging + from .staging import UK_STAGED_DATASET_PREFIX + + if not (output / EVIDENCE_MANIFEST_NAME).is_file(): + return None + mode = rowwise_staging.staged_dataset_mode(args) + if mode == "disabled": + return disabled_staged_dataset("--no-staging or --no-staged-dataset") + repository = None if mode == "local_only" else args.staged_dataset_repo_id + bundle = StagedDatasetBundle.from_manifest( + output, run_id=run_id, manifest_name=EVIDENCE_MANIFEST_NAME + ) + write_sidecars( + bundle, repository=repository, prefix=UK_STAGED_DATASET_PREFIX, telemetry=None + ) + if mode == "local_only": + return local_only_staged_dataset(bundle, prefix=UK_STAGED_DATASET_PREFIX) + return stage_bundle( + bundle, + storage=HuggingFaceDatasetStorage(repository, api=rowwise_staging._hub_api()), + prefix=UK_STAGED_DATASET_PREFIX, + ) diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_full_build_cli.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_full_build_cli.py index 6aea249db..9fc790b54 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_full_build_cli.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_full_build_cli.py @@ -372,7 +372,7 @@ def test_cli_cold_and_required_replay_recreate_dataset_and_sidecars( for node in snapshot["nodes"] ) assert [phase["phase"] for phase in overlay["phases"]][-2:] == ["export", "final"] - assert not list(out.glob("*.orrery.json")) + assert json.loads((out / "graph.orrery.json").read_bytes()) == snapshot assert (out / f"{STEM}.targets.csv").read_text().startswith("name,target_name") cold = json.loads((out / cli.MANIFEST_FILENAME).read_text())["execution"] assert cold["nodes_total"] > 0 and cold["nodes_reused"] == 0 diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_orrery_publication.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_orrery_publication.py new file mode 100644 index 000000000..221afedf3 --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_orrery_publication.py @@ -0,0 +1,147 @@ +"""UK exports and staging order, independent of the runs application's code.""" + +import json +from argparse import Namespace + +import pytest + +from microcosm.build.uk_runtime import ( + full_build_cli, + orrery_publication, + rowwise_staging, +) +from microcosm.build.uk_runtime.orrery_contract import save_uk_graph_schema +from microcosm.graph import Graph, Node, Owned, SourceRef, StructuralDelta + + +def fixture(tmp_path, **options): + out = tmp_path / "temporary" + out.mkdir() + graph = Graph( + "uk", + (SourceRef("uk_spine", "raw-bytes-v1"),), + ( + Node( + "uk.full.spine_checkpoint", + "fixture@1", + sources=("uk_spine",), + structural=StructuralDelta.CREATE, + outputs=(Owned("person", "age", "int64"),), + ), + ), + ) + save_uk_graph_schema(graph, out, scope="national") + args = Namespace( + out=out, + published_out=tmp_path / "out", + attempt_evidence=tmp_path / "attempt", + graph_store=tmp_path / "store", + no_staging=False, + no_staged_dataset=False, + staging_local_only=True, + publish_orrery=None, + **options, + ) + return args + + +def test_local_export_preserves_exact_graph_without_a_run_id_or_network( + tmp_path, monkeypatch +): + args = fixture(tmp_path) + monkeypatch.setattr(rowwise_staging, "_ACTIVE_EMITTER", None) + (args.out / "private.h5").write_bytes(b"licensed population") + (args.out / "raw.json").write_text('{"person_id":123}') + record = {} + receipt = orrery_publication.finalize_graph(args, record) + assert receipt["status"] == "skipped" + assert receipt["recorded_execution"] is False + preserved = args.attempt_evidence / "orrery-publication" + inventory = json.loads((preserved / "orrery.upload.json").read_bytes()) + assert "run_id" not in inventory + assert {f["name"] for f in inventory["files"]} == { + "graph.orrery.json", + "graph.schema.json", + } + assert (preserved / "graph.orrery.json").read_bytes() == ( + args.out / "graph.orrery.json" + ).read_bytes() + assert not (preserved / "private.h5").exists() + assert not (preserved / "raw.json").exists() + + +def test_publication_precedes_manifest_registration_and_receipt_is_a_snapshot( + tmp_path, monkeypatch +): + args = fixture(tmp_path) + args.publish_orrery = True + events = [] + + class Emitter: + def publish_graph(self, directory, inventory): + assert not (args.out / "orrery.publication.json").exists() + events.append("publication_attempt") + return { + "publication_id": inventory["publication_id"], + "status": "pending", + "error_code": "credential_missing", + } + + monkeypatch.setattr(rowwise_staging, "_ACTIVE_EMITTER", Emitter()) + manifest = {"outputs": {}} + record = {"manifest": manifest} + receipt = orrery_publication.finalize_graph(args, record) + events.append("registered") + assert events == ["publication_attempt", "registered"] + assert receipt["status"] == "pending" + names = {entry["path"].split("/")[-1] for entry in manifest["outputs"].values()} + assert { + "graph.orrery.json", + "graph.schema.json", + "orrery.publication.json", + "orrery.upload.json", + } <= names + snapshot = (args.out / "orrery.publication.json").read_bytes() + ( + args.attempt_evidence / "orrery-publication" / "publication.status.json" + ).write_text('{"status":"published"}') + assert (args.out / "orrery.publication.json").read_bytes() == snapshot + + +def test_failure_path_preserves_graph_before_temporary_cleanup(tmp_path, monkeypatch): + args = fixture(tmp_path) + calls = [] + monkeypatch.setattr( + orrery_publication, "stage_graph_evidence", lambda *a, **k: calls.append("hf") + ) + + def fail(*a, **k): + raise RuntimeError("kernel failed") + + with pytest.raises(RuntimeError, match="kernel failed"): + full_build_cli._run_with_graph_publication( + fail, None, args, telemetry=None, attempt=None, record={} + ) + assert calls == ["hf"] + assert (args.published_out / "graph.orrery.json").is_file() + assert (args.published_out / "orrery.evidence-manifest.json").is_file() + assert (args.published_out / "orrery.failure.json").is_file() + + +@pytest.mark.parametrize( + "explicit,local,disabled,expected", + [ + (None, False, False, True), + (None, True, False, False), + (None, False, True, False), + (True, True, True, True), + (False, False, False, False), + ], +) +def test_publication_has_explicit_local_only_and_opt_out_controls( + explicit, local, disabled, expected +): + args = Namespace( + publish_orrery=explicit, staging_local_only=local, no_staging=disabled + ) + assert orrery_publication.publication_enabled(args) is expected diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_rowwise_national_role.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_rowwise_national_role.py index 654a25b69..af83888fd 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_rowwise_national_role.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_rowwise_national_role.py @@ -305,7 +305,7 @@ def test_uk_national_role_builds_through_the_graph_and_stages_locally( json.loads(node["id"]) == ["operation", "uk.full.national.readback"] for node in document["nodes"] ) - assert not list(out.glob("*.orrery.json")) + assert (out / "graph.orrery.json").is_file() # The build record the release-cut certifier reads: the seam's schema. record = json.loads((out / "build_record.json").read_text()) @@ -417,7 +417,7 @@ def test_uk_national_role_builds_through_the_graph_and_stages_locally( assert manifest["build_kind"] == "uk_national_calibrated_candidate" assert manifest["release_role"] == "national" assert manifest["release_id"] == UK_NATIONAL_RELEASE_ID - assert set(manifest["outputs"]) == { + assert {name for name in manifest["outputs"] if not name.startswith("orrery/")} == { "dataset", "calibration_diagnostics", "build_record", diff --git a/test_support/microcosm_build/telemetry.py b/test_support/microcosm_build/telemetry.py index 65e8d7797..7ff55a247 100644 --- a/test_support/microcosm_build/telemetry.py +++ b/test_support/microcosm_build/telemetry.py @@ -24,6 +24,23 @@ def __init__(self, run: TelemetryRun) -> None: def _send(self, payload: Mapping[str, Any]) -> None: self.messages.append(payload) + def publish_graph( + self, directory: Path, inventory: dict, *, wait_seconds: float = 30.0 + ) -> dict: + self.messages.append( + { + "action": "graph_publication", + "directory": str(directory), + "inventory": inventory, + } + ) + return { + "version": 1, + "publication_id": inventory["publication_id"], + "status": "pending", + "error_code": None, + } + @property def events(self) -> list[Mapping[str, Any]]: return [message["event"] for message in self.messages if "event" in message] From 38f4f80da25e69b81764403290c56d6306d3716e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 9 Oct 2026 00:03:10 +0400 Subject: [PATCH 3/6] Document publication configuration and provider migration boundary --- README.md | 5 ++++ changelog.d/1150.added.md | 1 + docs/agent-guide.md | 6 ++++ docs/orrery-publication.md | 61 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 73 insertions(+) create mode 100644 changelog.d/1150.added.md create mode 100644 docs/orrery-publication.md diff --git a/README.md b/README.md index f25d358cd..820921a1f 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,11 @@ until you restart the kernel, load it by path with ## Build progress and staging run files +UK build completion automatically exports Orrery execution graphs, attempts +immutable publication before Hugging Face artifact staging, and preserves +durable retry jobs and receipt snapshots. See +[graph publication](docs/orrery-publication.md) for configuration and retry. + Supported US and UK build commands always start the local telemetry emitter service. It reports live progress to the hosted collector when the operator's existing Hugging Face login is accepted; otherwise it retains the events diff --git a/changelog.d/1150.added.md b/changelog.d/1150.added.md new file mode 100644 index 000000000..4cec6ee6b --- /dev/null +++ b/changelog.d/1150.added.md @@ -0,0 +1 @@ +Export UK build execution graphs automatically, queue immutable authenticated publication with durable retries, and include graph evidence and receipt snapshots in staged artifacts. diff --git a/docs/agent-guide.md b/docs/agent-guide.md index 11c9328ec..d121fe37f 100644 --- a/docs/agent-guide.md +++ b/docs/agent-guide.md @@ -14,6 +14,12 @@ the PEP 420 namespace `microcosm.`: `frame`, `fit`, `calibrate`, `build`, ## Commands +UK build completion exports and stages Orrery graph evidence. Publication uses +an independent worker inside the existing local service and a durable queue; +see [configuration and retry](orrery-publication.md). No new producer secret or +environment variable is required. Keep graph export/publication tests independent +of the hosted runs app and collector; use synthetic files and fake HTTP. + ```bash uv sync --all-packages # set up the whole workspace uv sync --all-packages --locked --extra us # US engine environment diff --git a/docs/orrery-publication.md b/docs/orrery-publication.md new file mode 100644 index 000000000..246699191 --- /dev/null +++ b/docs/orrery-publication.md @@ -0,0 +1,61 @@ +# UK graph publication + +UK full, dense, size and national builds export `graph.orrery.json` from the +saved graph schema and recorded execution evidence. They preserve the explicit +public-safe file inventory in the attempt directory, attempt publication before +Hugging Face artifact staging, and register graph/evidence files and an +`orrery.publication.json` receipt snapshot in the artifact manifest. + +The existing local service has two independent workers: live telemetry and +completed graph publication. Slow uploads cannot delay socket acknowledgement +or telemetry delivery. This change does not introduce another generic runtime, +rename the telemetry package, or add compatibility wrappers. + +Publication IDs do not depend on telemetry run IDs. Jobs use SQLAlchemy and +Alembic in the existing local SQLite database. Pending jobs survive missing or +rejected credentials, process restarts and telemetry event pruning. Retry +leases prevent a stale worker from replacing another worker's result. +The exact source bytes must remain available until delivery completes. + +The authenticated API creates upload permissions, receives exact files using +scoped signed URLs, then finalizes an immutable publication. Signed uploads +never receive the collector bearer credential. The API rejects different bytes +under an existing publication ID. No population H5 or content-store arrays are +included; only allowlisted graph declarations, execution receipts and reviewed +aggregate summaries can be published publicly. + +## Configuration + +| Setting | Source and consumer | Provisioning | +| --- | --- | --- | +| `HF_TOKEN` / `HUGGINGFACE_TOKEN` | Existing operator credential, exchanged by the service for a short-lived collector session | Existing environment or Hugging Face login; never committed | +| Collector origin | Tracked `PRODUCTION_COLLECTOR_URL` in the existing service | No producer variable or destination override | +| Runs origin | Tracked `DEFAULT_GRAPH_PUBLICATION_ORIGIN` | No new producer variable | +| `XDG_CACHE_HOME` | Optional existing cache-location setting | Defaults to the user's cache directory | + +The separately deployed runs application retains its Vercel Blob configuration. +Authenticated production publication also requires +PolicyEngine/calibration-diagnostics#208 deployed, and the runs application's +publication API. Producer tests do not contact either hosted service. + +No new Microcosm runtime variable or secret is required. Tests isolate ambient +credentials and caches and use fake HTTP or explicit loopback services. + +`--staging-local-only`, `--no-staging` and `--no-staged-dataset` export locally +without remote submission. `--publish-orrery` explicitly enables submission; +`--no-publish-orrery` disables it independently. A publication failure does not +replace the dataset's validation result or prevent HF artifact staging. +Failed builds preserve recorded graph evidence before temporary-file cleanup. + +## Retry + +```bash +microcosm-publish-graph --publication-id PUBLICATION_ID \ + --directory /path/to/attempt/orrery-publication +``` + +The preserved directory contains `orrery.upload.json` and the exact files. +Retries update a separate `publication.status.json`, never the original +HF receipt snapshot. Pending jobs also resume when a subsequent local service +starts. Package extraction into `microcosm-local-provider` is a separate, +deferred consumer change; this implementation does not require those packages. From fc3080540ef67442aa65829f65141ed39a16fae3 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 9 Oct 2026 01:47:35 +0400 Subject: [PATCH 4/6] Replace generic publication returns with explicit data contracts --- changelog.d/1150-publication-types.changed.md | 1 + .../src/microcosm/build/artifact_files.py | 15 +++- .../build/graph_publication_types.py | 57 ++++++++++++++ .../src/microcosm/build/py.typed | 1 + .../src/microcosm/build/staging_dataset.py | 45 ++++++++--- .../src/microcosm/build/telemetry_emitter.py | 19 ++++- .../graph_publication.py | 74 ++++++++++++------- .../build/uk_runtime/orrery_publication.py | 44 ++++++++--- .../shared/test_graph_publication.py | 62 +++++++++++++++- 9 files changed, 264 insertions(+), 54 deletions(-) create mode 100644 changelog.d/1150-publication-types.changed.md create mode 100644 packages/microcosm-build/src/microcosm/build/graph_publication_types.py create mode 100644 packages/microcosm-build/src/microcosm/build/py.typed diff --git a/changelog.d/1150-publication-types.changed.md b/changelog.d/1150-publication-types.changed.md new file mode 100644 index 000000000..094729823 --- /dev/null +++ b/changelog.d/1150-publication-types.changed.md @@ -0,0 +1 @@ +Define explicit types for publication inventories, results, and staged files, and return a frozen record when reserving a graph upload. diff --git a/packages/microcosm-build/src/microcosm/build/artifact_files.py b/packages/microcosm-build/src/microcosm/build/artifact_files.py index ff14fd7af..263b30e77 100644 --- a/packages/microcosm-build/src/microcosm/build/artifact_files.py +++ b/packages/microcosm-build/src/microcosm/build/artifact_files.py @@ -6,11 +6,20 @@ import tempfile from collections.abc import Mapping from pathlib import Path +from typing import TypedDict from .trace import sha256_file -def file_artifact(path: str | Path) -> dict[str, object]: +class FileArtifact(TypedDict): + """The measured identity of a regular file.""" + + filename: str + sha256: str + size_bytes: int + + +def file_artifact(path: str | Path) -> FileArtifact: """Bind one regular file without loading its entire payload into memory.""" source = Path(path) if not source.is_file(): @@ -27,7 +36,7 @@ def file_artifact(path: str | Path) -> dict[str, object]: return {"filename": source.name, "sha256": digest, "size_bytes": after.st_size} -def materialize_bytes(payload: bytes, path: str | Path) -> dict[str, object]: +def materialize_bytes(payload: bytes, path: str | Path) -> FileArtifact: """Write deterministic bytes atomically, including after a graph cache hit.""" if not isinstance(payload, bytes): raise TypeError("Artifact payload must be immutable bytes.") @@ -76,7 +85,7 @@ def publish_staged_bundle( *, completion_role: str = "manifest", expected: Mapping[str, Mapping[str, object]] | None = None, -) -> dict[str, dict[str, object]]: +) -> dict[str, FileArtifact]: """Publish a validated bundle with rollback and the completion marker last. The marker is absent while files change. Handled failures, including diff --git a/packages/microcosm-build/src/microcosm/build/graph_publication_types.py b/packages/microcosm-build/src/microcosm/build/graph_publication_types.py new file mode 100644 index 000000000..94f4452f3 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/graph_publication_types.py @@ -0,0 +1,57 @@ +"""Explicit data contracts for immutable graph publication.""" + +from dataclasses import dataclass +from pathlib import Path +from typing import Literal, NotRequired, TypedDict + +type PublicationFileRole = Literal[ + "graph", + "schema", + "index", + "declaration", + "manifest", + "binding", + "summaries", + "upstream", +] + + +class PublicationFile(TypedDict): + """One public file, identified by its exact bytes and allowed role.""" + + name: str + role: PublicationFileRole + bytes: int + sha256: str + + +class PublicationInventory(TypedDict): + """The versioned file inventory submitted to the publication API.""" + + version: Literal[1] + publication_id: str + files: list[PublicationFile] + + +class PublicationReceipt(TypedDict): + """Local publication result, including retry and build-snapshot fields.""" + + version: Literal[1] + publication_id: str | None + status: Literal["pending", "published", "skipped", "failed", "no_evidence"] + error_code: str | None + url: NotRequired[str | None] + attempts: NotRequired[int] + inventory_sha256: NotRequired[str] + preserved_directory: NotRequired[str] + recorded_execution: NotRequired[bool] + + +@dataclass(frozen=True) +class ClaimedGraphPublication: + """A job reserved until lease_until; the timestamp also identifies its lease.""" + + publication_id: str + directory: Path + inventory: PublicationInventory + lease_until: float diff --git a/packages/microcosm-build/src/microcosm/build/py.typed b/packages/microcosm-build/src/microcosm/build/py.typed new file mode 100644 index 000000000..8b1378917 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/py.typed @@ -0,0 +1 @@ + diff --git a/packages/microcosm-build/src/microcosm/build/staging_dataset.py b/packages/microcosm-build/src/microcosm/build/staging_dataset.py index 37e7ae97a..6a8cd23da 100644 --- a/packages/microcosm-build/src/microcosm/build/staging_dataset.py +++ b/packages/microcosm-build/src/microcosm/build/staging_dataset.py @@ -25,7 +25,7 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path, PurePosixPath -from typing import Any +from typing import Any, Literal, TypedDict, cast from microcosm.build.staging_storage import HuggingFaceDatasetStorage @@ -36,6 +36,7 @@ "STAGED_MANIFEST_FILENAME", "STAGED_MANIFEST_SCHEMA", "StagedDatasetBundle", + "StagedDatasetDelivery", "StagedDatasetError", "StagedFile", "delivery_covers_bundle", @@ -90,6 +91,30 @@ class StagedDatasetError(ValueError): """A staged-dataset bundle or delivery record is malformed.""" +class StagedFileDigest(TypedDict): + """The exact bytes registered for one staged file.""" + + sha256: str + bytes: int + + +class StagedDatasetDelivery(TypedDict): + """Validated result of staging a dataset or graph-evidence bundle.""" + + contract_version: Literal[1] + mode: Literal["local_and_remote", "local_only", "disabled"] + repository: str | None + prefix: str | None + run_id: str | None + revision: str | None + status: Literal["uploaded", "already_staged", "failed", "skipped"] + error_code: ( + Literal["UPLOAD_FAILED", "REMOTE_DIFFERS", "REPOSITORY_UNAVAILABLE"] | None + ) + opt_out_reason: str | None + files: dict[str, StagedFileDigest] + + def sha256_file(path: Path | str) -> str: digest = hashlib.sha256() with Path(path).open("rb") as stream: @@ -353,7 +378,9 @@ def parse_sha256sums(text: str) -> list[tuple[str, str]]: return entries -def validate_staged_dataset_delivery(payload: Mapping[str, Any]) -> dict[str, Any]: +def validate_staged_dataset_delivery( + payload: Mapping[str, Any], +) -> StagedDatasetDelivery: """Validate and normalize a staged-dataset delivery record.""" if not isinstance(payload, Mapping): @@ -423,7 +450,7 @@ def validate_staged_dataset_delivery(payload: Mapping[str, Any]) -> dict[str, An or files ): raise StagedDatasetError("a disabled staged dataset records nothing else.") - return record + return cast(StagedDatasetDelivery, record) if reason is not None: raise StagedDatasetError( "only a disabled staged dataset carries an opt-out reason." @@ -439,7 +466,7 @@ def validate_staged_dataset_delivery(payload: Mapping[str, Any]) -> dict[str, An ) if error_code is not None: raise StagedDatasetError("a local-only staged dataset has no error code.") - return record + return cast(StagedDatasetDelivery, record) if repository is None: raise StagedDatasetError("a remote staged dataset names its repository.") if status == "skipped": @@ -451,12 +478,12 @@ def validate_staged_dataset_delivery(payload: Mapping[str, Any]) -> dict[str, An raise StagedDatasetError( "a failed staged dataset carries an error code and no revision." ) - return record + return cast(StagedDatasetDelivery, record) if error_code is not None: raise StagedDatasetError("a delivered staged dataset carries no error code.") if status == "uploaded" and revision is None: raise StagedDatasetError("an uploaded staged dataset records its revision.") - return record + return cast(StagedDatasetDelivery, record) def delivery_covers_bundle( @@ -483,7 +510,7 @@ def delivery_covers_bundle( ) -def disabled_staged_dataset(reason: str) -> dict[str, Any]: +def disabled_staged_dataset(reason: str) -> StagedDatasetDelivery: """Explicit evidence for a deliberate staged-dataset opt-out.""" if not isinstance(reason, str) or not reason.strip(): @@ -506,7 +533,7 @@ def disabled_staged_dataset(reason: str) -> dict[str, Any]: def local_only_staged_dataset( bundle: StagedDatasetBundle, *, prefix: str = DEFAULT_STAGED_DATASET_PREFIX -) -> dict[str, Any]: +) -> StagedDatasetDelivery: """Evidence for a bundle verified and inventoried on disk but not uploaded.""" return validate_staged_dataset_delivery( @@ -535,7 +562,7 @@ def stage_bundle( storage: HuggingFaceDatasetStorage, prefix: str = DEFAULT_STAGED_DATASET_PREFIX, message: str | None = None, -) -> dict[str, Any]: +) -> StagedDatasetDelivery: """Upload the bundle and its sidecars in one commit; record, never raise. Idempotent on the outputs' digests: a remote ``staged_manifest.json`` whose diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter.py index 160f33b09..dd186f8d0 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter.py @@ -21,6 +21,10 @@ from pathlib import Path from typing import Any +from microcosm.build.graph_publication_types import ( + PublicationInventory, + PublicationReceipt, +) from microcosm.build.telemetry_emitter_constants import ( DEFAULT_HEARTBEAT_SECONDS, DEFAULT_SEND_TIMEOUT_SECONDS, @@ -460,8 +464,12 @@ def close(self) -> None: self._closed = True def publish_graph( - self, directory: Path, inventory: dict, *, wait_seconds: float = 30.0 - ) -> dict: + self, + directory: Path, + inventory: PublicationInventory, + *, + wait_seconds: float = 30.0, + ) -> PublicationReceipt: """Preserve a job first, then wait a bounded time for its first result. Unlike telemetry, missing credentials never discard this job. Enqueue @@ -479,10 +487,15 @@ def publish_graph( deadline = time.monotonic() + max(0, wait_seconds) while self.available and time.monotonic() < deadline: receipt = queue.receipt(inventory["publication_id"]) + assert receipt is not None, ( + "A queued graph publication must have a receipt." + ) if receipt.get("attempts", 0) or receipt["status"] == "published": return receipt time.sleep(0.1) - return queue.receipt(inventory["publication_id"]) + receipt = queue.receipt(inventory["publication_id"]) + assert receipt is not None, "A queued graph publication must have a receipt." + return receipt def _close_transition_stage( self, diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/graph_publication.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/graph_publication.py index 17072ed49..293dc4d1b 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/graph_publication.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/graph_publication.py @@ -13,11 +13,19 @@ import uuid from collections.abc import Callable, Mapping from pathlib import Path +from typing import cast from urllib.parse import urlsplit import httpx from sqlalchemy import select, update +from microcosm.build.graph_publication_types import ( + ClaimedGraphPublication, + PublicationFile, + PublicationFileRole, + PublicationInventory, + PublicationReceipt, +) from microcosm.build.telemetry_emitter_service.database import ( create_spool_engine, create_spool_session_factory, @@ -48,26 +56,34 @@ def _digest(payload: bytes) -> str: return hashlib.sha256(payload).hexdigest() -def inventory_digest(inventory: Mapping) -> str: +def inventory_digest(inventory: PublicationInventory) -> str: """Digest the versioned field order shared with the publication API.""" - inventory = { + ordered: PublicationInventory = { "version": inventory["version"], "publication_id": inventory["publication_id"], "files": [ - {key: item[key] for key in ("name", "role", "bytes", "sha256")} + { + "name": item["name"], + "role": item["role"], + "bytes": item["bytes"], + "sha256": item["sha256"], + } for item in sorted(inventory["files"], key=lambda file: file["name"]) ], } return _digest( - json.dumps(inventory, separators=(",", ":"), ensure_ascii=True).encode() + json.dumps(ordered, separators=(",", ":"), ensure_ascii=True).encode() ) def publication_inventory( - directory: Path, roles: Mapping[str, str], *, publication_id: str | None = None -) -> dict: + directory: Path, + roles: Mapping[str, PublicationFileRole], + *, + publication_id: str | None = None, +) -> PublicationInventory: """Hash an explicit allowlist of files; publication identity is not a run ID.""" - files = [] + files: list[PublicationFile] = [] for name, role in sorted(roles.items()): if role not in GRAPH_PUBLICATION_ROLES or not re.fullmatch( GRAPH_PUBLICATION_ROLES[role], name @@ -105,14 +121,14 @@ def publication_inventory( class GraphPublicationQueue: """SQLite job records with leases; retained until publication or operator action.""" - def __init__(self, path: Path): + def __init__(self, path: Path) -> None: self.path = Path(path) self.path.parent.mkdir(parents=True, exist_ok=True) self._engine = create_spool_engine(self.path) upgrade_spool_database(self._engine) self._sessions = create_spool_session_factory(self._engine) - def enqueue(self, directory: Path, inventory: dict) -> None: + def enqueue(self, directory: Path, inventory: PublicationInventory) -> None: """Durably register exact bytes without a telemetry run or credential.""" directory = Path(directory).resolve() expected = publication_inventory( @@ -156,8 +172,13 @@ def claim( *, lease_seconds: float = GRAPH_PUBLICATION_LEASE_SECONDS, publication_id: str | None = None, - ) -> dict | None: - """Atomically lease one eligible job across concurrent emitter processes.""" + ) -> ClaimedGraphPublication | None: + """Reserve one pending job whose retry delay and previous lease expired. + + A conditional database update allows only one current lease. Workers + that exit leave their jobs eligible again after lease expiry; completion + uses the returned lease timestamp to reject outdated workers. + """ now = time.time() with self._sessions.begin() as session: statement = select(GraphPublicationJob.publication_id) @@ -186,18 +207,19 @@ def claim( if not result.rowcount: return None job = session.get(GraphPublicationJob, id_) - return { - "publication_id": id_, - "directory": job.directory, - "inventory": job.inventory, - "lease_until": job.lease_until, - } + return ClaimedGraphPublication( + publication_id=id_, + directory=Path(job.directory), + inventory=cast(PublicationInventory, job.inventory), + lease_until=job.lease_until, + ) - def receipt(self, id_: str) -> dict | None: + def receipt(self, id_: str) -> PublicationReceipt | None: """Read the latest local result; old HF receipt snapshots remain unchanged.""" with self._sessions() as session: job = session.get(GraphPublicationJob, id_) - return None if job is None else dict(job.receipt) + # ORM JSON columns are generic at the persistence boundary. + return None if job is None else cast(PublicationReceipt, dict(job.receipt)) def finish( self, @@ -235,7 +257,9 @@ def finish( "error_code": error_code, "url": url, "attempts": job.attempts, - "inventory_sha256": inventory_digest(job.inventory), + "inventory_sha256": inventory_digest( + cast(PublicationInventory, job.inventory) + ), } # This is a separate latest-result file, not the build/HF receipt. path = Path(job.directory) / GRAPH_PUBLICATION_STATUS_FILENAME @@ -265,7 +289,7 @@ def __init__( origin: str, client: httpx.Client | None = None, invalidate_credential: Callable[[], None] | None = None, - ): + ) -> None: parsed = urlsplit(origin) if parsed.scheme != "https" and not ( parsed.scheme == "http" @@ -292,12 +316,12 @@ def flush_once(self, *, publication_id: str | None = None) -> bool: if job is None: return False id_, inventory, directory = ( - job["publication_id"], - job["inventory"], - Path(job["directory"]), + job.publication_id, + job.inventory, + job.directory, ) error = "publication_unavailable" - lease = job["lease_until"] + lease = job.lease_until try: status, token = self.credential() if token is None: diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_publication.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_publication.py index c2bd2047a..7034e266b 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_publication.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_publication.py @@ -9,9 +9,17 @@ import json import re +from argparse import Namespace +from collections.abc import Iterator from pathlib import Path +from typing import Any, TypedDict from microcosm.build.artifact_files import file_artifact, materialize_bytes +from microcosm.build.graph_publication_types import ( + PublicationFileRole, + PublicationReceipt, +) +from microcosm.build.staging_dataset import StagedDatasetDelivery from microcosm.build.telemetry_emitter_service.graph_publication import ( publication_inventory, ) @@ -25,7 +33,15 @@ EVIDENCE_MANIFEST_NAME = "orrery.evidence-manifest.json" -def publication_enabled(args) -> bool: +class GraphEvidenceOutput(TypedDict): + """A preserved public file registered in the build's output manifest.""" + + path: str + sha256: str + bytes: int + + +def publication_enabled(args: Namespace) -> bool: explicit = getattr(args, "publish_orrery", None) if explicit is not None: return explicit @@ -36,7 +52,7 @@ def publication_enabled(args) -> bool: ) -def _urls(value): +def _urls(value: object) -> Iterator[str]: if isinstance(value, dict): for key, item in value.items(): if key == "url" and isinstance(item, str): @@ -48,7 +64,7 @@ def _urls(value): yield from _urls(item) -def _output(path, published_root): +def _output(path: Path, published_root: Path) -> GraphEvidenceOutput: info = file_artifact(path) return { "path": str(published_root / path.name), @@ -57,7 +73,7 @@ def _output(path, published_root): } -def finalize_graph(args, record: dict) -> dict: +def finalize_graph(args: Namespace, record: dict[str, Any]) -> PublicationReceipt: """Freeze bytes and a receipt snapshot; retry never rewrites that snapshot. This runs inside the temporary build directory, including on exceptions. @@ -68,7 +84,7 @@ def finalize_graph(args, record: dict) -> dict: output = Path(args.out) published_root = Path(getattr(args, "published_out", output)) schema_path = output / "graph.schema.json" - receipt = { + receipt: PublicationReceipt = { "version": 1, "status": "no_evidence", "publication_id": None, @@ -81,7 +97,10 @@ def finalize_graph(args, record: dict) -> dict: schema = json.loads(schema_path.read_bytes()) index_path = output / "execution.evidence.json" execution = None - roles = {schema_path.name: "schema", "graph.orrery.json": "graph"} + roles: dict[str, PublicationFileRole] = { + schema_path.name: "schema", + "graph.orrery.json": "graph", + } if index_path.is_file(): store = ContentStore(args.graph_store, create=False) runs = load_run_evidence(index_path, store=store, reference_base=output) @@ -89,12 +108,13 @@ def finalize_graph(args, record: dict) -> dict: roles[index_path.name] = "index" index = json.loads(index_path.read_bytes()) for run in index["runs"]: - for field, role in ( + evidence_roles: tuple[tuple[str, PublicationFileRole], ...] = ( ("graph", "declaration"), ("manifest", "manifest"), ("binding", "binding"), ("summaries", "summaries"), - ): + ) + for field, role in evidence_roles: roles[run[field]["path"]] = role payload = orrery_json_from_schema(schema, execution=execution).encode() materialize_bytes(payload, output / "graph.orrery.json") @@ -140,7 +160,9 @@ def finalize_graph(args, record: dict) -> dict: _cache_dir() / TELEMETRY_SPOOL_FILENAME ) queue.enqueue(preserved, inventory) - receipt.update(queue.receipt(inventory["publication_id"])) + queued_receipt = queue.receipt(inventory["publication_id"]) + assert queued_receipt is not None + receipt.update(queued_receipt) else: receipt.update(emitter.publish_graph(preserved, inventory)) except Exception: @@ -181,7 +203,9 @@ def finalize_graph(args, record: dict) -> dict: return receipt -def stage_graph_evidence(args, output: Path, *, run_id: str) -> dict | None: +def stage_graph_evidence( + args: Namespace, output: Path, *, run_id: str +) -> StagedDatasetDelivery | None: """Stage failure evidence through the existing generic HF bundle adapter.""" from microcosm.build.staging_dataset import ( StagedDatasetBundle, diff --git a/packages/microcosm-build/tests/engine_free/shared/test_graph_publication.py b/packages/microcosm-build/tests/engine_free/shared/test_graph_publication.py index 1696a56af..8be249899 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_graph_publication.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_graph_publication.py @@ -1,10 +1,12 @@ """Durable graph jobs use fake HTTP, independently of the visualization app.""" import json +from dataclasses import FrozenInstanceError import httpx import pytest +from microcosm.build.graph_publication_types import ClaimedGraphPublication from microcosm.build.telemetry_emitter_service.graph_publication import ( GraphPublicationDelivery, GraphPublicationQueue, @@ -26,12 +28,64 @@ def test_jobs_survive_restart_without_run_registration_or_event_retention(tmp_pa queue.enqueue(directory, inventory) restarted = GraphPublicationQueue(queue.path) job = restarted.claim() - assert job["inventory"] == inventory + assert isinstance(job, ClaimedGraphPublication) + assert job.inventory == inventory assert "run_id" not in inventory - assert job["publication_id"] == inventory["publication_id"] + assert job.publication_id == inventory["publication_id"] + assert job.directory == directory.resolve() assert restarted.claim() is None +def test_claim_is_a_frozen_record_and_expired_leases_can_be_reclaimed( + tmp_path, monkeypatch +): + from microcosm.build.telemetry_emitter_service.graph_publication import time + + directory, inventory = source(tmp_path) + queue = GraphPublicationQueue(tmp_path / "jobs.sqlite3") + queue.enqueue(directory, inventory) + clock = [1000.0] + monkeypatch.setattr(time, "time", lambda: clock[0]) + claimed = queue.claim(lease_seconds=5) + assert isinstance(claimed, ClaimedGraphPublication) + assert claimed.lease_until == 1005.0 + with pytest.raises(FrozenInstanceError): + claimed.lease_until = 0 + assert GraphPublicationQueue(queue.path).claim() is None + clock[0] = 1005.0 + replacement = GraphPublicationQueue(queue.path).claim() + assert replacement.publication_id == claimed.publication_id + assert replacement.lease_until > claimed.lease_until + + +def test_claim_respects_retry_delay_and_published_jobs_are_not_reclaimed( + tmp_path, monkeypatch +): + from microcosm.build.telemetry_emitter_service.graph_publication import time + + directory, inventory = source(tmp_path) + queue = GraphPublicationQueue(tmp_path / "jobs.sqlite3") + queue.enqueue(directory, inventory) + clock = [1000.0] + monkeypatch.setattr(time, "time", lambda: clock[0]) + claimed = queue.claim() + queue.finish( + claimed.publication_id, + error_code="publication_network_error", + expected_lease=claimed.lease_until, + ) + assert queue.claim() is None + clock[0] = 1002.0 + retried = queue.claim() + assert retried.publication_id == claimed.publication_id + queue.finish( + retried.publication_id, + url="https://runs.example/runs/published/", + expected_lease=retried.lease_until, + ) + assert queue.claim() is None + + def test_queue_rejects_changed_bytes_and_conflicting_inventory(tmp_path): directory, inventory = source(tmp_path) queue = GraphPublicationQueue(tmp_path / "jobs.sqlite3") @@ -162,14 +216,14 @@ def test_expired_worker_cannot_release_a_replacement_workers_lease(tmp_path): queue.finish( id_, error_code="publication_network_error", - expected_lease=expired["lease_until"], + expected_lease=expired.lease_until, ) assert queue.receipt(id_) == initial assert queue.claim() is None replacement.finish( id_, url=f"https://runs.example/runs/{id_}/", - expected_lease=current["lease_until"], + expected_lease=current.lease_until, ) assert queue.receipt(id_)["status"] == "published" From c6a725b8cf2577c69ebc85fa092fc9cefde0fb6b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 9 Oct 2026 01:51:05 +0400 Subject: [PATCH 5/6] Keep publication database and staging results type-safe --- .../build/telemetry_emitter_service/graph_publication.py | 8 ++++++-- .../src/microcosm/build/uk_runtime/orrery_publication.py | 4 +++- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/graph_publication.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/graph_publication.py index 293dc4d1b..92582bee4 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/graph_publication.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/graph_publication.py @@ -195,7 +195,7 @@ def claim( ) if id_ is None: return None - result = session.execute( + claimed_id = session.scalar( update(GraphPublicationJob) .where( GraphPublicationJob.publication_id == id_, @@ -203,10 +203,14 @@ def claim( GraphPublicationJob.lease_until <= now, ) .values(lease_until=now + lease_seconds) + .returning(GraphPublicationJob.publication_id) ) - if not result.rowcount: + if claimed_id is None: return None job = session.get(GraphPublicationJob, id_) + assert job is not None, ( + "A claimed publication must exist in its transaction." + ) return ClaimedGraphPublication( publication_id=id_, directory=Path(job.directory), diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_publication.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_publication.py index 7034e266b..bd5e876d2 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_publication.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_publication.py @@ -168,7 +168,8 @@ def finalize_graph(args: Namespace, record: dict[str, Any]) -> PublicationReceip except Exception: # No credential, exception text, paths to raw input data, or signed URL # is copied from an exception into public evidence or telemetry. - receipt.update(status="failed", error_code="graph_export_or_enqueue_failed") + receipt["status"] = "failed" + receipt["error_code"] = "graph_export_or_enqueue_failed" materialize_bytes(canonical_json(receipt), output / RECEIPT_NAME) names.append(RECEIPT_NAME) record["orrery_files"] = names @@ -233,6 +234,7 @@ def stage_graph_evidence( ) if mode == "local_only": return local_only_staged_dataset(bundle, prefix=UK_STAGED_DATASET_PREFIX) + assert repository is not None, "Remote staging requires a repository." return stage_bundle( bundle, storage=HuggingFaceDatasetStorage(repository, api=rowwise_staging._hub_api()), From 95c4cc9c1e076c7ac4f8fd261dd53f377f208eaa Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 9 Oct 2026 18:36:43 +0400 Subject: [PATCH 6/6] Keep failed-build evidence out of the output and graph delivery running A failed UK build now publishes its graph evidence to orrery-failure/ in the attempt evidence directory instead of the output directory, so an earlier successful bundle and its completion marker stay consistent. Interrupts keep local evidence and the queue entry without waiting for an upload or staging to Hugging Face, and a failure while saving evidence no longer replaces the build's own error. The graph worker survives local queue errors. A job whose preserved directory was removed stays pending with preserved_files_missing; the status file is written after the queue commit, and claims take the earliest-due job first. Remove the unused socket publication action, share origin validation with the collector, return typed collector credentials and graph finalization results, type the ORM publication columns, and replace private cross-module access with public accessors. Update the publication guide and add regression tests. Co-Authored-By: Claude Opus 5.5 --- .../1150-publication-failures.fixed.md | 1 + docs/orrery-publication.md | 13 +- .../microcosm/build/graph_publication_cli.py | 12 +- .../build/graph_publication_types.py | 3 + .../src/microcosm/build/telemetry_emitter.py | 15 +- .../build/telemetry_emitter_service/auth.py | 57 ++++--- .../telemetry_emitter_service/collector.py | 8 +- .../telemetry_emitter_service/constants.py | 6 +- .../graph_publication.py | 77 +++++---- .../build/telemetry_emitter_service/models.py | 13 +- .../telemetry_emitter_service/runtime.py | 39 +++-- .../build/uk_runtime/full_build_cli.py | 158 ++++++++++++------ .../build/uk_runtime/orrery_publication.py | 62 ++++--- .../build/uk_runtime/rowwise_staging.py | 6 + .../shared/test_graph_publication.py | 59 +++++++ .../shared/test_publication_workers.py | 37 ++++ .../uk/test_uk_orrery_publication.py | 109 ++++++++++-- 17 files changed, 502 insertions(+), 173 deletions(-) create mode 100644 changelog.d/1150-publication-failures.fixed.md diff --git a/changelog.d/1150-publication-failures.fixed.md b/changelog.d/1150-publication-failures.fixed.md new file mode 100644 index 000000000..497aacf3b --- /dev/null +++ b/changelog.d/1150-publication-failures.fixed.md @@ -0,0 +1 @@ +Keep a failed UK build's graph evidence out of the output directory so an earlier successful bundle stays intact, and keep graph delivery running when a preserved directory is removed or the local queue reports an error. diff --git a/docs/orrery-publication.md b/docs/orrery-publication.md index 246699191..c7e6c65d9 100644 --- a/docs/orrery-publication.md +++ b/docs/orrery-publication.md @@ -45,7 +45,13 @@ credentials and caches and use fake HTTP or explicit loopback services. without remote submission. `--publish-orrery` explicitly enables submission; `--no-publish-orrery` disables it independently. A publication failure does not replace the dataset's validation result or prevent HF artifact staging. + Failed builds preserve recorded graph evidence before temporary-file cleanup. +They publish it, with an `orrery.failure.json` marker, to `orrery-failure/` in +the attempt evidence directory named by `failure.json`, and leave the output +directory, including an earlier successful bundle, unchanged. An interrupted +build preserves the same local evidence and queue entry but does not wait for +an upload or stage to Hugging Face. ## Retry @@ -57,5 +63,8 @@ microcosm-publish-graph --publication-id PUBLICATION_ID \ The preserved directory contains `orrery.upload.json` and the exact files. Retries update a separate `publication.status.json`, never the original HF receipt snapshot. Pending jobs also resume when a subsequent local service -starts. Package extraction into `microcosm-local-provider` is a separate, -deferred consumer change; this implementation does not require those packages. +starts. A job whose preserved directory has been removed stays pending with +error code `preserved_files_missing` until its files are restored to the same +directory; it does not stop delivery of other jobs. Moving this +service into the `microcosm-emitter` package is a separate, deferred consumer +change; this implementation does not require that package. diff --git a/packages/microcosm-build/src/microcosm/build/graph_publication_cli.py b/packages/microcosm-build/src/microcosm/build/graph_publication_cli.py index 95f6d2e20..89318d838 100644 --- a/packages/microcosm-build/src/microcosm/build/graph_publication_cli.py +++ b/packages/microcosm-build/src/microcosm/build/graph_publication_cli.py @@ -6,11 +6,10 @@ import json from pathlib import Path -from microcosm.build.telemetry_emitter import _cache_dir -from microcosm.build.telemetry_emitter_constants import TELEMETRY_SPOOL_FILENAME +from microcosm.build.telemetry_emitter import default_spool_path from microcosm.build.telemetry_emitter_service.auth import ( CollectorSession, - _collector_origin, + validated_origin, ) from microcosm.build.telemetry_emitter_service.constants import PRODUCTION_COLLECTOR_URL from microcosm.build.telemetry_emitter_service.graph_publication import ( @@ -22,9 +21,7 @@ def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--spool", type=Path, default=_cache_dir() / TELEMETRY_SPOOL_FILENAME - ) + parser.add_argument("--spool", type=Path, default=default_spool_path()) parser.add_argument("--publication-id", required=True) parser.add_argument( "--directory", @@ -42,7 +39,7 @@ def main(argv: list[str] | None = None) -> int: if queue.receipt(args.publication_id) is None: parser.error("Publication job not found; supply its preserved directory.") queue.retry(args.publication_id) - session = CollectorSession(_collector_origin(PRODUCTION_COLLECTOR_URL)) + session = CollectorSession(validated_origin(PRODUCTION_COLLECTOR_URL)) delivery = GraphPublicationDelivery( queue, credential=session.credential, @@ -52,6 +49,7 @@ def main(argv: list[str] | None = None) -> int: # Explicit retries process this job only, not unrelated pending jobs. delivery.flush_once(publication_id=args.publication_id) receipt = queue.receipt(args.publication_id) + assert receipt is not None, "A queued graph publication must have a receipt." print(json.dumps(receipt, sort_keys=True)) return 0 if receipt["status"] == "published" else 1 diff --git a/packages/microcosm-build/src/microcosm/build/graph_publication_types.py b/packages/microcosm-build/src/microcosm/build/graph_publication_types.py index 94f4452f3..a0f03f6a8 100644 --- a/packages/microcosm-build/src/microcosm/build/graph_publication_types.py +++ b/packages/microcosm-build/src/microcosm/build/graph_publication_types.py @@ -16,6 +16,9 @@ ] +type GraphJobStatus = Literal["pending", "published"] + + class PublicationFile(TypedDict): """One public file, identified by its exact bytes and allowed role.""" diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter.py index dd186f8d0..336b846fe 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter.py @@ -91,6 +91,11 @@ def _cache_dir() -> Path: return root.joinpath(*TELEMETRY_CACHE_PARTS) +def default_spool_path() -> Path: + """The local SQLite queue shared by telemetry events and graph publication.""" + return _cache_dir() / TELEMETRY_SPOOL_FILENAME + + @dataclass(frozen=True) class TelemetryRun: """Identity registered with the hosted collector for one build.""" @@ -183,11 +188,7 @@ def start( ) runtime_dir.chmod(0o700) socket_path = runtime_dir / SOCKET_FILENAME - queue_path = ( - Path(spool_path) - if spool_path - else _cache_dir() / TELEMETRY_SPOOL_FILENAME - ) + queue_path = Path(spool_path) if spool_path else default_spool_path() command = [ sys.executable, "-m", @@ -480,9 +481,7 @@ def publish_graph( GraphPublicationQueue, ) - queue = GraphPublicationQueue( - self._spool_path or _cache_dir() / TELEMETRY_SPOOL_FILENAME - ) + queue = GraphPublicationQueue(self._spool_path or default_spool_path()) queue.enqueue(directory, inventory) deadline = time.monotonic() + max(0, wait_seconds) while self.available and time.monotonic() < deadline: diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/auth.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/auth.py index a8e0726f7..205fef3dd 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/auth.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/auth.py @@ -4,20 +4,19 @@ import json import os +import threading import time import urllib.error import urllib.request -from collections.abc import Mapping +from collections.abc import Callable, Mapping from http import HTTPStatus -from typing import Any +from typing import Any, Literal, NamedTuple from urllib.parse import urlsplit from huggingface_hub import get_token from microcosm.build.telemetry_emitter_service.constants import ( COLLECTOR_RESPONSE_TOO_LARGE_ERROR, - COLLECTOR_URL_HTTPS_ERROR, - COLLECTOR_URL_ORIGIN_ERROR, DEFAULT_TOKEN_LIFETIME_SECONDS, DEVELOPMENT_COLLECTOR_LOOPBACK_ERROR, HTTP_TIMEOUT_SECONDS, @@ -25,6 +24,8 @@ LOOPBACK_HOSTS, MAX_HTTP_RESPONSE_BYTES, MINIMUM_TOKEN_LIFETIME_SECONDS, + SERVICE_ORIGIN_FORMAT_ERROR, + SERVICE_ORIGIN_HTTPS_ERROR, TOKEN_EXCHANGE_PATH, TOKEN_REFRESH_MARGIN_SECONDS, ) @@ -37,14 +38,25 @@ def redirect_request(self, req, fp, code, msg, headers, newurl): return None -def _collector_origin(value: str, *, allow_loopback_http: bool = False) -> str: +type CredentialStatus = Literal["ok", "missing", "rejected", "unavailable"] + + +class CollectorCredential(NamedTuple): + """A session token, or the reason none is available for this attempt.""" + + status: CredentialStatus + token: str | None + + +def validated_origin(value: str, *, allow_loopback_http: bool = False) -> str: + """Return an HTTPS origin, or plain HTTP only for an allowed loopback host.""" parsed = urlsplit(value) loopback = parsed.hostname in LOOPBACK_HOSTS valid_scheme = parsed.scheme == "https" or ( allow_loopback_http and loopback and parsed.scheme == "http" ) if not parsed.hostname or not valid_scheme: - raise ValueError(COLLECTOR_URL_HTTPS_ERROR) + raise ValueError(SERVICE_ORIGIN_HTTPS_ERROR) if ( parsed.username or parsed.password @@ -52,7 +64,7 @@ def _collector_origin(value: str, *, allow_loopback_http: bool = False) -> str: or parsed.fragment or parsed.path not in {"", "/"} ): - raise ValueError(COLLECTOR_URL_ORIGIN_ERROR) + raise ValueError(SERVICE_ORIGIN_FORMAT_ERROR) return value.rstrip("/") @@ -60,7 +72,7 @@ def _development_collector_url(value: str) -> str: parsed = urlsplit(value) if parsed.hostname not in LOOPBACK_HOSTS: raise ValueError(DEVELOPMENT_COLLECTOR_LOOPBACK_ERROR) - return _collector_origin(value, allow_loopback_http=True) + return validated_origin(value, allow_loopback_http=True) def _decode_response(body: bytes) -> dict[str, Any]: @@ -102,6 +114,9 @@ def _http_post( return error.code, _decode_response(body) +type HttpPost = Callable[[str, Mapping[str, Any], str], tuple[int, dict[str, Any]]] + + def _huggingface_token() -> str | None: return ( os.environ.get("HF_TOKEN", "").strip() @@ -113,9 +128,13 @@ def _huggingface_token() -> str | None: class CollectorSession: """One thread-safe short-lived credential shared by independent publishers.""" - def __init__(self, origin: str, *, post=None, credential=None): - import threading - + def __init__( + self, + origin: str, + *, + post: HttpPost | None = None, + credential: Callable[[], str | None] | None = None, + ) -> None: self.origin = origin self._post = post or _http_post self._credential = credential or _huggingface_token @@ -126,23 +145,23 @@ def invalidate(self) -> None: with self._lock: self._session_token = None - def credential(self) -> tuple[str, str | None]: + def credential(self) -> CollectorCredential: with self._lock: if ( self._session_token is not None and self._session_token[1] > time.monotonic() + TOKEN_REFRESH_MARGIN_SECONDS ): - return "ok", self._session_token[0] + return CollectorCredential("ok", self._session_token[0]) hf_token = self._credential() if not hf_token: - return "missing", None + return CollectorCredential("missing", None) try: status, response = self._post( self.origin + TOKEN_EXCHANGE_PATH, {}, hf_token ) except (OSError, TimeoutError): - return "unavailable", None + return CollectorCredential("unavailable", None) if status == HTTPStatus.OK and isinstance( response.get("access_token"), str ): @@ -152,10 +171,10 @@ def credential(self) -> tuple[str, str | None]: int(response.get("expires_in", DEFAULT_TOKEN_LIFETIME_SECONDS)), ) except (ValueError, TypeError): - return "unavailable", None + return CollectorCredential("unavailable", None) token = response["access_token"] self._session_token = (token, time.monotonic() + expires_in) - return "ok", token + return CollectorCredential("ok", token) if status in {HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN}: - return "rejected", None - return "unavailable", None + return CollectorCredential("rejected", None) + return CollectorCredential("unavailable", None) diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/collector.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/collector.py index 1b96dfab7..8bc19a3b7 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/collector.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/collector.py @@ -9,11 +9,12 @@ from typing import Any from microcosm.build.telemetry_emitter_service.auth import ( + CollectorCredential, CollectorSession, - _collector_origin, _development_collector_url, _http_post, _huggingface_token, + validated_origin, ) from microcosm.build.telemetry_emitter_service.constants import ( INITIAL_RETRY_SECONDS, @@ -32,7 +33,6 @@ ) from microcosm.build.telemetry_emitter_service.spool import EventSpool - #: Client errors worth retrying: a timeout and a rate limit pass with time. #: Every other 4xx is the collector's settled answer to the same request, so #: retrying it only wedges the queue (a payload shape it does not accept, say). @@ -62,7 +62,7 @@ def __init__( self.collector_url = ( _development_collector_url(development_collector_url) if development_collector_url is not None - else _collector_origin(PRODUCTION_COLLECTOR_URL) + else validated_origin(PRODUCTION_COLLECTOR_URL) ) self.spool = spool self.session = session or CollectorSession( @@ -137,7 +137,7 @@ def _collector_token(self, registration: Mapping[str, Any]) -> str | None: self._defer_retry() return token - def session_credential(self) -> tuple[str, str | None]: + def session_credential(self) -> CollectorCredential: """Compatibility accessor; authentication has no telemetry queue policy.""" return self.session.credential() diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/constants.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/constants.py index 83163fe53..bd7559467 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/constants.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/constants.py @@ -42,9 +42,9 @@ "events (HTTP {status}), so they will not be retried. The dataset build will " "continue." ) -COLLECTOR_URL_HTTPS_ERROR: Final = "collector URL must be an HTTPS origin" -COLLECTOR_URL_ORIGIN_ERROR: Final = ( - "collector URL must be an origin without credentials or path data" +SERVICE_ORIGIN_HTTPS_ERROR: Final = "service URL must be an HTTPS origin" +SERVICE_ORIGIN_FORMAT_ERROR: Final = ( + "service URL must be an origin without credentials or path data" ) DEVELOPMENT_COLLECTOR_LOOPBACK_ERROR: Final = ( "development collector URL must use a loopback address" diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/graph_publication.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/graph_publication.py index 92582bee4..f1f95ed13 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/graph_publication.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/graph_publication.py @@ -13,7 +13,6 @@ import uuid from collections.abc import Callable, Mapping from pathlib import Path -from typing import cast from urllib.parse import urlsplit import httpx @@ -26,6 +25,10 @@ PublicationInventory, PublicationReceipt, ) +from microcosm.build.telemetry_emitter_service.auth import ( + CollectorCredential, + validated_origin, +) from microcosm.build.telemetry_emitter_service.database import ( create_spool_engine, create_spool_session_factory, @@ -35,10 +38,13 @@ MAX_GRAPH_FILE_BYTES = 64 * 1024 * 1024 MAX_GRAPH_TOTAL_BYTES = 512 * 1024 * 1024 -GRAPH_PUBLICATION_ACTION = "graph_publication" GRAPH_PUBLICATION_WAIT_SECONDS = 30.0 GRAPH_PUBLICATION_LEASE_SECONDS = 300.0 GRAPH_PUBLICATION_STATUS_FILENAME = "publication.status.json" +GRAPH_WORKER_ERROR_MESSAGE = ( + "Microcosm graph publication hit a local queue error; pending jobs remain " + "queued for a later attempt." +) DEFAULT_GRAPH_PUBLICATION_ORIGIN = "https://microcosm-runs.vercel.app" GRAPH_PUBLICATION_ROLES = { "graph": r"graph\.orrery\.json", @@ -118,6 +124,16 @@ def publication_inventory( return {"version": 1, "publication_id": id_, "files": files} +def _write_status(directory: Path, receipt: PublicationReceipt) -> None: + path = directory / GRAPH_PUBLICATION_STATUS_FILENAME + temporary = path.with_suffix(f".{uuid.uuid4().hex}.tmp") + try: + temporary.write_text(json.dumps(receipt, sort_keys=True) + "\n") + temporary.replace(path) + except OSError: + temporary.unlink(missing_ok=True) + + class GraphPublicationQueue: """SQLite job records with leases; retained until publication or operator action.""" @@ -191,7 +207,11 @@ def claim( GraphPublicationJob.status == "pending", GraphPublicationJob.next_attempt_at <= now, GraphPublicationJob.lease_until <= now, - ).limit(1) + ) + # Earliest-due first, so a job that keeps failing cannot + # repeatedly displace other pending publications. + .order_by(GraphPublicationJob.next_attempt_at) + .limit(1) ) if id_ is None: return None @@ -214,7 +234,7 @@ def claim( return ClaimedGraphPublication( publication_id=id_, directory=Path(job.directory), - inventory=cast(PublicationInventory, job.inventory), + inventory=job.inventory, lease_until=job.lease_until, ) @@ -222,8 +242,7 @@ def receipt(self, id_: str) -> PublicationReceipt | None: """Read the latest local result; old HF receipt snapshots remain unchanged.""" with self._sessions() as session: job = session.get(GraphPublicationJob, id_) - # ORM JSON columns are generic at the persistence boundary. - return None if job is None else cast(PublicationReceipt, dict(job.receipt)) + return None if job is None else job.receipt def finish( self, @@ -254,22 +273,21 @@ def finish( if job is None: return job.next_attempt_at = time.time() + min(300, 2 ** min(job.attempts, 8)) - job.receipt = { + receipt: PublicationReceipt = { "version": 1, "publication_id": id_, "status": job.status, "error_code": error_code, "url": url, "attempts": job.attempts, - "inventory_sha256": inventory_digest( - cast(PublicationInventory, job.inventory) - ), + "inventory_sha256": inventory_digest(job.inventory), } - # This is a separate latest-result file, not the build/HF receipt. - path = Path(job.directory) / GRAPH_PUBLICATION_STATUS_FILENAME - temporary = path.with_suffix(f".{uuid.uuid4().hex}.tmp") - temporary.write_text(json.dumps(job.receipt, sort_keys=True) + "\n") - temporary.replace(path) + job.receipt = receipt + directory = Path(job.directory) + # The queue row is authoritative. This separate latest-result file (not + # the build/HF receipt) is a convenience copy, written only after the + # commit: a deleted preserved directory must not roll back the result. + _write_status(directory, receipt) def retry(self, id_: str | None = None) -> None: """Make preserved pending jobs eligible again, including missing credentials.""" @@ -289,28 +307,14 @@ def __init__( self, queue: GraphPublicationQueue, *, - credential: Callable[[], tuple[str, str | None]], + credential: Callable[[], CollectorCredential], origin: str, client: httpx.Client | None = None, invalidate_credential: Callable[[], None] | None = None, ) -> None: - parsed = urlsplit(origin) - if parsed.scheme != "https" and not ( - parsed.scheme == "http" - and parsed.hostname in {"localhost", "127.0.0.1", "::1"} - ): - raise ValueError( - "Graph publication origin must be HTTPS or local development." - ) - if ( - parsed.username - or parsed.password - or parsed.query - or parsed.fragment - or parsed.path not in {"", "/"} - ): - raise ValueError("Graph publication destination must be an origin.") - self.queue, self.credential, self.origin = queue, credential, origin.rstrip("/") + self.queue = queue + self.credential = credential + self.origin = validated_origin(origin, allow_loopback_http=True) self.client = client or httpx.Client(timeout=60, follow_redirects=False) self.invalidate_credential = invalidate_credential @@ -326,6 +330,13 @@ def flush_once(self, *, publication_id: str | None = None) -> bool: ) error = "publication_unavailable" lease = job.lease_until + if not directory.is_dir(): + # Retained until an operator restores the files to this directory; + # the retry backoff lets other pending jobs proceed meanwhile. + self.queue.finish( + id_, error_code="preserved_files_missing", expected_lease=lease + ) + return False try: status, token = self.credential() if token is None: diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/models.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/models.py index d49bb8277..610576a94 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/models.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/models.py @@ -16,6 +16,11 @@ from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship from sqlalchemy.types import TypeDecorator +from microcosm.build.graph_publication_types import ( + GraphJobStatus, + PublicationInventory, + PublicationReceipt, +) from microcosm.build.telemetry_emitter_service.constants import ( UPLOAD_STATE_PENDING, ) @@ -136,10 +141,12 @@ class GraphPublicationJob(SpoolModel): __tablename__ = "graph_publication_jobs" publication_id: Mapped[str] = mapped_column(Text, primary_key=True) - inventory: Mapped[dict[str, Any]] = mapped_column(JsonObjectText, nullable=False) + inventory: Mapped[PublicationInventory] = mapped_column( + JsonObjectText, nullable=False + ) directory: Mapped[str] = mapped_column(Text, nullable=False) - status: Mapped[str] = mapped_column(Text, nullable=False) + status: Mapped[GraphJobStatus] = mapped_column(Text, nullable=False) attempts: Mapped[int] = mapped_column(Integer, nullable=False) next_attempt_at: Mapped[float] = mapped_column(Float, nullable=False) lease_until: Mapped[float] = mapped_column(Float, nullable=False) - receipt: Mapped[dict[str, Any]] = mapped_column(JsonObjectText, nullable=False) + receipt: Mapped[PublicationReceipt] = mapped_column(JsonObjectText, nullable=False) diff --git a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/runtime.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/runtime.py index 3794d65a1..3a95be815 100644 --- a/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/runtime.py +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/runtime.py @@ -5,6 +5,7 @@ import json import os import socket +import sys import threading import time from collections.abc import Mapping @@ -26,7 +27,7 @@ WORKER_INTERVAL_SECONDS, ) from microcosm.build.telemetry_emitter_service.graph_publication import ( - GRAPH_PUBLICATION_ACTION, + GRAPH_WORKER_ERROR_MESSAGE, GraphPublicationDelivery, ) from microcosm.build.telemetry_emitter_service.resources import ProcessTreeSampler @@ -122,10 +123,20 @@ def run(self) -> None: server.settimeout(SOCKET_ACCEPT_TIMEOUT_SECONDS) workers = [] if self.registration is not None: - workers.append(threading.Thread(target=self._worker, daemon=True)) + workers.append( + threading.Thread( + target=self._worker, + args=(self.registration,), + daemon=True, + ) + ) if self.graph_delivery is not None: workers.append( - threading.Thread(target=self._graph_worker, daemon=True) + threading.Thread( + target=self._graph_worker, + args=(self.graph_delivery,), + daemon=True, + ) ) for worker in workers: worker.start() @@ -190,11 +201,7 @@ def _read_message(connection: socket.socket) -> dict[str, Any]: def _handle(self, message: Mapping[str, Any]) -> None: action = message.get("action") - if action == GRAPH_PUBLICATION_ACTION and self.graph_delivery is not None: - self.graph_delivery.queue.enqueue( - Path(message["directory"]), message["inventory"] - ) - elif action == ACTION_EVENT and self.registration is not None: + if action == ACTION_EVENT and self.registration is not None: event = message.get("event") if not isinstance(event, Mapping): raise ValueError(EVENT_OBJECT_ERROR) @@ -217,7 +224,7 @@ def _handle(self, message: Mapping[str, Any]) -> None: else: raise ValueError(UNSUPPORTED_ACTION_ERROR) - def _worker(self) -> None: + def _worker(self, registration: Mapping[str, Any]) -> None: next_heartbeat = time.monotonic() + self.heartbeat_seconds while not self._stop.wait(WORKER_INTERVAL_SECONDS): now = time.monotonic() @@ -226,7 +233,7 @@ def _worker(self) -> None: self.sampler.sample() if now >= next_heartbeat: self.spool.append( - self.registration, + registration, _heartbeat_event(self._last_stage), resources=self.sampler.sample(), ) @@ -234,11 +241,19 @@ def _worker(self) -> None: self.delivery.flush_once() self._drain() - def _graph_worker(self) -> None: + def _graph_worker(self, delivery: GraphPublicationDelivery) -> None: # Uploads may take much longer than a heartbeat. They never run on the # telemetry worker or the local socket thread. Pending jobs survive exit. + warned = False while not self._stop.wait(WORKER_INTERVAL_SECONDS): - self.graph_delivery.flush_once() + try: + delivery.flush_once() + except Exception: + # A database or filesystem error leaves the job leased for a + # later retry; it must not end delivery of every other job. + if not warned: + print(GRAPH_WORKER_ERROR_MESSAGE, file=sys.stderr, flush=True) + warned = True def _drain(self) -> None: deadline = time.monotonic() + self.drain_seconds diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/full_build_cli.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/full_build_cli.py index 70725f23a..69cdfef27 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/full_build_cli.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/full_build_cli.py @@ -45,6 +45,7 @@ from dataclasses import asdict, dataclass, replace from datetime import UTC, date, datetime from pathlib import Path +from typing import TYPE_CHECKING, Protocol import numpy as np @@ -209,6 +210,9 @@ from .size_checkpoint import uk_size_checkpoint_identity from .staging import UK_STAGED_DATASET_REPOSITORY, UK_STAGING_REPOSITORY +if TYPE_CHECKING: + from microcosm.build.staging_v2 import StagingRunBundleWriterV2 + # The names the rowwise tool's tests call on either driver. _validate_cli_args = validate_cli_args _refuse_national_role_arguments = refuse_national_role_arguments @@ -1224,69 +1228,129 @@ def _release_stem(posture) -> tuple[str, str]: ) -def _run_with_graph_publication(execute, prepared, args, *, telemetry, attempt, record): +class _BuildKernel[Prepared](Protocol): + """A UK build body run inside its temporary output directory.""" + + def __call__( + self, + prepared: Prepared, + args: argparse.Namespace, + *, + telemetry: StagingRunBundleWriterV2 | None = None, + attempt: dict | None = None, + record: dict | None = None, + ) -> int: ... + + +def _run_with_graph_publication[Prepared]( + execute: _BuildKernel[Prepared], + prepared: Prepared, + args: argparse.Namespace, + *, + telemetry: StagingRunBundleWriterV2 | None, + attempt: dict | None, + record: dict, +) -> int: """Export before temporary files disappear, even if a kernel raises.""" - from .orrery_publication import ( - EVIDENCE_MANIFEST_NAME, - finalize_graph, - stage_graph_evidence, - ) + from .orrery_publication import finalize_graph try: status = execute( prepared, args, telemetry=telemetry, attempt=attempt, record=record ) - except BaseException: - finalize_graph(args, record) - # A failed build may have no dataset. Preserve only its explicit graph - # inventory, not incomplete H5 files or arbitrary temporary JSON. - from microcosm.build.artifact_files import publish_staged_bundle - - output = Path(args.published_out) - names = [*record.get("orrery_files", []), EVIDENCE_MANIFEST_NAME] - staged = { - name: args.out / name for name in names if (args.out / name).is_file() - } - marker = args.out / "orrery.failure.json" - materialize_bytes( - canonical_json( - { - "schema_version": 1, - "kind": "uk_build_failed", - "releasable": False, - "outputs": { - name: file_artifact(path) for name, path in staged.items() - }, - } - ), - marker, - ) - staged["manifest"] = marker - publish_staged_bundle( - staged, - {role: output / path.name for role, path in staged.items()}, - completion_role="manifest", - ) + except BaseException as failure: + # An interrupt still preserves local evidence and its durable queue + # entry, but does not wait for an upload or start Hugging Face staging. + interrupted = not isinstance(failure, Exception) try: - stage_graph_evidence( - args, - output, - run_id=( - telemetry.run_id - if telemetry is not None - else Path(args.attempt_evidence).name - ), + _publish_failure_evidence( + args, record, telemetry=telemetry, interrupted=interrupted ) except Exception as error: print( - f"Graph evidence staging failed ({type(error).__name__}); files remain local.", + f"Graph failure evidence was not preserved ({type(error).__name__}); " + "the build's own error follows.", file=sys.stderr, ) raise - finalize_graph(args, record) + finalize_graph(args, manifest=record.get("manifest")) return status +def _publish_failure_evidence( + args: argparse.Namespace, + record: dict, + *, + telemetry: StagingRunBundleWriterV2 | None, + interrupted: bool, +) -> None: + """Publish a failed build's graph evidence beside its attempt evidence. + + A failed build leaves the output directory, including any earlier + successful bundle and its completion marker, unchanged. The graph files go + to a new directory under this attempt's durable evidence, which + ``failure.json`` already names. Incomplete H5 files and arbitrary + temporary JSON are never published. + """ + from microcosm.build.artifact_files import publish_staged_bundle + from microcosm.build.telemetry_emitter_service.graph_publication import ( + GRAPH_PUBLICATION_WAIT_SECONDS, + ) + + from .orrery_publication import ( + EVIDENCE_MANIFEST_NAME, + FAILURE_BUNDLE_DIRECTORY, + FAILURE_MARKER_NAME, + finalize_graph, + stage_graph_evidence, + ) + + output = Path(args.attempt_evidence) / FAILURE_BUNDLE_DIRECTORY + graph = finalize_graph( + args, + manifest=record.get("manifest"), + published_root=output, + wait_seconds=0 if interrupted else GRAPH_PUBLICATION_WAIT_SECONDS, + ) + names = [*graph.files, EVIDENCE_MANIFEST_NAME] + staged = {name: args.out / name for name in names if (args.out / name).is_file()} + marker = args.out / FAILURE_MARKER_NAME + materialize_bytes( + canonical_json( + { + "schema_version": 1, + "kind": "uk_build_failed", + "releasable": False, + "outputs": {name: file_artifact(path) for name, path in staged.items()}, + } + ), + marker, + ) + staged["manifest"] = marker + publish_staged_bundle( + staged, + {role: output / path.name for role, path in staged.items()}, + completion_role="manifest", + ) + if interrupted: + return + try: + stage_graph_evidence( + args, + output, + run_id=( + telemetry.run_id + if telemetry is not None + else Path(args.attempt_evidence).name + ), + ) + except Exception as error: + print( + f"Graph evidence staging failed ({type(error).__name__}); files remain local.", + file=sys.stderr, + ) + + def execute_full_build( prepared: PreparedUKFullBuild, args: argparse.Namespace, diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_publication.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_publication.py index bd5e876d2..81422e9dc 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_publication.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_publication.py @@ -10,7 +10,8 @@ import json import re from argparse import Namespace -from collections.abc import Iterator +from collections.abc import Iterator, MutableMapping +from dataclasses import dataclass from pathlib import Path from typing import Any, TypedDict @@ -20,7 +21,10 @@ PublicationReceipt, ) from microcosm.build.staging_dataset import StagedDatasetDelivery +from microcosm.build.telemetry_emitter import default_spool_path from microcosm.build.telemetry_emitter_service.graph_publication import ( + GRAPH_PUBLICATION_WAIT_SECONDS, + GraphPublicationQueue, publication_inventory, ) from microcosm.graph import ContentStore @@ -31,6 +35,8 @@ RECEIPT_NAME = "orrery.publication.json" INVENTORY_NAME = "orrery.upload.json" EVIDENCE_MANIFEST_NAME = "orrery.evidence-manifest.json" +FAILURE_MARKER_NAME = "orrery.failure.json" +FAILURE_BUNDLE_DIRECTORY = "orrery-failure" class GraphEvidenceOutput(TypedDict): @@ -41,6 +47,14 @@ class GraphEvidenceOutput(TypedDict): bytes: int +@dataclass(frozen=True) +class GraphFinalization: + """The build-directory files written for one graph and its receipt snapshot.""" + + files: tuple[str, ...] + receipt: PublicationReceipt + + def publication_enabled(args: Namespace) -> bool: explicit = getattr(args, "publish_orrery", None) if explicit is not None: @@ -73,16 +87,27 @@ def _output(path: Path, published_root: Path) -> GraphEvidenceOutput: } -def finalize_graph(args: Namespace, record: dict[str, Any]) -> PublicationReceipt: +def finalize_graph( + args: Namespace, + *, + manifest: MutableMapping[str, Any] | None, + published_root: Path | None = None, + wait_seconds: float = GRAPH_PUBLICATION_WAIT_SECONDS, +) -> GraphFinalization: """Freeze bytes and a receipt snapshot; retry never rewrites that snapshot. This runs inside the temporary build directory, including on exceptions. - Graph publication failure does not replace a dataset's scientific verdict. + A dataset ``manifest``, when the build produced one, registers the graph + files and receipt. Graph publication failure does not replace a dataset's + scientific verdict. ``published_root`` is where the registered files will + be published, by default the build's output directory. ``wait_seconds`` + bounds the wait for a first upload result after the job is durably queued. """ from . import rowwise_staging output = Path(args.out) - published_root = Path(getattr(args, "published_out", output)) + if published_root is None: + published_root = Path(getattr(args, "published_out", output)) schema_path = output / "graph.schema.json" receipt: PublicationReceipt = { "version": 1, @@ -144,27 +169,23 @@ def finalize_graph(args: Namespace, record: dict[str, Any]) -> PublicationReceip "recorded_execution": execution is not None, } if publication_enabled(args): - emitter = rowwise_staging._ACTIVE_EMITTER + emitter = rowwise_staging.active_telemetry_emitter() if emitter is None: # Durable enqueue is still possible without a running # emitter; the retry command requires no telemetry run. - from microcosm.build.telemetry_emitter import _cache_dir - from microcosm.build.telemetry_emitter_constants import ( - TELEMETRY_SPOOL_FILENAME, - ) - from microcosm.build.telemetry_emitter_service.graph_publication import ( - GraphPublicationQueue, - ) - - queue = GraphPublicationQueue( - _cache_dir() / TELEMETRY_SPOOL_FILENAME - ) + queue = GraphPublicationQueue(default_spool_path()) queue.enqueue(preserved, inventory) queued_receipt = queue.receipt(inventory["publication_id"]) - assert queued_receipt is not None + assert queued_receipt is not None, ( + "A queued graph publication must have a receipt." + ) receipt.update(queued_receipt) else: - receipt.update(emitter.publish_graph(preserved, inventory)) + receipt.update( + emitter.publish_graph( + preserved, inventory, wait_seconds=wait_seconds + ) + ) except Exception: # No credential, exception text, paths to raw input data, or signed URL # is copied from an exception into public evidence or telemetry. @@ -172,9 +193,6 @@ def finalize_graph(args: Namespace, record: dict[str, Any]) -> PublicationReceip receipt["error_code"] = "graph_export_or_enqueue_failed" materialize_bytes(canonical_json(receipt), output / RECEIPT_NAME) names.append(RECEIPT_NAME) - record["orrery_files"] = names - record["orrery_publication"] = receipt - manifest = record.get("manifest") if manifest is not None: from .rowwise_cli import MANIFEST_FILENAME, json_text @@ -201,7 +219,7 @@ def finalize_graph(args: Namespace, record: dict[str, Any]) -> PublicationReceip "outputs": {name: _output(output / name, published_root) for name in names}, } materialize_bytes(canonical_json(evidence), output / EVIDENCE_MANIFEST_NAME) - return receipt + return GraphFinalization(files=tuple(names), receipt=receipt) def stage_graph_evidence( diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/rowwise_staging.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/rowwise_staging.py index 5833fa391..e9a3be5de 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/rowwise_staging.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/rowwise_staging.py @@ -108,6 +108,12 @@ _ACTIVE_EMITTER: LocalTelemetryEmitter | None = None +def active_telemetry_emitter() -> LocalTelemetryEmitter | None: + """The process-local emitter of the current UK build attempt, if started.""" + + return _ACTIVE_EMITTER + + def _hub_api() -> Any: """The Hub client used for the run bundle and staged dataset (test seam).""" diff --git a/packages/microcosm-build/tests/engine_free/shared/test_graph_publication.py b/packages/microcosm-build/tests/engine_free/shared/test_graph_publication.py index 8be249899..f3997d945 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_graph_publication.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_graph_publication.py @@ -297,3 +297,62 @@ def flush_once(self, *, publication_id): ) assert calls == [inventory["publication_id"]] assert snapshot.read_text() == '{"status":"pending"}' + + +def test_deleted_preserved_directory_is_retained_without_stopping_delivery( + tmp_path, monkeypatch +): + import shutil + + from microcosm.build.telemetry_emitter_service.graph_publication import time + + clock = [1000.0] + monkeypatch.setattr(time, "time", lambda: clock[0]) + queue = GraphPublicationQueue(tmp_path / "jobs.sqlite3") + (tmp_path / "removed").mkdir() + (tmp_path / "kept").mkdir() + removed, removed_inventory = source(tmp_path / "removed") + queue.enqueue(removed, removed_inventory) + clock[0] = 1001.0 + kept, kept_inventory = source(tmp_path / "kept") + queue.enqueue(kept, kept_inventory) + shutil.rmtree(removed) + delivery = GraphPublicationDelivery( + queue, credential=lambda: ("missing", None), origin="https://runs.example" + ) + + assert delivery.flush_once() is False + receipt = queue.receipt(removed_inventory["publication_id"]) + assert receipt["status"] == "pending" + assert receipt["error_code"] == "preserved_files_missing" + assert receipt["attempts"] == 1 + # The failed job backs off, so the next claim reaches the other job. + assert queue.claim().publication_id == kept_inventory["publication_id"] + + +def test_status_file_failure_does_not_roll_back_the_queue_result(tmp_path): + directory, inventory = source(tmp_path) + queue = GraphPublicationQueue(tmp_path / "jobs.sqlite3") + queue.enqueue(directory, inventory) + claimed = queue.claim() + (directory / "publication.status.json").mkdir() + queue.finish( + claimed.publication_id, + url="https://runs.example/runs/published/", + expected_lease=claimed.lease_until, + ) + receipt = queue.receipt(claimed.publication_id) + assert receipt["status"] == "published" + assert list(directory.glob("*.tmp")) == [] + + +@pytest.mark.parametrize( + "origin", + ["http://runs.example", "https://runs.example/path", "https://user@runs.example"], +) +def test_delivery_rejects_non_origin_destinations(tmp_path, origin): + queue = GraphPublicationQueue(tmp_path / "jobs.sqlite3") + with pytest.raises(ValueError, match="service URL"): + GraphPublicationDelivery( + queue, credential=lambda: ("missing", None), origin=origin + ) diff --git a/packages/microcosm-build/tests/engine_free/shared/test_publication_workers.py b/packages/microcosm-build/tests/engine_free/shared/test_publication_workers.py index 0289c8012..71bfe1566 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_publication_workers.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_publication_workers.py @@ -103,3 +103,40 @@ def exchange(*args): thread.join(2) assert results == [("ok", "synthetic-session")] * 2 assert len(calls) == 1 + + +def test_graph_worker_survives_a_local_queue_error(tmp_path, capsys): + from sqlalchemy.exc import OperationalError + + attempts = [] + recovered = threading.Event() + + class GraphDelivery: + def flush_once(self): + attempts.append(None) + if len(attempts) == 1: + raise OperationalError("UPDATE", {}, Exception("database is locked")) + recovered.set() + return False + + service = EmitterService( + socket_path=tmp_path / "unused.sock", + registration=None, + spool=EventSpool(tmp_path / "events.sqlite3"), + delivery=None, + graph_delivery=GraphDelivery(), + sampler=None, + heartbeat_seconds=60, + drain_seconds=0, + ) + worker = threading.Thread( + target=service._graph_worker, args=(service.graph_delivery,) + ) + worker.start() + try: + assert recovered.wait(10) + finally: + service._stop.set() + worker.join(5) + assert not worker.is_alive() + assert "graph publication hit a local queue error" in capsys.readouterr().err diff --git a/packages/microcosm-build/tests/engine_free/uk/test_uk_orrery_publication.py b/packages/microcosm-build/tests/engine_free/uk/test_uk_orrery_publication.py index 221afedf3..1e6b3c8da 100644 --- a/packages/microcosm-build/tests/engine_free/uk/test_uk_orrery_publication.py +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_orrery_publication.py @@ -52,8 +52,8 @@ def test_local_export_preserves_exact_graph_without_a_run_id_or_network( monkeypatch.setattr(rowwise_staging, "_ACTIVE_EMITTER", None) (args.out / "private.h5").write_bytes(b"licensed population") (args.out / "raw.json").write_text('{"person_id":123}') - record = {} - receipt = orrery_publication.finalize_graph(args, record) + result = orrery_publication.finalize_graph(args, manifest=None) + receipt = result.receipt assert receipt["status"] == "skipped" assert receipt["recorded_execution"] is False preserved = args.attempt_evidence / "orrery-publication" @@ -78,7 +78,7 @@ def test_publication_precedes_manifest_registration_and_receipt_is_a_snapshot( events = [] class Emitter: - def publish_graph(self, directory, inventory): + def publish_graph(self, directory, inventory, *, wait_seconds): assert not (args.out / "orrery.publication.json").exists() events.append("publication_attempt") return { @@ -89,8 +89,7 @@ def publish_graph(self, directory, inventory): monkeypatch.setattr(rowwise_staging, "_ACTIVE_EMITTER", Emitter()) manifest = {"outputs": {}} - record = {"manifest": manifest} - receipt = orrery_publication.finalize_graph(args, record) + receipt = orrery_publication.finalize_graph(args, manifest=manifest).receipt events.append("registered") assert events == ["publication_attempt", "registered"] assert receipt["status"] == "pending" @@ -108,24 +107,108 @@ def publish_graph(self, directory, inventory): assert (args.out / "orrery.publication.json").read_bytes() == snapshot +def fail(*args, **kwargs): + raise RuntimeError("kernel failed") + + def test_failure_path_preserves_graph_before_temporary_cleanup(tmp_path, monkeypatch): args = fixture(tmp_path) - calls = [] + staged = [] monkeypatch.setattr( - orrery_publication, "stage_graph_evidence", lambda *a, **k: calls.append("hf") + orrery_publication, + "stage_graph_evidence", + lambda args, output, **kwargs: staged.append(output), ) - def fail(*a, **k): - raise RuntimeError("kernel failed") + with pytest.raises(RuntimeError, match="kernel failed"): + full_build_cli._run_with_graph_publication( + fail, None, args, telemetry=None, attempt=None, record={} + ) + bundle = args.attempt_evidence / orrery_publication.FAILURE_BUNDLE_DIRECTORY + assert staged == [bundle] + assert (bundle / "graph.orrery.json").is_file() + assert (bundle / "orrery.evidence-manifest.json").is_file() + assert (bundle / "orrery.failure.json").is_file() + evidence = json.loads((bundle / "orrery.evidence-manifest.json").read_bytes()) + assert {entry["path"] for entry in evidence["outputs"].values()} == { + str(bundle / name) for name in evidence["outputs"] + } + + +def test_failure_leaves_an_earlier_successful_bundle_unchanged(tmp_path, monkeypatch): + args = fixture(tmp_path) + monkeypatch.setattr( + orrery_publication, "stage_graph_evidence", lambda *a, **k: None + ) + args.published_out.mkdir() + previous = { + "build.json": b'{"kind":"previous-complete-build"}', + "graph.orrery.json": b'{"previous":"graph"}', + "orrery.publication.json": b'{"previous":"receipt"}', + } + for name, payload in previous.items(): + (args.published_out / name).write_bytes(payload) with pytest.raises(RuntimeError, match="kernel failed"): full_build_cli._run_with_graph_publication( fail, None, args, telemetry=None, attempt=None, record={} ) - assert calls == ["hf"] - assert (args.published_out / "graph.orrery.json").is_file() - assert (args.published_out / "orrery.evidence-manifest.json").is_file() - assert (args.published_out / "orrery.failure.json").is_file() + assert { + path.name: path.read_bytes() for path in args.published_out.iterdir() + } == previous + + +def test_interrupt_preserves_local_evidence_without_waiting_or_staging( + tmp_path, monkeypatch +): + args = fixture(tmp_path) + args.publish_orrery = True + waits = [] + + class Emitter: + def publish_graph(self, directory, inventory, *, wait_seconds): + waits.append(wait_seconds) + return { + "publication_id": inventory["publication_id"], + "status": "pending", + "error_code": None, + } + + monkeypatch.setattr(rowwise_staging, "_ACTIVE_EMITTER", Emitter()) + monkeypatch.setattr( + orrery_publication, + "stage_graph_evidence", + lambda *a, **k: pytest.fail("An interrupt started Hugging Face staging"), + ) + + def interrupt(*args, **kwargs): + raise KeyboardInterrupt + + with pytest.raises(KeyboardInterrupt): + full_build_cli._run_with_graph_publication( + interrupt, None, args, telemetry=None, attempt=None, record={} + ) + assert waits == [0] + bundle = args.attempt_evidence / orrery_publication.FAILURE_BUNDLE_DIRECTORY + assert (bundle / "graph.orrery.json").is_file() + + +def test_failure_evidence_error_does_not_replace_the_build_error( + tmp_path, monkeypatch, capsys +): + args = fixture(tmp_path) + + def broken(*args, **kwargs): + raise OSError("disk full") + + monkeypatch.setattr(orrery_publication, "finalize_graph", broken) + with pytest.raises(RuntimeError, match="kernel failed"): + full_build_cli._run_with_graph_publication( + fail, None, args, telemetry=None, attempt=None, record={} + ) + assert "Graph failure evidence was not preserved (OSError)" in ( + capsys.readouterr().err + ) @pytest.mark.parametrize(