diff --git a/README.md b/README.md index f25d358cd6..820921a1f3 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-publication-failures.fixed.md b/changelog.d/1150-publication-failures.fixed.md new file mode 100644 index 0000000000..497aacf3b9 --- /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/changelog.d/1150-publication-types.changed.md b/changelog.d/1150-publication-types.changed.md new file mode 100644 index 0000000000..0947298231 --- /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/changelog.d/1150.added.md b/changelog.d/1150.added.md new file mode 100644 index 0000000000..4cec6ee6be --- /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 11c9328eca..d121fe37f0 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 0000000000..c7e6c65d96 --- /dev/null +++ b/docs/orrery-publication.md @@ -0,0 +1,70 @@ +# 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. +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 + +```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. 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/pyproject.toml b/packages/microcosm-build/pyproject.toml index c184758388..ed07ea12da 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/artifact_files.py b/packages/microcosm-build/src/microcosm/build/artifact_files.py index ff14fd7af4..263b30e777 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_cli.py b/packages/microcosm-build/src/microcosm/build/graph_publication_cli.py new file mode 100644 index 0000000000..89318d838a --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/graph_publication_cli.py @@ -0,0 +1,58 @@ +"""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 default_spool_path +from microcosm.build.telemetry_emitter_service.auth import ( + CollectorSession, + validated_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=default_spool_path()) + 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(validated_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) + 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 + + +if __name__ == "__main__": + raise SystemExit(main()) 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 0000000000..a0f03f6a8c --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/graph_publication_types.py @@ -0,0 +1,60 @@ +"""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", +] + + +type GraphJobStatus = Literal["pending", "published"] + + +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 0000000000..8b13789179 --- /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 37e7ae97ab..6a8cd23da4 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 c9fb5aaf7d..336b846fed 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, @@ -87,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.""" @@ -122,12 +131,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 @@ -177,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", @@ -215,6 +222,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 +254,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 +464,38 @@ def close(self) -> None: self._send({"action": ACTION_CLOSE}) self._closed = True + def publish_graph( + 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 + 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 default_spool_path()) + queue.enqueue(directory, inventory) + 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) + 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/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 0000000000..14ae0318a6 --- /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 0000000000..205fef3dd5 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/auth.py @@ -0,0 +1,180 @@ +"""HF-to-collector session shared by telemetry and graph publication.""" + +from __future__ import annotations + +import json +import os +import threading +import time +import urllib.error +import urllib.request +from collections.abc import Callable, Mapping +from http import HTTPStatus +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, + DEFAULT_TOKEN_LIFETIME_SECONDS, + DEVELOPMENT_COLLECTOR_LOOPBACK_ERROR, + HTTP_TIMEOUT_SECONDS, + HTTP_USER_AGENT, + 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, +) + + +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 + + +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(SERVICE_ORIGIN_HTTPS_ERROR) + if ( + parsed.username + or parsed.password + or parsed.query + or parsed.fragment + or parsed.path not in {"", "/"} + ): + raise ValueError(SERVICE_ORIGIN_FORMAT_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 validated_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) + + +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() + 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: 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 + 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) -> CollectorCredential: + with self._lock: + if ( + self._session_token is not None + and self._session_token[1] + > time.monotonic() + TOKEN_REFRESH_MARGIN_SECONDS + ): + return CollectorCredential("ok", self._session_token[0]) + hf_token = self._credential() + if not hf_token: + return CollectorCredential("missing", None) + try: + status, response = self._post( + self.origin + TOKEN_EXCHANGE_PATH, {}, hf_token + ) + except (OSError, TimeoutError): + return CollectorCredential("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 CollectorCredential("unavailable", None) + token = response["access_token"] + self._session_token = (token, time.monotonic() + expires_in) + return CollectorCredential("ok", token) + if status in {HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN}: + 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 0d7e4d294a..8bc19a3b73 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 ( + CollectorCredential, + CollectorSession, + _development_collector_url, + _http_post, + _huggingface_token, + validated_origin, +) 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,14 +57,19 @@ def __init__( spool: EventSpool, *, development_collector_url: str | None = None, + session: CollectorSession | None = None, ) -> None: 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_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) -> CollectorCredential: + """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/constants.py b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/constants.py index 83163fe53c..bd7559467b 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 new file mode 100644 index 0000000000..f1f95ed137 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/telemetry_emitter_service/graph_publication.py @@ -0,0 +1,414 @@ +"""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.graph_publication_types import ( + ClaimedGraphPublication, + PublicationFile, + PublicationFileRole, + 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, +) +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_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", + "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: PublicationInventory) -> str: + """Digest the versioned field order shared with the publication API.""" + ordered: PublicationInventory = { + "version": inventory["version"], + "publication_id": inventory["publication_id"], + "files": [ + { + "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(ordered, separators=(",", ":"), ensure_ascii=True).encode() + ) + + +def publication_inventory( + 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: 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 + ): + 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} + + +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.""" + + 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: PublicationInventory) -> 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, + ) -> 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) + 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, + ) + # 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 + claimed_id = session.scalar( + update(GraphPublicationJob) + .where( + GraphPublicationJob.publication_id == id_, + GraphPublicationJob.status == "pending", + GraphPublicationJob.lease_until <= now, + ) + .values(lease_until=now + lease_seconds) + .returning(GraphPublicationJob.publication_id) + ) + 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), + inventory=job.inventory, + lease_until=job.lease_until, + ) + + 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 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)) + receipt: PublicationReceipt = { + "version": 1, + "publication_id": id_, + "status": job.status, + "error_code": error_code, + "url": url, + "attempts": job.attempts, + "inventory_sha256": inventory_digest(job.inventory), + } + 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.""" + 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[[], CollectorCredential], + origin: str, + client: httpx.Client | None = None, + invalidate_credential: Callable[[], None] | None = None, + ) -> None: + 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 + + 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, + job.directory, + ) + 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: + 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 2872bf91a8..51d18fa742 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 c1c69c61c2..8446a5a2f8 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 87f1c678d4..610576a94f 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, @@ -15,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, ) @@ -128,3 +134,19 @@ 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[PublicationInventory] = mapped_column( + JsonObjectText, nullable=False + ) + directory: 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[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 0bb1d9c2a4..3a95be8157 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,10 +1,11 @@ -"""Unix-socket runtime for the telemetry emitter service.""" +"""Existing local service with independent telemetry and graph workers.""" from __future__ import annotations import json import os import socket +import sys import threading import time from collections.abc import Mapping @@ -25,6 +26,10 @@ UNSUPPORTED_ACTION_ERROR, WORKER_INTERVAL_SECONDS, ) +from microcosm.build.telemetry_emitter_service.graph_publication import ( + GRAPH_WORKER_ERROR_MESSAGE, + 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 +86,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 +111,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 +121,33 @@ 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, + args=(self.registration,), + daemon=True, + ) + ) + if self.graph_delivery is not None: + workers.append( + threading.Thread( + target=self._graph_worker, + args=(self.graph_delivery,), + 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 +156,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 +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 == ACTION_EVENT: + if action == ACTION_EVENT and self.registration is not None: event = message.get("event") if not isinstance(event, Mapping): raise ValueError(EVENT_OBJECT_ERROR) @@ -185,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() @@ -194,22 +233,28 @@ 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(), ) 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, 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): + 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 while self.spool.has_deliverable() and time.monotonic() < deadline: 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 c39f63901b..69cdfef272 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 @@ -619,6 +623,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 +1228,129 @@ def _release_stem(posture) -> tuple[str, str]: ) +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 finalize_graph + + try: + status = execute( + prepared, args, telemetry=telemetry, attempt=attempt, record=record + ) + 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: + _publish_failure_evidence( + args, record, telemetry=telemetry, interrupted=interrupted + ) + except Exception as error: + print( + f"Graph failure evidence was not preserved ({type(error).__name__}); " + "the build's own error follows.", + file=sys.stderr, + ) + raise + 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, @@ -1245,8 +1386,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 +1495,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 +2333,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 +2881,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 0000000000..81422e9dc7 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/orrery_publication.py @@ -0,0 +1,260 @@ +"""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 argparse import Namespace +from collections.abc import Iterator, MutableMapping +from dataclasses import dataclass +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 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 +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" +FAILURE_MARKER_NAME = "orrery.failure.json" +FAILURE_BUNDLE_DIRECTORY = "orrery-failure" + + +class GraphEvidenceOutput(TypedDict): + """A preserved public file registered in the build's output manifest.""" + + path: str + sha256: str + 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: + 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: object) -> Iterator[str]: + 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: Path, published_root: Path) -> GraphEvidenceOutput: + info = file_artifact(path) + return { + "path": str(published_root / path.name), + "sha256": info["sha256"], + "bytes": info["size_bytes"], + } + + +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. + 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) + if published_root is None: + published_root = Path(getattr(args, "published_out", output)) + schema_path = output / "graph.schema.json" + receipt: PublicationReceipt = { + "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: 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) + 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"]: + 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") + 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_telemetry_emitter() + if emitter is None: + # Durable enqueue is still possible without a running + # emitter; the retry command requires no telemetry run. + queue = GraphPublicationQueue(default_spool_path()) + queue.enqueue(preserved, inventory) + queued_receipt = queue.receipt(inventory["publication_id"]) + 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, 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. + receipt["status"] = "failed" + receipt["error_code"] = "graph_export_or_enqueue_failed" + materialize_bytes(canonical_json(receipt), output / RECEIPT_NAME) + names.append(RECEIPT_NAME) + 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 GraphFinalization(files=tuple(names), receipt=receipt) + + +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, + 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) + assert repository is not None, "Remote staging requires a repository." + return stage_bundle( + bundle, + storage=HuggingFaceDatasetStorage(repository, api=rowwise_staging._hub_api()), + prefix=UK_STAGED_DATASET_PREFIX, + ) 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 5833fa3911..e9a3be5de8 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/src/microcosm/build/us_runtime/worker_identity.py b/packages/microcosm-build/src/microcosm/build/us_runtime/worker_identity.py index 6183307342..2d5aff5b0b 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 0000000000..f3997d9451 --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/shared/test_graph_publication.py @@ -0,0 +1,358 @@ +"""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, + 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 isinstance(job, ClaimedGraphPublication) + assert job.inventory == inventory + assert "run_id" not in inventory + 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") + (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"}' + + +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 new file mode 100644 index 0000000000..71bfe1566b --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/shared/test_publication_workers.py @@ -0,0 +1,142 @@ +"""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 + + +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/shared/test_telemetry_emitter.py b/packages/microcosm-build/tests/engine_free/shared/test_telemetry_emitter.py index cee28f4349..a2cb501161 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/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 6aea249dbc..9fc790b549 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 0000000000..1e6b3c8da1 --- /dev/null +++ b/packages/microcosm-build/tests/engine_free/uk/test_uk_orrery_publication.py @@ -0,0 +1,230 @@ +"""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}') + 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" + 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, *, wait_seconds): + 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": {}} + receipt = orrery_publication.finalize_graph(args, manifest=manifest).receipt + 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 fail(*args, **kwargs): + raise RuntimeError("kernel failed") + + +def test_failure_path_preserves_graph_before_temporary_cleanup(tmp_path, monkeypatch): + args = fixture(tmp_path) + staged = [] + monkeypatch.setattr( + orrery_publication, + "stage_graph_evidence", + lambda args, output, **kwargs: staged.append(output), + ) + + 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 { + 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( + "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 654a25b695..af83888fd6 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 65e8d77974..7ff55a2475 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] diff --git a/uv.lock b/uv.lock index 9ea420ac3f..8da37c6d36 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" },