From 20c6ace7bd3ba5691161da0a9bcb8b9768c1cbe0 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:47:45 +0400 Subject: [PATCH 01/23] Add the durable Hugging Face hand-over queue and repository calls - Migration 20261020_03 adds four tables to the spool database: staged_run_files (the newest handed-over bytes per repository path), staged_run_deliveries (per-run counts and the pending runs-index entry), repository_commit_schedules (commit spacing, backoff and a per-repository lease) and handover_jobs (one row per hand-over and destination, with leases, retries and a receipt, like graph jobs). - `HandoverQueue` records hand-overs without touching the Hub: bytes equal to those last committed are not pending, unforced hand-overs wait for their interval after the repository's last commit, forced ones keep a 10-second spacing, and failures back off. - `DatasetRepository` ports the calls of Microcosm's HuggingFaceDatasetStorage that the emitter needs, including credential_can_write, and recognises parent-commit conflicts (HTTP 409 or 412). `hugging_face_api()` uses the build's credential in the same order as core/auth.py. - Tests use an in-memory Hugging Face API (tests/support). Co-Authored-By: Claude Opus 5.5 --- README.md | 2 +- .../huggingface-handover-queue.added.md | 1 + pyproject.toml | 2 + .../versions/v20261020_03_handover_queue.py | 68 +++ src/microcosm_emitter/core/models.py | 71 +++ src/microcosm_emitter/huggingface/__init__.py | 1 + .../huggingface/service/__init__.py | 1 + .../huggingface/service/hub.py | 210 +++++++ .../huggingface/service/queue.py | 561 ++++++++++++++++++ tests/huggingface/test_hub.py | 192 ++++++ tests/huggingface/test_queue.py | 281 +++++++++ tests/support/fake_hugging_face.py | 284 +++++++++ tests/telemetry/test_persistence.py | 6 +- tools/installed-smoke.py | 2 +- 14 files changed, 1679 insertions(+), 3 deletions(-) create mode 100644 changelog.d/huggingface-handover-queue.added.md create mode 100644 src/microcosm_emitter/core/alembic/versions/v20261020_03_handover_queue.py create mode 100644 src/microcosm_emitter/huggingface/__init__.py create mode 100644 src/microcosm_emitter/huggingface/service/__init__.py create mode 100644 src/microcosm_emitter/huggingface/service/hub.py create mode 100644 src/microcosm_emitter/huggingface/service/queue.py create mode 100644 tests/huggingface/test_hub.py create mode 100644 tests/huggingface/test_queue.py create mode 100644 tests/support/fake_hugging_face.py diff --git a/README.md b/README.md index cf20166..44d9b7a 100644 --- a/README.md +++ b/README.md @@ -42,7 +42,7 @@ The adapter bounds socket calls and reports telemetry failures without failing t The service uses the existing Hugging Face credential from the environment or local login cache. The collector validates PolicyEngine membership before issuing a short-lived write credential. Missing or rejected credentials permanently mark that producer's events as local-only. Permanent HTTP client errors when registering a run or uploading events also keep only the affected producer's records local; HTTP 408, 429, server errors and network failures remain retryable. HTTP 401 invalidates the short-lived credential for renewal. The production destination belongs to the package; development overrides accept loopback addresses only. -The queue upgrades the existing `20261007_01` Alembic revision to `20261008_02`, adding graph-job records without changing producer sequences, event IDs or upload eligibility. Unversioned databases are rejected without schema modification. There is no raw SQL, automatic schema adoption, or parallel legacy implementation. The hosted telemetry collector and calibration dashboard do not change. +The queue upgrades the existing `20261007_01` Alembic revision to `20261008_02`, adding graph-job records, and then to `20261020_03`, adding Hugging Face hand-over records, without changing producer sequences, event IDs or upload eligibility. Unversioned databases are rejected without schema modification. There is no raw SQL, automatic schema adoption, or parallel legacy implementation. The hosted telemetry collector and calibration dashboard do not change. ## Graph publication diff --git a/changelog.d/huggingface-handover-queue.added.md b/changelog.d/huggingface-handover-queue.added.md new file mode 100644 index 0000000..471245f --- /dev/null +++ b/changelog.d/huggingface-handover-queue.added.md @@ -0,0 +1 @@ +Add a durable Hugging Face hand-over queue to the emitter's spool database (Alembic revision `20261020_03`). Staging run files keep only the newest handed-over bytes per repository path; dataset bundles and releases are jobs with leases, retries and receipts; and a per-repository schedule spaces commits and lets one writer on a machine hold a repository at a time. Hugging Face calls use the build's own credential, found in the same order as telemetry's. diff --git a/pyproject.toml b/pyproject.toml index eddfeab..3f3b96b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,6 +43,8 @@ dev = ["pytest>=8,<9", "ruff>=0.11,<1", "build>=1,<2", "twine>=7,<8", "pyyaml>=6 [tool.pytest.ini_options] testpaths = ["tests"] +# Shared test doubles, such as the in-memory Hugging Face Hub. +pythonpath = ["tests/support"] addopts = "-q --import-mode=importlib" markers = ["collector: requires the pinned collector and disposable PostgreSQL"] diff --git a/src/microcosm_emitter/core/alembic/versions/v20261020_03_handover_queue.py b/src/microcosm_emitter/core/alembic/versions/v20261020_03_handover_queue.py new file mode 100644 index 0000000..6abfa3f --- /dev/null +++ b/src/microcosm_emitter/core/alembic/versions/v20261020_03_handover_queue.py @@ -0,0 +1,68 @@ +"""Add durable Hugging Face hand-overs without changing telemetry or graph jobs.""" + +import sqlalchemy as sa +from alembic import op + +revision = "20261020_03" +down_revision = "20261008_02" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.create_table( + "staged_run_files", + sa.Column("repository", sa.Text(), primary_key=True), + sa.Column("path_in_repo", sa.Text(), primary_key=True), + sa.Column("run_id", sa.Text(), nullable=False), + sa.Column("content", sa.LargeBinary(), nullable=False), + sa.Column("content_sha256", sa.Text(), nullable=False), + sa.Column("committed_sha256", sa.Text(), nullable=True), + sa.Column("pending", sa.Boolean(), nullable=False), + sa.Column("due_at", sa.Float(), nullable=False), + sa.Column("updated_at", sa.Float(), nullable=False), + ) + op.create_table( + "staged_run_deliveries", + sa.Column("repository", sa.Text(), primary_key=True), + sa.Column("run_id", sa.Text(), primary_key=True), + sa.Column("files_attempted", sa.Integer(), nullable=False), + sa.Column("files_committed", sa.Integer(), nullable=False), + sa.Column("commit_attempts", sa.Integer(), nullable=False), + sa.Column("commits", sa.Integer(), nullable=False), + sa.Column("last_error_code", sa.Text(), nullable=True), + sa.Column("last_commit", sa.Text(), nullable=True), + sa.Column("run_index_entry", sa.Text(), nullable=True), + sa.Column("run_index_pending", sa.Boolean(), nullable=False), + sa.Column("run_index_due_at", sa.Float(), nullable=False), + sa.Column("updated_at", sa.Float(), nullable=False), + ) + op.create_table( + "repository_commit_schedules", + sa.Column("repository", sa.Text(), primary_key=True), + sa.Column("last_commit_at", sa.Float(), nullable=False), + sa.Column("retry_after", sa.Float(), nullable=False), + sa.Column("consecutive_failures", sa.Integer(), nullable=False), + sa.Column("lease_until", sa.Float(), nullable=False), + ) + op.create_table( + "handover_jobs", + sa.Column("handover_id", sa.Text(), primary_key=True), + sa.Column("destination", sa.Text(), primary_key=True), + sa.Column("kind", sa.Text(), nullable=False), + sa.Column("target", sa.Text(), nullable=False), + sa.Column("request", 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), + sa.Column("created_at", sa.Float(), nullable=False), + ) + + +def downgrade() -> None: + op.drop_table("handover_jobs") + op.drop_table("repository_commit_schedules") + op.drop_table("staged_run_deliveries") + op.drop_table("staged_run_files") diff --git a/src/microcosm_emitter/core/models.py b/src/microcosm_emitter/core/models.py index 0bca2a9..d116285 100644 --- a/src/microcosm_emitter/core/models.py +++ b/src/microcosm_emitter/core/models.py @@ -6,10 +6,12 @@ from typing import Any from sqlalchemy import ( + Boolean, Float, ForeignKeyConstraint, Index, Integer, + LargeBinary, Text, UniqueConstraint, ) @@ -143,3 +145,72 @@ class GraphPublicationJob(SpoolModel): next_attempt_at: Mapped[float] = mapped_column(Float, nullable=False) lease_until: Mapped[float] = mapped_column(Float, nullable=False) receipt: Mapped[dict[str, Any]] = mapped_column(JsonObjectText, nullable=False) + + +class StagedRunFile(SpoolModel): + """The latest handed-over bytes for one repository path of a staging run. + + A later hand-over of the same path replaces pending bytes, so each commit + carries only the newest state. Committed rows keep the committed digest. + """ + + __tablename__ = "staged_run_files" + repository: Mapped[str] = mapped_column(Text, primary_key=True) + path_in_repo: Mapped[str] = mapped_column(Text, primary_key=True) + run_id: Mapped[str] = mapped_column(Text, nullable=False) + content: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + content_sha256: Mapped[str] = mapped_column(Text, nullable=False) + committed_sha256: Mapped[str | None] = mapped_column(Text) + pending: Mapped[bool] = mapped_column(Boolean, nullable=False) + due_at: Mapped[float] = mapped_column(Float, nullable=False) + updated_at: Mapped[float] = mapped_column(Float, nullable=False) + + +class StagedRunDelivery(SpoolModel): + """Delivery counts for one staging run, and its pending runs-index entry.""" + + __tablename__ = "staged_run_deliveries" + repository: Mapped[str] = mapped_column(Text, primary_key=True) + run_id: Mapped[str] = mapped_column(Text, primary_key=True) + files_attempted: Mapped[int] = mapped_column(Integer, nullable=False) + files_committed: Mapped[int] = mapped_column(Integer, nullable=False) + commit_attempts: Mapped[int] = mapped_column(Integer, nullable=False) + commits: Mapped[int] = mapped_column(Integer, nullable=False) + last_error_code: Mapped[str | None] = mapped_column(Text) + last_commit: Mapped[str | None] = mapped_column(Text) + run_index_entry: Mapped[dict[str, Any] | None] = mapped_column(JsonObjectText) + run_index_pending: Mapped[bool] = mapped_column(Boolean, nullable=False) + run_index_due_at: Mapped[float] = mapped_column(Float, nullable=False) + updated_at: Mapped[float] = mapped_column(Float, nullable=False) + + +class RepositoryCommitSchedule(SpoolModel): + """When a repository may next receive a commit from this machine.""" + + __tablename__ = "repository_commit_schedules" + repository: Mapped[str] = mapped_column(Text, primary_key=True) + last_commit_at: Mapped[float] = mapped_column(Float, nullable=False) + retry_after: Mapped[float] = mapped_column(Float, nullable=False) + consecutive_failures: Mapped[int] = mapped_column(Integer, nullable=False) + lease_until: Mapped[float] = mapped_column(Float, nullable=False) + + +class HandoverJob(SpoolModel): + """One durable hand-over for one destination, retried until it settles. + + The same hand-over can have a row per destination (today only + ``huggingface``), so a later destination needs no new table. + """ + + __tablename__ = "handover_jobs" + handover_id: Mapped[str] = mapped_column(Text, primary_key=True) + destination: Mapped[str] = mapped_column(Text, primary_key=True) + kind: Mapped[str] = mapped_column(Text, nullable=False) + target: Mapped[str] = mapped_column(Text, nullable=False) + request: Mapped[dict[str, Any]] = mapped_column(JsonObjectText, nullable=False) + status: Mapped[str] = mapped_column(Text, nullable=False) + attempts: Mapped[int] = mapped_column(Integer, nullable=False) + next_attempt_at: Mapped[float] = mapped_column(Float, nullable=False) + lease_until: Mapped[float] = mapped_column(Float, nullable=False) + receipt: Mapped[dict[str, Any]] = mapped_column(JsonObjectText, nullable=False) + created_at: Mapped[float] = mapped_column(Float, nullable=False) diff --git a/src/microcosm_emitter/huggingface/__init__.py b/src/microcosm_emitter/huggingface/__init__.py new file mode 100644 index 0000000..c3d3b9c --- /dev/null +++ b/src/microcosm_emitter/huggingface/__init__.py @@ -0,0 +1 @@ +"""Hugging Face hand-overs: staging run files, staged dataset bundles and releases.""" diff --git a/src/microcosm_emitter/huggingface/service/__init__.py b/src/microcosm_emitter/huggingface/service/__init__.py new file mode 100644 index 0000000..cdb589d --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/__init__.py @@ -0,0 +1 @@ +"""The Hugging Face hand-over module; constructed only by the service host.""" diff --git a/src/microcosm_emitter/huggingface/service/hub.py b/src/microcosm_emitter/huggingface/service/hub.py new file mode 100644 index 0000000..29382d5 --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/hub.py @@ -0,0 +1,210 @@ +"""Hugging Face dataset-repository calls, ported from Microcosm's staging storage. + +``DatasetRepository`` keeps the behaviour of Microcosm's +``HuggingFaceDatasetStorage`` (``staging_storage.py``) for the calls the +emitter makes, with an injected ``HfApi``-shaped client so tests use an +in-memory Hub. The client carries the build's own credential, found exactly as +``microcosm_emitter.core.auth`` finds it. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + +from microcosm_emitter.core.auth import _huggingface_token + +#: HTTP statuses the Hub returns when ``parent_commit`` is no longer the head. +PARENT_CONFLICT_STATUSES = frozenset({409, 412}) +_ABSENT_ERROR_NAMES = frozenset( + { + "EntryNotFoundError", + "RemoteEntryNotFoundError", + "RevisionNotFoundError", + "FileNotFoundError", + } +) + + +def hugging_face_api() -> Any | None: + """An ``HfApi`` with the build's credential, or ``None`` without one. + + The credential is ``HF_TOKEN``, then ``HUGGINGFACE_TOKEN``, then the token + saved by ``hf auth login``: the same order telemetry uses. It is read again + for every cycle, so a login during a long build is picked up. + """ + token = _huggingface_token() + if not token: + return None + from huggingface_hub import HfApi + + return HfApi(token=token) + + +def add_operation(path_in_repo: str, content: bytes | Path) -> Any: + """One file write in a commit, from bytes or a preserved local file.""" + from huggingface_hub import CommitOperationAdd + + return CommitOperationAdd( + path_in_repo=path_in_repo, + path_or_fileobj=content if isinstance(content, bytes) else str(content), + ) + + +def status_code(error: BaseException) -> int | None: + response = getattr(error, "response", None) + code = getattr(response, "status_code", None) + return code if isinstance(code, int) else None + + +def is_parent_conflict(error: BaseException) -> bool: + """Another commit landed after the head this commit was built on.""" + return status_code(error) in PARENT_CONFLICT_STATUSES + + +def is_absent(error: BaseException) -> bool: + """The Hub answered that the file or revision does not exist, nothing else.""" + return bool({cls.__name__ for cls in type(error).__mro__} & _ABSENT_ERROR_NAMES) + + +class DatasetRepository: + """Read and write files in one Hugging Face dataset repository.""" + + def __init__(self, repo_id: str, api: Any) -> None: + if not isinstance(repo_id, str) or not repo_id.strip(): + raise ValueError("A Hugging Face hand-over requires a repository.") + self.repo_id = repo_id.strip() + self.api = api + + def head_revision(self) -> str: + info = self.api.repo_info( + repo_id=self.repo_id, repo_type="dataset", revision="main" + ) + sha = ( + info.get("sha") if isinstance(info, Mapping) else getattr(info, "sha", None) + ) + if not sha: + raise RuntimeError("The Hub reported no head revision.") + return str(sha) + + def commit( + self, + operations: Sequence[Any], + *, + message: str, + parent_commit: str | None = None, + ) -> str: + """Write several files in one commit and return the new revision. + + Pinning ``parent_commit`` makes a concurrent write fail with a + conflict (see ``is_parent_conflict``) instead of interleaving. + """ + info = self.api.create_commit( + repo_id=self.repo_id, + operations=list(operations), + commit_message=message, + repo_type="dataset", + parent_commit=parent_commit, + ) + oid = ( + info.get("oid") if isinstance(info, Mapping) else getattr(info, "oid", None) + ) + if not isinstance(oid, str) or not oid.strip(): + raise RuntimeError("The Hub commit reported no revision.") + return oid.strip() + + def download(self, path_in_repo: str, *, revision: str | None = None) -> bytes: + local = self.api.hf_hub_download( + repo_id=self.repo_id, + filename=path_in_repo, + repo_type="dataset", + revision=revision, + force_download=revision is None, + ) + return Path(local).read_bytes() + + def file_exists(self, path_in_repo: str, *, revision: str | None = None) -> bool: + return bool( + self.api.file_exists( + repo_id=self.repo_id, + filename=path_in_repo, + repo_type="dataset", + revision=revision, + ) + ) + + def list_files(self, *, revision: str) -> list[str]: + return sorted( + self.api.list_repo_files( + repo_id=self.repo_id, repo_type="dataset", revision=revision + ) + ) + + def last_commit(self, path_in_repo: str) -> str | None: + """The commit that last changed ``path_in_repo``, when the Hub reports it.""" + entries = self.api.get_paths_info( + repo_id=self.repo_id, + paths=[path_in_repo], + expand=True, + repo_type="dataset", + ) + for entry in entries or []: + commit = ( + entry.get("last_commit") + if isinstance(entry, Mapping) + else getattr(entry, "last_commit", None) + ) + oid = ( + commit.get("oid") + if isinstance(commit, Mapping) + else getattr(commit, "oid", None) + ) + if oid: + return str(oid) + return None + + def _access_token(self) -> Mapping[str, Any] | None: + info = self.api.whoami() + auth = info.get("auth") if isinstance(info, Mapping) else None + token = auth.get("accessToken") if isinstance(auth, Mapping) else None + return token if isinstance(token, Mapping) else None + + def credential_role(self) -> str | None: + """``read``, ``write`` or ``fineGrained``, as the Hub reports the token.""" + token = self._access_token() + role = token.get("role") if token else None + return str(role) if role else None + + def credential_can_write(self) -> bool | None: + """Whether the credential may write this repository; ``None`` if unknown. + + ``True`` for a classic write token or a fine-grained token whose scopes + grant ``repo.write`` on the repository, its owner or globally. + """ + token = self._access_token() + if token is None: + return None + role = str(token.get("role") or "") + if role == "read": + return False + if role == "write": + return True + if role != "fineGrained": + return None + fine = token.get("fineGrained") + fine = fine if isinstance(fine, Mapping) else {} + if "repo.write" in (fine.get("global") or []): + return True + owner = self.repo_id.split("/", 1)[0] + for scope in fine.get("scoped") or []: + if not isinstance(scope, Mapping): + continue + entity = scope.get("entity") + entity = entity if isinstance(entity, Mapping) else {} + name = str(entity.get("name") or "") + kind = str(entity.get("type") or "") + covers = name == self.repo_id or (kind in {"user", "org"} and name == owner) + if covers and "repo.write" in (scope.get("permissions") or []): + return True + return False diff --git a/src/microcosm_emitter/huggingface/service/queue.py b/src/microcosm_emitter/huggingface/service/queue.py new file mode 100644 index 0000000..ca1cd9f --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/queue.py @@ -0,0 +1,561 @@ +"""Durable hand-overs, kept in the emitter's spool database. + +Two record shapes share one SQLite database (and its Alembic history) with +telemetry and graph jobs: + +- **Staging run files** coalesce: each repository path keeps only the newest + handed-over bytes, so a commit carries the latest state of every pending + path at once. One commit per repository per cycle replaces Microcosm's one + ``upload_file`` commit per file. +- **Jobs** (dataset bundles, and later releases) follow the graph queue's + pattern: a row per hand-over and destination, a lease while a worker holds + it, retries with backoff, and a receipt. + +A per-repository schedule spaces commits and leases the repository to one +writer on this machine at a time. +""" + +from __future__ import annotations + +import hashlib +import time +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from sqlalchemy import select, update +from sqlalchemy.orm import Session + +from microcosm_emitter.core.database import ( + create_spool_engine, + create_spool_session_factory, +) +from microcosm_emitter.core.migrations import upgrade_spool_database +from microcosm_emitter.core.models import ( + HandoverJob, + RepositoryCommitSchedule, + StagedRunDelivery, + StagedRunFile, +) + +#: The shortest time between two commits to one repository from this machine, +#: even for hand-overs that ask to be committed at the next cycle. +MIN_COMMIT_SPACING_SECONDS = 10.0 +MAX_RETRY_SECONDS = 300.0 +DESTINATION_HUGGING_FACE = "huggingface" +JOB_PENDING = "pending" + + +def sha256_bytes(content: bytes) -> str: + return hashlib.sha256(content).hexdigest() + + +def retry_delay(failures: int) -> float: + return min(MAX_RETRY_SECONDS, float(2 ** min(failures, 8))) + + +@dataclass(frozen=True) +class PendingRunFile: + """The bytes a commit will write at one repository path.""" + + path_in_repo: str + run_id: str + content: bytes + sha256: str + + +@dataclass(frozen=True) +class ClaimedJob: + """A job leased until ``lease_until``; the value also identifies the lease.""" + + handover_id: str + destination: str + kind: str + target: str + request: dict[str, Any] + attempts: int + lease_until: float + + +class HandoverQueue: + """Hand-over records, retained until delivered or settled.""" + + def __init__(self, path: Path | str) -> 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 close(self) -> None: + """Release this handle's database connections without deleting records.""" + self._engine.dispose() + + # Repository schedules ------------------------------------------------ + + @staticmethod + def _schedule(session: Session, repository: str) -> RepositoryCommitSchedule: + schedule = session.get(RepositoryCommitSchedule, repository) + if schedule is None: + schedule = RepositoryCommitSchedule( + repository=repository, + last_commit_at=0.0, + retry_after=0.0, + consecutive_failures=0, + lease_until=0.0, + ) + session.add(schedule) + session.flush() + return schedule + + def claim_repository( + self, repository: str, *, now: float, lease_seconds: float + ) -> float | None: + """Lease the repository to this writer; ``None`` while another holds it.""" + with self._sessions.begin() as session: + self._schedule(session, repository) + lease = now + lease_seconds + claimed = session.scalar( + update(RepositoryCommitSchedule) + .where( + RepositoryCommitSchedule.repository == repository, + RepositoryCommitSchedule.lease_until <= now, + ) + .values(lease_until=lease) + .returning(RepositoryCommitSchedule.repository) + ) + return lease if claimed is not None else None + + def release_repository(self, repository: str, lease: float) -> None: + with self._sessions.begin() as session: + session.execute( + update(RepositoryCommitSchedule) + .where( + RepositoryCommitSchedule.repository == repository, + RepositoryCommitSchedule.lease_until == lease, + ) + .values(lease_until=0.0) + ) + + def record_commit(self, repository: str, *, now: float) -> None: + with self._sessions.begin() as session: + self._schedule(session, repository).last_commit_at = now + + def due_repositories(self, now: float) -> list[str]: + """Repositories with work due now that may receive a commit now.""" + with self._sessions() as session: + run_file_repositories = set( + session.scalars( + select(StagedRunFile.repository).where( + StagedRunFile.pending.is_(True), StagedRunFile.due_at <= now + ) + ) + ) | set( + session.scalars( + select(StagedRunDelivery.repository).where( + StagedRunDelivery.run_index_pending.is_(True), + StagedRunDelivery.run_index_due_at <= now, + ) + ) + ) + job_repositories = set( + session.scalars( + select(HandoverJob.target).where( + HandoverJob.destination == DESTINATION_HUGGING_FACE, + HandoverJob.status == JOB_PENDING, + HandoverJob.next_attempt_at <= now, + HandoverJob.lease_until <= now, + ) + ) + ) + due = [] + for repository in sorted(run_file_repositories | job_repositories): + schedule = session.get(RepositoryCommitSchedule, repository) + if schedule is not None and ( + schedule.lease_until > now + or schedule.last_commit_at + MIN_COMMIT_SPACING_SECONDS > now + ): + continue + run_files_allowed = schedule is None or schedule.retry_after <= now + if repository in job_repositories or ( + repository in run_file_repositories and run_files_allowed + ): + due.append(repository) + return due + + # Staging run files --------------------------------------------------- + + @staticmethod + def _delivery( + session: Session, repository: str, run_id: str, now: float + ) -> StagedRunDelivery: + delivery = session.get(StagedRunDelivery, (repository, run_id)) + if delivery is None: + delivery = StagedRunDelivery( + repository=repository, + run_id=run_id, + files_attempted=0, + files_committed=0, + commit_attempts=0, + commits=0, + last_error_code=None, + last_commit=None, + run_index_entry=None, + run_index_pending=False, + run_index_due_at=0.0, + updated_at=now, + ) + session.add(delivery) + session.flush() + return delivery + + def stage_run_files( + self, + *, + repository: str, + run_id: str, + files: Sequence[tuple[str, bytes]], + run_index_entry: Mapping[str, Any] | None, + flush: bool, + interval: float, + now: float, + ) -> None: + """Record the newest bytes per path; nothing reaches the Hub here. + + An unforced hand-over is due one ``interval`` after the repository's + last commit; a forced one (``flush``) at the next cycle. + """ + with self._sessions.begin() as session: + schedule = self._schedule(session, repository) + due_at = now if flush else max(now, schedule.last_commit_at + interval) + for path_in_repo, content in files: + digest = sha256_bytes(content) + row = session.get(StagedRunFile, (repository, path_in_repo)) + if row is None: + session.add( + StagedRunFile( + repository=repository, + path_in_repo=path_in_repo, + run_id=run_id, + content=content, + content_sha256=digest, + committed_sha256=None, + pending=True, + due_at=due_at, + updated_at=now, + ) + ) + continue + if digest == row.committed_sha256: + # The repository already holds exactly these bytes. + row.pending, row.content = False, b"" + elif not (row.pending and digest == row.content_sha256): + row.due_at = min(row.due_at, due_at) if row.pending else due_at + row.content, row.pending = content, True + row.content_sha256 = digest + row.run_id, row.updated_at = run_id, now + delivery = self._delivery(session, repository, run_id, now) + if run_index_entry is not None: + entry = dict(run_index_entry) + if delivery.run_index_entry != entry or delivery.run_index_pending: + delivery.run_index_due_at = ( + min(delivery.run_index_due_at, due_at) + if delivery.run_index_pending + else due_at + ) + delivery.run_index_entry = entry + delivery.run_index_pending = True + if flush: + session.execute( + update(StagedRunFile) + .where( + StagedRunFile.repository == repository, + StagedRunFile.run_id == run_id, + StagedRunFile.pending.is_(True), + StagedRunFile.due_at > now, + ) + .values(due_at=now) + ) + if delivery.run_index_pending: + delivery.run_index_due_at = min(delivery.run_index_due_at, now) + delivery.updated_at = now + + def pending_run_files(self, repository: str) -> list[PendingRunFile]: + with self._sessions() as session: + rows = session.scalars( + select(StagedRunFile) + .where( + StagedRunFile.repository == repository, + StagedRunFile.pending.is_(True), + ) + .order_by(StagedRunFile.path_in_repo) + ) + return [ + PendingRunFile( + row.path_in_repo, row.run_id, row.content, row.content_sha256 + ) + for row in rows + ] + + def pending_run_index_entries( + self, repository: str + ) -> list[tuple[str, dict[str, Any]]]: + with self._sessions() as session: + rows = session.scalars( + select(StagedRunDelivery) + .where( + StagedRunDelivery.repository == repository, + StagedRunDelivery.run_index_pending.is_(True), + ) + .order_by(StagedRunDelivery.run_id) + ) + return [(row.run_id, dict(row.run_index_entry or {})) for row in rows] + + @staticmethod + def _files_per_run(entries: Iterable[PendingRunFile]) -> dict[str, int]: + counts: dict[str, int] = {} + for entry in entries: + counts[entry.run_id] = counts.get(entry.run_id, 0) + 1 + return counts + + def record_run_files_committed( + self, + repository: str, + entries: Sequence[PendingRunFile], + index_entries: Sequence[tuple[str, Mapping[str, Any]]], + *, + revision: str, + now: float, + ) -> None: + """Settle what the commit wrote; newer hand-overs stay pending.""" + with self._sessions.begin() as session: + for entry in entries: + row = session.get(StagedRunFile, (repository, entry.path_in_repo)) + if row is None: + continue + row.committed_sha256 = entry.sha256 + if row.content_sha256 == entry.sha256: + row.pending, row.content = False, b"" + counts = self._files_per_run(entries) + index_runs = {run_id for run_id, _ in index_entries} + for run_id in sorted(set(counts) | index_runs): + delivery = self._delivery(session, repository, run_id, now) + delivery.files_attempted += counts.get(run_id, 0) + delivery.files_committed += counts.get(run_id, 0) + delivery.commit_attempts += 1 + delivery.commits += 1 + delivery.last_commit = revision + delivery.last_error_code = None + delivery.updated_at = now + for run_id, entry in index_entries: + delivery = self._delivery(session, repository, run_id, now) + if delivery.run_index_entry == dict(entry): + delivery.run_index_pending = False + schedule = self._schedule(session, repository) + schedule.last_commit_at = now + schedule.consecutive_failures = 0 + schedule.retry_after = 0.0 + + def record_run_files_failed( + self, + repository: str, + entries: Sequence[PendingRunFile], + index_entries: Sequence[tuple[str, Mapping[str, Any]]], + *, + error_code: str, + now: float, + ) -> None: + """Keep everything pending and back off before the next attempt.""" + with self._sessions.begin() as session: + counts = self._files_per_run(entries) + index_runs = {run_id for run_id, _ in index_entries} + for run_id in sorted(set(counts) | index_runs): + delivery = self._delivery(session, repository, run_id, now) + delivery.files_attempted += counts.get(run_id, 0) + delivery.commit_attempts += 1 + delivery.last_error_code = error_code + delivery.updated_at = now + schedule = self._schedule(session, repository) + schedule.consecutive_failures += 1 + schedule.retry_after = now + retry_delay(schedule.consecutive_failures) + + def run_receipt(self, repository: str, run_id: str) -> dict[str, Any]: + with self._sessions() as session: + delivery = session.get(StagedRunDelivery, (repository, run_id)) + pending = len( + list( + session.scalars( + select(StagedRunFile.path_in_repo).where( + StagedRunFile.repository == repository, + StagedRunFile.run_id == run_id, + StagedRunFile.pending.is_(True), + ) + ) + ) + ) + if delivery is not None and delivery.run_index_pending: + pending += 1 + return { + "repository": repository, + "run_id": run_id, + "pending_files": pending, + "files_attempted": delivery.files_attempted if delivery else 0, + "files_committed": delivery.files_committed if delivery else 0, + "commit_attempts": delivery.commit_attempts if delivery else 0, + "commits": delivery.commits if delivery else 0, + "last_error_code": delivery.last_error_code if delivery else None, + "last_commit": delivery.last_commit if delivery else None, + } + + def committed_run_files(self, repository: str, run_id: str) -> dict[str, str]: + """Each committed path of the run and the digest of its committed bytes.""" + with self._sessions() as session: + rows = session.scalars( + select(StagedRunFile).where( + StagedRunFile.repository == repository, + StagedRunFile.run_id == run_id, + StagedRunFile.committed_sha256.is_not(None), + ) + ) + return {row.path_in_repo: str(row.committed_sha256) for row in rows} + + # Jobs ---------------------------------------------------------------- + + def enqueue_job( + self, + *, + handover_id: str, + kind: str, + target: str, + request: Mapping[str, Any], + receipt: Mapping[str, Any], + now: float, + destination: str = DESTINATION_HUGGING_FACE, + ) -> dict[str, Any]: + """Durably record a hand-over once; repeating it returns its receipt.""" + with self._sessions.begin() as session: + existing = session.get(HandoverJob, (handover_id, destination)) + if existing is not None: + if ( + existing.request != dict(request) + or existing.kind != kind + or existing.target != target + ): + raise ValueError( + "Hand-over identity conflicts with a queued hand-over." + ) + return dict(existing.receipt) + session.add( + HandoverJob( + handover_id=handover_id, + destination=destination, + kind=kind, + target=target, + request=dict(request), + status=JOB_PENDING, + attempts=0, + next_attempt_at=0.0, + lease_until=0.0, + receipt=dict(receipt), + created_at=now, + ) + ) + return dict(receipt) + + def claim_job( + self, + target: str, + *, + now: float, + lease_seconds: float, + destination: str = DESTINATION_HUGGING_FACE, + ) -> ClaimedJob | None: + """Lease the earliest-due pending job for ``target``. + + A conditional update allows one current lease; a worker that exits + leaves its job eligible again once the lease expires. + """ + with self._sessions.begin() as session: + candidate = session.scalar( + select(HandoverJob) + .where( + HandoverJob.destination == destination, + HandoverJob.target == target, + HandoverJob.status == JOB_PENDING, + HandoverJob.next_attempt_at <= now, + HandoverJob.lease_until <= now, + ) + .order_by(HandoverJob.next_attempt_at, HandoverJob.created_at) + .limit(1) + ) + if candidate is None: + return None + lease = now + lease_seconds + claimed = session.scalar( + update(HandoverJob) + .where( + HandoverJob.handover_id == candidate.handover_id, + HandoverJob.destination == destination, + HandoverJob.status == JOB_PENDING, + HandoverJob.lease_until <= now, + ) + .values(lease_until=lease) + .returning(HandoverJob.handover_id) + ) + if claimed is None: + return None + session.refresh(candidate) + return ClaimedJob( + handover_id=candidate.handover_id, + destination=destination, + kind=candidate.kind, + target=candidate.target, + request=dict(candidate.request), + attempts=candidate.attempts, + lease_until=lease, + ) + + def finish_job( + self, + job: ClaimedJob, + *, + receipt: Mapping[str, Any], + retry: bool, + now: float, + committed: bool = False, + ) -> None: + """Record one attempt; only the current lease holder may do so.""" + with self._sessions.begin() as session: + row = session.scalar( + update(HandoverJob) + .where( + HandoverJob.handover_id == job.handover_id, + HandoverJob.destination == job.destination, + HandoverJob.status == JOB_PENDING, + HandoverJob.lease_until == job.lease_until, + ) + .values(attempts=HandoverJob.attempts + 1, lease_until=0.0) + .returning(HandoverJob) + ) + if row is None: + return + row.receipt = dict(receipt) + if retry: + row.next_attempt_at = now + retry_delay(row.attempts) + else: + row.status = str(receipt["status"]) + if committed: + self._schedule(session, job.target).last_commit_at = now + + def job_receipt( + self, handover_id: str, destination: str = DESTINATION_HUGGING_FACE + ) -> dict[str, Any] | None: + with self._sessions() as session: + job = session.get(HandoverJob, (handover_id, destination)) + return None if job is None else dict(job.receipt) + + +def now_seconds() -> float: + """Wall-clock seconds, shared by every persisted queue time.""" + return time.time() diff --git a/tests/huggingface/test_hub.py b/tests/huggingface/test_hub.py new file mode 100644 index 0000000..a1e7d0f --- /dev/null +++ b/tests/huggingface/test_hub.py @@ -0,0 +1,192 @@ +"""The emitter's Hugging Face repository calls, ported from Microcosm's storage.""" + +import pytest +from fake_hugging_face import ( + EntryNotFoundError, + FakeHubHTTPError, + FakeHuggingFaceApi, + RevisionNotFoundError, +) + +from microcosm_emitter.huggingface.service.hub import ( + DatasetRepository, + add_operation, + hugging_face_api, + is_absent, + is_parent_conflict, +) + +REPOSITORY = "policyengine/fake-staging" + + +@pytest.fixture +def hub(): + api = FakeHuggingFaceApi() + api.create_repository(REPOSITORY, {"README.md": b"fixture"}) + return api + + +def test_commit_writes_every_operation_in_one_commit_and_returns_its_revision( + hub, tmp_path +): + local = tmp_path / "progress.json" + local.write_bytes(b'{"stage":"created"}') + repository = DatasetRepository(REPOSITORY, hub) + revision = repository.commit( + [ + add_operation("runs/run-1/progress.json", local), + add_operation("runs/run-1/events.ndjson", b"{}\n"), + ], + message="Stage run files", + ) + (commit,) = hub.commits(REPOSITORY) + assert commit.oid == revision == repository.head_revision() + assert commit.added == { + "runs/run-1/progress.json": b'{"stage":"created"}', + "runs/run-1/events.ndjson": b"{}\n", + } + assert commit.message == "Stage run files" + + +def test_commit_on_a_stale_parent_is_a_conflict(hub): + repository = DatasetRepository(REPOSITORY, hub) + stale = repository.head_revision() + hub.write_directly(REPOSITORY, {"other.json": b"{}"}) + with pytest.raises(FakeHubHTTPError) as conflict: + repository.commit( + [add_operation("a.json", b"{}")], message="m", parent_commit=stale + ) + assert is_parent_conflict(conflict.value) + assert not is_parent_conflict(FakeHubHTTPError(500)) + assert not is_parent_conflict(ValueError("other")) + + +def test_downloads_read_exact_bytes_at_a_revision(hub): + repository = DatasetRepository(REPOSITORY, hub) + first = repository.head_revision() + hub.write_directly(REPOSITORY, {"README.md": b"changed"}) + assert repository.download("README.md") == b"changed" + assert repository.download("README.md", revision=first) == b"fixture" + assert repository.file_exists("README.md") + assert not repository.file_exists("missing.json") + assert repository.list_files(revision=first) == ["README.md"] + with pytest.raises(EntryNotFoundError) as missing: + repository.download("missing.json") + assert is_absent(missing.value) + assert is_absent(RevisionNotFoundError()) + assert not is_absent(FakeHubHTTPError(503)) + + +def test_last_commit_names_the_commit_that_last_changed_a_path(hub): + repository = DatasetRepository(REPOSITORY, hub) + changed = hub.write_directly(REPOSITORY, {"a.json": b"1"}) + hub.write_directly(REPOSITORY, {"b.json": b"2"}) + assert repository.last_commit("a.json") == changed + assert repository.last_commit("never.json") is None + + +@pytest.mark.parametrize( + ("access_token", "can_write"), + [ + ({"role": "read"}, False), + ({"role": "write"}, True), + ({"role": "fineGrained", "fineGrained": {"global": ["repo.write"]}}, True), + ( + { + "role": "fineGrained", + "fineGrained": { + "scoped": [ + { + "entity": {"name": REPOSITORY, "type": "dataset"}, + "permissions": ["repo.write"], + } + ] + }, + }, + True, + ), + ( + { + "role": "fineGrained", + "fineGrained": { + "scoped": [ + { + "entity": {"name": "policyengine", "type": "org"}, + "permissions": ["repo.write"], + } + ] + }, + }, + True, + ), + ( + { + "role": "fineGrained", + "fineGrained": { + "scoped": [ + { + "entity": {"name": "someone-else", "type": "org"}, + "permissions": ["repo.write"], + } + ] + }, + }, + False, + ), + ({"role": "something-new"}, None), + ], + ids=[ + "read", + "write", + "fine-grained-global", + "fine-grained-repository", + "fine-grained-owner", + "fine-grained-elsewhere", + "unknown-role", + ], +) +def test_write_scope_is_read_from_the_hub_description_of_the_token( + access_token, can_write +): + api = FakeHuggingFaceApi(whoami={"auth": {"accessToken": access_token}}) + repository = DatasetRepository(REPOSITORY, api) + assert repository.credential_can_write() is can_write + assert repository.credential_role() == access_token["role"] + + +def test_write_scope_is_unknown_when_the_hub_does_not_describe_the_token(): + api = FakeHuggingFaceApi(whoami={"name": "fixture"}) + assert DatasetRepository(REPOSITORY, api).credential_can_write() is None + + +@pytest.mark.parametrize( + ("variables", "saved_login", "expected_token"), + [ + ( + {"HF_TOKEN": "hf_fake_first", "HUGGINGFACE_TOKEN": "hf_fake_second"}, + "hf_fake_saved", + "hf_fake_first", + ), + ({"HUGGINGFACE_TOKEN": "hf_fake_second"}, "hf_fake_saved", "hf_fake_second"), + ({}, "hf_fake_saved", "hf_fake_saved"), + ({}, None, None), + ], + ids=["hf-token", "huggingface-token", "saved-login", "none"], +) +def test_uploads_use_the_build_credential_found_in_the_core_auth_order( + monkeypatch, variables, saved_login, expected_token +): + import os + + for name, value in variables.items(): + monkeypatch.setenv(name, value) + if saved_login is not None: + token_path = os.environ["HF_TOKEN_PATH"] + os.makedirs(os.path.dirname(token_path), exist_ok=True) + with open(token_path, "w") as token_file: + token_file.write(saved_login) + api = hugging_face_api() + if expected_token is None: + assert api is None + else: + assert api.token == expected_token diff --git a/tests/huggingface/test_queue.py b/tests/huggingface/test_queue.py new file mode 100644 index 0000000..bcb6801 --- /dev/null +++ b/tests/huggingface/test_queue.py @@ -0,0 +1,281 @@ +"""The durable hand-over queue shares the spool database and its migrations.""" + +import shutil +from pathlib import Path + +import pytest +from sqlalchemy import func, select +from sqlalchemy.orm import Session + +from microcosm_emitter.core.database import create_spool_engine +from microcosm_emitter.core.migrations import current_database_revision +from microcosm_emitter.core.models import TelemetryEventRecord +from microcosm_emitter.huggingface.service.queue import ( + MIN_COMMIT_SPACING_SECONDS, + HandoverQueue, +) + +FIXTURES = Path(__file__).parents[1] / "fixtures" +REPOSITORY = "policyengine/fake-staging" +OTHER_REPOSITORY = "policyengine/fake-private" + + +@pytest.fixture +def queue(tmp_path): + handle = HandoverQueue(tmp_path / "events.sqlite3") + yield handle + handle.close() + + +def stage(queue, path_in_repo, content, *, now, flush=True, interval=30.0, **extra): + queue.stage_run_files( + repository=extra.pop("repository", REPOSITORY), + run_id=extra.pop("run_id", "run-1"), + files=[(path_in_repo, content)], + run_index_entry=extra.pop("run_index_entry", None), + flush=flush, + interval=interval, + now=now, + ) + + +def commit_everything(queue, *, repository=REPOSITORY, now, revision="commit-1"): + queue.record_run_files_committed( + repository, + queue.pending_run_files(repository), + queue.pending_run_index_entries(repository), + revision=revision, + now=now, + ) + + +def test_migration_adds_hand_over_tables_without_changing_telemetry(tmp_path): + path = tmp_path / "events.sqlite3" + shutil.copyfile(FIXTURES / "microcosm-1099.sqlite3", path) + engine = create_spool_engine(path) + with Session(engine) as session: + events_before = session.scalar( + select(func.count(TelemetryEventRecord.event_id)) + ) + engine.dispose() + HandoverQueue(path).close() + assert current_database_revision(path) == "20261020_03" + engine = create_spool_engine(path) + with Session(engine) as session: + assert ( + session.scalar(select(func.count(TelemetryEventRecord.event_id))) + == events_before + ) + engine.dispose() + + +def test_latest_bytes_for_a_path_replace_earlier_pending_bytes(queue): + stage(queue, "runs/run-1/progress.json", b'{"stage":"created"}', now=100) + stage(queue, "runs/run-1/progress.json", b'{"stage":"calibrating"}', now=101) + (pending,) = queue.pending_run_files(REPOSITORY) + assert pending.path_in_repo == "runs/run-1/progress.json" + assert pending.content == b'{"stage":"calibrating"}' + assert pending.run_id == "run-1" + + +def test_bytes_equal_to_those_last_committed_are_not_pending(queue): + stage(queue, "runs/run-1/run_manifest.json", b"{}", now=100) + commit_everything(queue, now=100) + stage(queue, "runs/run-1/run_manifest.json", b"{}", now=200) + assert queue.pending_run_files(REPOSITORY) == [] + stage(queue, "runs/run-1/run_manifest.json", b'{"changed":true}', now=201) + assert [entry.content for entry in queue.pending_run_files(REPOSITORY)] == [ + b'{"changed":true}' + ] + + +def test_unforced_hand_overs_wait_for_the_interval_after_the_last_commit(queue): + stage(queue, "runs/run-1/progress.json", b"1", now=1000, flush=False, interval=300) + assert queue.due_repositories(1000) == [REPOSITORY] + commit_everything(queue, now=1000) + stage(queue, "runs/run-1/progress.json", b"2", now=1010, flush=False, interval=300) + assert queue.due_repositories(1010) == [] + assert queue.due_repositories(1299) == [] + assert queue.due_repositories(1300) == [REPOSITORY] + + +def test_forced_hand_overs_keep_the_minimum_spacing_between_commits(queue): + stage(queue, "runs/run-1/progress.json", b"1", now=1000) + commit_everything(queue, now=1000) + stage(queue, "runs/run-1/progress.json", b"2", now=1001) + assert queue.due_repositories(1001) == [] + assert queue.due_repositories(1000 + MIN_COMMIT_SPACING_SECONDS) == [REPOSITORY] + + +def test_a_failed_commit_backs_off_before_the_repository_is_due_again(queue): + stage(queue, "runs/run-1/progress.json", b"1", now=1000) + queue.record_run_files_failed( + REPOSITORY, + queue.pending_run_files(REPOSITORY), + queue.pending_run_index_entries(REPOSITORY), + error_code="UPLOAD_FAILED", + now=1000, + ) + assert queue.due_repositories(1001) == [] + assert queue.due_repositories(1002) == [REPOSITORY] + receipt = queue.run_receipt(REPOSITORY, "run-1") + assert receipt["last_error_code"] == "UPLOAD_FAILED" + assert receipt["commit_attempts"] == 1 and receipt["commits"] == 0 + assert receipt["pending_files"] == 1 + + +def test_repositories_are_due_independently(queue): + stage(queue, "runs/run-1/progress.json", b"1", now=1000) + stage(queue, "runs/run-2/progress.json", b"1", now=1000, run_id="run-2") + stage( + queue, + "runs/run-3/progress.json", + b"1", + now=1000, + run_id="run-3", + repository=OTHER_REPOSITORY, + ) + assert queue.due_repositories(1000) == sorted([REPOSITORY, OTHER_REPOSITORY]) + assert {entry.run_id for entry in queue.pending_run_files(REPOSITORY)} == { + "run-1", + "run-2", + } + + +def test_one_writer_holds_a_repository_at_a_time(tmp_path): + first = HandoverQueue(tmp_path / "events.sqlite3") + second = HandoverQueue(tmp_path / "events.sqlite3") + try: + lease = first.claim_repository(REPOSITORY, now=1000, lease_seconds=60) + assert lease is not None + assert second.claim_repository(REPOSITORY, now=1001, lease_seconds=60) is None + assert second.claim_repository(OTHER_REPOSITORY, now=1001, lease_seconds=60) + first.release_repository(REPOSITORY, lease) + assert second.claim_repository(REPOSITORY, now=1002, lease_seconds=60) + assert first.claim_repository(REPOSITORY, now=1003, lease_seconds=60) is None + assert first.claim_repository(REPOSITORY, now=1063, lease_seconds=60) + finally: + first.close() + second.close() + + +def test_run_index_entries_stay_pending_until_committed(queue): + stage( + queue, + "runs/run-1/progress.json", + b"1", + now=1000, + run_index_entry={"run_id": "run-1", "status": "running"}, + ) + assert queue.pending_run_index_entries(REPOSITORY) == [ + ("run-1", {"run_id": "run-1", "status": "running"}) + ] + commit_everything(queue, now=1000) + assert queue.pending_run_index_entries(REPOSITORY) == [] + + +def test_a_newer_hand_over_during_a_commit_stays_pending(queue): + stage(queue, "runs/run-1/progress.json", b"1", now=1000) + claimed = queue.pending_run_files(REPOSITORY) + stage(queue, "runs/run-1/progress.json", b"2", now=1001) + queue.record_run_files_committed( + REPOSITORY, claimed, [], revision="commit-1", now=1002 + ) + assert [entry.content for entry in queue.pending_run_files(REPOSITORY)] == [b"2"] + assert queue.committed_run_files(REPOSITORY, "run-1") == { + "runs/run-1/progress.json": claimed[0].sha256 + } + + +def test_receipt_counts_committed_files_and_commits_per_run(queue): + queue.stage_run_files( + repository=REPOSITORY, + run_id="run-1", + files=[("runs/run-1/a.json", b"1"), ("runs/run-1/b.json", b"2")], + run_index_entry=None, + flush=True, + interval=30, + now=1000, + ) + commit_everything(queue, now=1000, revision="commit-1") + receipt = queue.run_receipt(REPOSITORY, "run-1") + assert receipt == { + "repository": REPOSITORY, + "run_id": "run-1", + "pending_files": 0, + "files_attempted": 2, + "files_committed": 2, + "commit_attempts": 1, + "commits": 1, + "last_error_code": None, + "last_commit": "commit-1", + } + + +def test_jobs_are_idempotent_by_hand_over_and_survive_a_restart(tmp_path): + path = tmp_path / "events.sqlite3" + queue = HandoverQueue(path) + request = {"repository": REPOSITORY, "run_id": "run-1"} + first = queue.enqueue_job( + handover_id="bundle-1", + kind="dataset_bundle", + target=REPOSITORY, + request=request, + receipt={"status": "pending"}, + now=1000, + ) + again = queue.enqueue_job( + handover_id="bundle-1", + kind="dataset_bundle", + target=REPOSITORY, + request=request, + receipt={"status": "pending"}, + now=1001, + ) + assert first == again == {"status": "pending"} + with pytest.raises(ValueError, match="conflicts"): + queue.enqueue_job( + handover_id="bundle-1", + kind="dataset_bundle", + target=REPOSITORY, + request={"repository": REPOSITORY, "run_id": "other"}, + receipt={"status": "pending"}, + now=1002, + ) + queue.close() + restarted = HandoverQueue(path) + try: + job = restarted.claim_job(REPOSITORY, now=1003, lease_seconds=60) + assert job is not None and job.handover_id == "bundle-1" + assert job.request == request and job.kind == "dataset_bundle" + assert restarted.claim_job(REPOSITORY, now=1004, lease_seconds=60) is None + restarted.finish_job( + job, receipt={"status": "pending", "error_code": "X"}, retry=True, now=1005 + ) + assert restarted.claim_job(REPOSITORY, now=1005, lease_seconds=60) is None + retried = restarted.claim_job(REPOSITORY, now=1007, lease_seconds=60) + assert retried is not None and retried.attempts == 1 + restarted.finish_job( + retried, receipt={"status": "uploaded"}, retry=False, now=1008 + ) + assert restarted.claim_job(REPOSITORY, now=5000, lease_seconds=60) is None + assert restarted.job_receipt("bundle-1") == {"status": "uploaded"} + assert restarted.due_repositories(5000) == [] + finally: + restarted.close() + + +def test_an_expired_worker_cannot_finish_a_reclaimed_job(queue): + queue.enqueue_job( + handover_id="bundle-1", + kind="dataset_bundle", + target=REPOSITORY, + request={}, + receipt={"status": "pending"}, + now=1000, + ) + expired = queue.claim_job(REPOSITORY, now=1000, lease_seconds=5) + current = queue.claim_job(REPOSITORY, now=1006, lease_seconds=60) + assert current is not None + queue.finish_job(expired, receipt={"status": "uploaded"}, retry=False, now=1007) + assert queue.job_receipt("bundle-1") == {"status": "pending"} diff --git a/tests/support/fake_hugging_face.py b/tests/support/fake_hugging_face.py new file mode 100644 index 0000000..14c0af4 --- /dev/null +++ b/tests/support/fake_hugging_face.py @@ -0,0 +1,284 @@ +"""An in-memory Hugging Face Hub for tests: no network, no real repositories. + +It implements the ``HfApi`` methods the emitter calls, with the Hub's +observable behaviour where the emitter depends on it: ``create_commit`` with a +``parent_commit`` that is no longer the head fails with HTTP 412, missing files +and revisions raise errors named like ``huggingface_hub``'s, and every commit +records its exact operations. +""" + +from __future__ import annotations + +import hashlib +import tempfile +import threading +import uuid +from collections.abc import Callable, Iterable, Mapping +from dataclasses import dataclass, field +from pathlib import Path +from types import SimpleNamespace +from typing import Any + + +class FakeHubHTTPError(OSError): + """Carries ``response.status_code`` like ``huggingface_hub.HfHubHTTPError``.""" + + def __init__(self, status_code: int, message: str = "fake Hub error") -> None: + super().__init__(message) + self.response = SimpleNamespace(status_code=status_code) + + +class RevisionNotFoundError(FakeHubHTTPError): + def __init__(self, message: str = "revision not found") -> None: + super().__init__(404, message) + + +class EntryNotFoundError(FakeHubHTTPError): + def __init__(self, message: str = "entry not found") -> None: + super().__init__(404, message) + + +class RepositoryNotFoundError(FakeHubHTTPError): + def __init__(self, message: str = "repository not found") -> None: + super().__init__(404, message) + + +@dataclass +class FakeCommit: + oid: str + parent: str | None + message: str + files: dict[str, bytes] + added: dict[str, bytes] + + +@dataclass +class FakeRepository: + commits: list[FakeCommit] = field(default_factory=list) + tags: dict[str, str] = field(default_factory=dict) + branches: dict[str, str] = field(default_factory=dict) + + @property + def head(self) -> FakeCommit: + return self.commits[-1] + + +def _content(path_or_fileobj: Any) -> bytes: + if isinstance(path_or_fileobj, bytes): + return path_or_fileobj + if isinstance(path_or_fileobj, (str, Path)): + return Path(path_or_fileobj).read_bytes() + return path_or_fileobj.read() + + +class FakeHuggingFaceApi: + """The subset of ``HfApi`` the emitter uses, backed by dictionaries.""" + + def __init__(self, *, whoami: Mapping[str, Any] | None = None) -> None: + self.repositories: dict[str, FakeRepository] = {} + self.calls: list[tuple[str, dict[str, Any]]] = [] + self.whoami_response = whoami + self._failures: dict[str, list[BaseException]] = {} + self._before_commit: list[Callable[[], None]] = [] + self._lock = threading.RLock() + self._download_directory = Path(tempfile.mkdtemp(prefix="fake-hub-")) + self._sequence = 0 + + # Test controls ------------------------------------------------------- + + def create_repository( + self, repository: str, files: Mapping[str, bytes] | None = None + ) -> str: + with self._lock: + self.repositories[repository] = FakeRepository() + return self._append_commit( + repository, dict(files or {}), "Initial commit", parent=None + ) + + def fail_next(self, method: str, error: BaseException, *, times: int = 1) -> None: + """Raise ``error`` from the next ``times`` calls of ``method``.""" + self._failures.setdefault(method, []).extend([error] * times) + + def run_before_next_commit(self, action: Callable[[], None]) -> None: + """Run ``action`` (for example another writer's commit) first.""" + self._before_commit.append(action) + + def files(self, repository: str, revision: str | None = None) -> dict[str, bytes]: + return dict(self._commit_at(repository, revision).files) + + def commits(self, repository: str) -> list[FakeCommit]: + return list(self.repositories[repository].commits[1:]) + + def write_directly( + self, repository: str, files: Mapping[str, bytes], message: str = "Direct" + ) -> str: + """Another writer's commit, outside the code under test.""" + with self._lock: + head = self.repositories[repository].head + return self._append_commit( + repository, {**head.files, **files}, message, parent=head.oid + ) + + # HfApi surface ------------------------------------------------------- + + def _record(self, method: str, **arguments: Any) -> None: + self.calls.append((method, arguments)) + pending = self._failures.get(method) + if pending: + raise pending.pop(0) + + def _repository(self, repo_id: str) -> FakeRepository: + repository = self.repositories.get(repo_id) + if repository is None: + raise RepositoryNotFoundError(repo_id) + return repository + + def _commit_at(self, repo_id: str, revision: str | None) -> FakeCommit: + repository = self._repository(repo_id) + if revision in (None, "main"): + return repository.head + if revision in repository.tags: + revision = repository.tags[revision] + if revision in repository.branches: + revision = repository.branches[revision] + for commit in repository.commits: + if commit.oid == revision: + return commit + raise RevisionNotFoundError(str(revision)) + + def _append_commit( + self, + repository: str, + files: dict[str, bytes], + message: str, + *, + parent: str | None, + added: dict[str, bytes] | None = None, + ) -> str: + self._sequence += 1 + oid = hashlib.sha1(f"{repository}:{self._sequence}".encode()).hexdigest() + self.repositories[repository].commits.append( + FakeCommit(oid, parent, message, files, dict(added or {})) + ) + return oid + + def repo_info( + self, repo_id: str, *, repo_type: str | None = None, revision=None, **_: Any + ): + self._record("repo_info", repo_id=repo_id, revision=revision) + return SimpleNamespace(sha=self._commit_at(repo_id, revision).oid) + + def create_commit( + self, + repo_id: str, + operations: Iterable[Any], + *, + commit_message: str, + repo_type: str | None = None, + revision: str | None = None, + parent_commit: str | None = None, + **_: Any, + ): + operations = list(operations) + self._record( + "create_commit", + repo_id=repo_id, + operations=[operation.path_in_repo for operation in operations], + commit_message=commit_message, + revision=revision, + parent_commit=parent_commit, + ) + while self._before_commit: + self._before_commit.pop(0)() + with self._lock: + repository = self._repository(repo_id) + head = repository.head + if parent_commit is not None and parent_commit != head.oid: + raise FakeHubHTTPError(412, "A commit has happened since.") + added = { + operation.path_in_repo: _content(operation.path_or_fileobj) + for operation in operations + } + oid = self._append_commit( + repo_id, + {**head.files, **added}, + commit_message, + parent=head.oid, + added=added, + ) + return SimpleNamespace(oid=oid, commit_url=f"https://fake/{oid}") + + def hf_hub_download( + self, + repo_id: str, + filename: str, + *, + repo_type: str | None = None, + revision: str | None = None, + **_: Any, + ) -> str: + self._record( + "hf_hub_download", repo_id=repo_id, filename=filename, revision=revision + ) + files = self._commit_at(repo_id, revision).files + if filename not in files: + raise EntryNotFoundError(filename) + target = self._download_directory / uuid_name(filename) + target.write_bytes(files[filename]) + return str(target) + + def file_exists( + self, + repo_id: str, + filename: str, + *, + repo_type: str | None = None, + revision: str | None = None, + **_: Any, + ) -> bool: + self._record("file_exists", repo_id=repo_id, filename=filename) + return filename in self._commit_at(repo_id, revision).files + + def list_repo_files( + self, repo_id: str, *, revision: str | None = None, **_: Any + ) -> list[str]: + self._record("list_repo_files", repo_id=repo_id, revision=revision) + return sorted(self._commit_at(repo_id, revision).files) + + def get_paths_info( + self, repo_id: str, paths, *, expand: bool = False, revision=None, **_: Any + ): + self._record("get_paths_info", repo_id=repo_id, paths=list(paths)) + repository = self._repository(repo_id) + entries = [] + for path in paths: + last = None + previous = None + for commit in repository.commits: + content = commit.files.get(path) + if content is not None and content != previous: + last = commit.oid + previous = content + if last is not None: + entries.append( + SimpleNamespace(path=path, last_commit=SimpleNamespace(oid=last)) + ) + return entries + + def whoami(self, *args: Any, **kwargs: Any) -> dict[str, Any]: + self._record("whoami") + if self.whoami_response is None: + raise FakeHubHTTPError(401, "no credential") + return dict(self.whoami_response) + + def create_tag(self, repo_id: str, **arguments: Any) -> None: + self._record("create_tag", repo_id=repo_id, **arguments) + raise AssertionError("The emitter must never create a Hugging Face tag.") + + def create_branch(self, repo_id: str, **arguments: Any) -> None: + self._record("create_branch", repo_id=repo_id, **arguments) + raise AssertionError("The emitter must never create a Hugging Face branch.") + + +def uuid_name(filename: str) -> str: + return f"{uuid.uuid4().hex}-{Path(filename).name}" diff --git a/tests/telemetry/test_persistence.py b/tests/telemetry/test_persistence.py index c8fdbe3..a850847 100644 --- a/tests/telemetry/test_persistence.py +++ b/tests/telemetry/test_persistence.py @@ -37,7 +37,7 @@ def test_migrations_match_models_and_reopening_preserves_sequences(tmp_path): ] == [1, 2] assert spool.batch("fixture-run", "producer-a")[0] == first assert first["event_id"] != second["event_id"] - assert current_database_revision(path) == "20261008_02" + assert current_database_revision(path) == "20261020_03" with ( spool._engine.begin() as connection, alembic_config(connection=connection) as config, @@ -83,5 +83,9 @@ def test_multiple_producers_and_local_only_survive_reopening(tmp_path): "telemetry_runs", "telemetry_events", "graph_publication_jobs", + "staged_run_files", + "staged_run_deliveries", + "repository_commit_schedules", + "handover_jobs", } spool.close() diff --git a/tools/installed-smoke.py b/tools/installed-smoke.py index a886317..1ab3b5c 100644 --- a/tools/installed-smoke.py +++ b/tools/installed-smoke.py @@ -56,7 +56,7 @@ def telemetry_smoke(): ) from microcosm_emitter.telemetry.service.spool import EventSpool - assert current_database_revision(path) == "20261008_02" + assert current_database_revision(path) == "20261020_03" spool = EventSpool(path) try: assert not spool.has_deliverable() From 6e8ced3e274abfa72ba4dc7b0ec382f04cacb100 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:57:50 +0400 Subject: [PATCH 02/23] Stage run files with one Hugging Face commit per repository per cycle The new `huggingface` emitter module takes over the transport of Microcosm's staging writers (staging.py, staging_v2.py): - `stage_run_files` copies the handed-over files' bytes into the queue before answering. They are written at `//` with the bytes the build wrote; staging contract version 1 also writes `latest_staging.json` and merges its run summary into the repository's `runs.json` with Microcosm's ordering and formatting. - Each module tick makes at most one commit per due repository, with the newest bytes of every pending path. Unforced hand-overs wait for the build's upload interval after the last commit; forced ones go at the next cycle, at least 10 seconds after the previous commit. - Failures keep files queued and back off; without a credential one warning is printed. Close commits pending files once more. - A `runs.json` merge commits with `parent_commit` and, on a conflict, reads the index again and merges again. - `read_back_run_files` compares the Hub's copies with the committed bytes. - Requests are checked with the interface's shared path rules; staging never writes under `releases/`. - Protocol fixtures for the new messages; the installed-wheel check starts the module without a credential. Co-Authored-By: Claude Opus 5.5 --- README.md | 13 + .../huggingface-staging-run-files.added.md | 1 + .../microcosm_emitter_interface/__init__.py | 8 + .../src/microcosm_emitter_interface/client.py | 81 +++- .../huggingface-configuration.json | 8 + .../huggingface-read-back-run-files.json | 15 + .../huggingface-run-files-receipt.json | 15 + ...uggingface-stage-run-files-with-index.json | 33 ++ .../huggingface-stage-run-files.json | 34 ++ .../identifiers.py | 72 +++ .../microcosm_emitter_interface/protocol.py | 7 +- .../microcosm_emitter_interface/staging.py | 192 ++++++++ pyproject.toml | 1 + src/microcosm_emitter/core/models.py | 3 +- .../huggingface/service/__init__.py | 9 + .../huggingface/service/delivery.py | 147 +++++++ .../huggingface/service/module.py | 182 ++++++++ .../huggingface/service/queue.py | 35 +- .../huggingface/service/requests.py | 143 ++++++ .../huggingface/service/run_files.py | 154 +++++++ tests/huggingface/conftest.py | 61 +++ tests/huggingface/test_run_files.py | 409 ++++++++++++++++++ tests/interface/test_client_messages.py | 60 ++- tests/interface/test_protocol_fixtures.py | 44 ++ tests/interface/test_staging_client.py | 205 +++++++++ tests/support/fake_hugging_face.py | 3 + tests/test_package_names.py | 7 + tools/installed-interface-smoke.py | 19 +- 28 files changed, 1952 insertions(+), 9 deletions(-) create mode 100644 changelog.d/huggingface-staging-run-files.added.md create mode 100644 packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-configuration.json create mode 100644 packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-read-back-run-files.json create mode 100644 packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-run-files-receipt.json create mode 100644 packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-stage-run-files-with-index.json create mode 100644 packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-stage-run-files.json create mode 100644 packages/microcosm-emitter-interface/src/microcosm_emitter_interface/identifiers.py create mode 100644 packages/microcosm-emitter-interface/src/microcosm_emitter_interface/staging.py create mode 100644 src/microcosm_emitter/huggingface/service/delivery.py create mode 100644 src/microcosm_emitter/huggingface/service/module.py create mode 100644 src/microcosm_emitter/huggingface/service/requests.py create mode 100644 src/microcosm_emitter/huggingface/service/run_files.py create mode 100644 tests/huggingface/conftest.py create mode 100644 tests/huggingface/test_run_files.py create mode 100644 tests/interface/test_staging_client.py diff --git a/README.md b/README.md index 44d9b7a..3bab13a 100644 --- a/README.md +++ b/README.md @@ -54,6 +54,19 @@ The host can also start telemetry and graph modules together with `start_service Preserve the selected source files until publication succeeds. `microcosm-emitter-publish-graph` asks the service to retry a publication by ID and can re-enqueue a preserved directory on another machine. It never rewrites the original build receipt; the queue and separate latest-status file record retries. +## Hugging Face hand-overs + +The `huggingface` module takes over the uploads Microcosm made itself. Microcosm still decides what to hand over and validates its own contracts; the module only writes the bytes it was given, with the build's own Hugging Face credential. + +**Staging run files** (Microcosm's `staging.py` and `staging_v2.py` uploads). A build hands over its current run files; the module copies their bytes into the spool database before answering, so the build may rewrite them at once. + +- Files land at `//` (by default `runs//...`), exactly as the build's writer lays them out. Staging contract version 1 also writes `latest_staging.json` at the repository root, and its `runs.json` index: the module merges the run's summary into the repository's current copy at commit time, as Microcosm did, keeping other runs and the same JSON formatting. +- Each path keeps only the newest bytes, and bytes equal to those last committed are not written again. +- **One commit per repository per cycle.** Each module tick commits, for each repository with work due, every pending path in one commit. An unforced hand-over is due one interval after the repository's last commit (Microcosm's upload interval); a forced one (final states, artifacts) at the next cycle, but never sooner than 10 seconds after the previous commit. Microcosm's writers made one commit per file, up to eight per upload cycle, and each commit triggered the dataset repository's webhook. +- Uploads stay best-effort. Requests answer within seconds or report that they were not accepted; a failed commit keeps the files queued and retries with backoff (up to five minutes); without a credential the files stay queued and one warning is printed. On shutdown the module commits pending files once more. Nothing raises into the build. +- A commit that merges `runs.json` is pinned to the head it read (`parent_commit`); if another build commits first, the module reads the index again and merges again. +- `read_back_run_files` commits pending files, then compares the Hub's copy of each committed file with the committed bytes. + ## Development Use Python 3.13 or 3.14 on Linux or macOS. diff --git a/changelog.d/huggingface-staging-run-files.added.md b/changelog.d/huggingface-staging-run-files.added.md new file mode 100644 index 0000000..a7966a0 --- /dev/null +++ b/changelog.d/huggingface-staging-run-files.added.md @@ -0,0 +1 @@ +Add the `huggingface` emitter module and the interface's `stage_run_files`, `run_files_receipt` and `read_back_run_files` calls. A build hands over its staging run files; the emitter writes them at the same paths with the same bytes as Microcosm's staging writers (`runs//...`, and for staging contract version 1 `latest_staging.json` and a merged `runs.json`), in one Hugging Face commit per repository per cycle instead of one commit per file. Uploads stay best-effort and are retried from the durable queue. diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/__init__.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/__init__.py index c02750d..ccb71f1 100644 --- a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/__init__.py +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/__init__.py @@ -12,6 +12,11 @@ from microcosm_emitter_interface.orrery import OrrerySettings from microcosm_emitter_interface.protocol import PROTOCOL_VERSION from microcosm_emitter_interface.spool import default_spool_path +from microcosm_emitter_interface.staging import ( + HuggingFaceSettings, + RunFile, + StageRunFiles, +) from microcosm_emitter_interface.telemetry import TelemetryRun, TelemetrySettings __all__ = [ @@ -20,8 +25,11 @@ "PROTOCOL_VERSION", "EmitterClient", "EmitterStartupError", + "HuggingFaceSettings", "OrrerySettings", "RequestFailure", + "RunFile", + "StageRunFiles", "TelemetryRun", "TelemetrySettings", "default_spool_path", diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py index 1ba3f3e..cddf20e 100644 --- a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py @@ -17,7 +17,7 @@ from collections.abc import Callable, Mapping from dataclasses import dataclass from pathlib import Path -from typing import Any, Literal +from typing import Any, Literal, cast from microcosm_emitter_interface.launcher import ( EmitterProcess, @@ -41,11 +41,23 @@ from microcosm_emitter_interface.protocol import ( ACKNOWLEDGEMENT_TIMEOUT_SECONDS, FILE_REQUEST_TIMEOUT_SECONDS, + HUGGING_FACE_MODULE, ORRERY_MODULE, + READ_BACK_TIMEOUT_SECONDS, + STAGING_REQUEST_TIMEOUT_SECONDS, STARTUP_TIMEOUT_SECONDS, TELEMETRY_MODULE, ) from microcosm_emitter_interface.spool import default_spool_path +from microcosm_emitter_interface.staging import ( + HuggingFaceSettings, + ReadBackResult, + RunFilesReceipt, + StageRunFiles, + run_reference_message, + unavailable_read_back, + unavailable_run_files_receipt, +) from microcosm_emitter_interface.telemetry import ( TelemetryCommand, TelemetryRecorder, @@ -109,6 +121,7 @@ def start( *, telemetry: TelemetrySettings | None = None, orrery: OrrerySettings | None = None, + hugging_face: HuggingFaceSettings | None = None, spool_path: Path | str | None = None, command: str | os.PathLike[str] | None = None, startup_timeout: float = STARTUP_TIMEOUT_SECONDS, @@ -121,6 +134,8 @@ def start( modules[TELEMETRY_MODULE] = telemetry.configuration(path) if orrery is not None: modules[ORRERY_MODULE] = orrery.configuration(path) + if hugging_face is not None: + modules[HUGGING_FACE_MODULE] = hugging_face.configuration(path) except Exception: failure = EmitterStartupError("invalid_configuration") print(STARTUP_WARNING.format(reason=failure.reason), file=sys.stderr) @@ -323,6 +338,70 @@ def _orrery_receipt_request( answer, publication_id, error_code=_receipt_error_code(failure) ) + def stage_run_files(self, request: StageRunFiles) -> RunFilesReceipt: + """Hand over a staging run's files; returns once they are queued. + + The emitter copies the files' bytes before answering, so the build + may rewrite them immediately afterwards. + """ + message = request.message() + response = self._query( + HUGGING_FACE_MODULE, + lambda: message, + kind="staging", + timeout=STAGING_REQUEST_TIMEOUT_SECONDS, + ) + value = None if response is None else response.get("receipt") + if isinstance(value, dict): + return cast(RunFilesReceipt, value) + return unavailable_run_files_receipt(request.repository, request.run_id) + + def run_files_receipt(self, repository: str, run_id: str) -> RunFilesReceipt: + message = run_reference_message("run_files_receipt", repository, run_id) + response = self._query( + HUGGING_FACE_MODULE, + lambda: message, + kind="staging", + timeout=STAGING_REQUEST_TIMEOUT_SECONDS, + ) + value = None if response is None else response.get("receipt") + if isinstance(value, dict): + return cast(RunFilesReceipt, value) + return unavailable_run_files_receipt(repository, run_id) + + def read_back_run_files( + self, repository: str, run_id: str, *, wait_seconds: float = 30.0 + ) -> ReadBackResult: + """Commit the run's pending files now, then compare the Hub's copies. + + Waits up to ``wait_seconds`` for pending files to be committed; a run + whose files are still pending then reads back as ``pending``. + """ + if HUGGING_FACE_MODULE not in self.modules: + return unavailable_read_back() + receipt = self.stage_run_files( + StageRunFiles(repository=repository, run_id=run_id, flush=True) + ) + deadline = time.monotonic() + max(0, wait_seconds) + while ( + receipt["status"] == "recorded" + and receipt["pending_files"] + and time.monotonic() < deadline + ): + time.sleep(0.1) + receipt = self.run_files_receipt(repository, run_id) + message = run_reference_message("read_back_run_files", repository, run_id) + response = self._query( + HUGGING_FACE_MODULE, + lambda: message, + kind="staging", + timeout=READ_BACK_TIMEOUT_SECONDS, + ) + value = None if response is None else response.get("read_back") + if isinstance(value, dict): + return cast(ReadBackResult, value) + return unavailable_read_back() + def close(self) -> None: """Ask the emitter to drain and stop; repeated calls do nothing.""" if self._closed: diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-configuration.json b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-configuration.json new file mode 100644 index 0000000..0288ad2 --- /dev/null +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-configuration.json @@ -0,0 +1,8 @@ +{ + "protocol_version": 1, + "module": "huggingface", + "kind": "configuration", + "body": { + "spool_path": "{spool_path}" + } +} diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-read-back-run-files.json b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-read-back-run-files.json new file mode 100644 index 0000000..7f428fc --- /dev/null +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-read-back-run-files.json @@ -0,0 +1,15 @@ +{ + "protocol_version": 1, + "module": "huggingface", + "kind": "query", + "sequence": 4, + "body": { + "protocol_version": 1, + "command": "read_back_run_files", + "repository": "policyengine/fixture-staging", + "run_id": "fixture-run" + }, + "response_keys": [ + "read_back" + ] +} diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-run-files-receipt.json b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-run-files-receipt.json new file mode 100644 index 0000000..4e2004f --- /dev/null +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-run-files-receipt.json @@ -0,0 +1,15 @@ +{ + "protocol_version": 1, + "module": "huggingface", + "kind": "query", + "sequence": 3, + "body": { + "protocol_version": 1, + "command": "run_files_receipt", + "repository": "policyengine/fixture-staging", + "run_id": "fixture-run" + }, + "response_keys": [ + "receipt" + ] +} diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-stage-run-files-with-index.json b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-stage-run-files-with-index.json new file mode 100644 index 0000000..f7c6ae5 --- /dev/null +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-stage-run-files-with-index.json @@ -0,0 +1,33 @@ +{ + "protocol_version": 1, + "module": "huggingface", + "kind": "query", + "sequence": 2, + "files": { + "progress.json": "{\"stage\":\"calibrating\"}", + "latest_staging.json": "{\"run_id\":\"fixture-run\"}" + }, + "body": { + "protocol_version": 1, + "command": "stage_run_files", + "repository": "policyengine/fixture-staging", + "run_id": "fixture-run", + "prefix": "runs", + "files": [ + { + "relative_path": "progress.json", + "local_path": "{directory}/progress.json" + } + ], + "latest_pointer": "{directory}/latest_staging.json", + "run_index_entry": { + "run_id": "fixture-run", + "status": "running" + }, + "flush": false, + "commit_interval_seconds": 300.0 + }, + "response_keys": [ + "receipt" + ] +} diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-stage-run-files.json b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-stage-run-files.json new file mode 100644 index 0000000..6dff8eb --- /dev/null +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-stage-run-files.json @@ -0,0 +1,34 @@ +{ + "protocol_version": 1, + "module": "huggingface", + "kind": "query", + "sequence": 1, + "files": { + "progress.json": "{\"stage\":\"created\"}", + "fit_summary.json": "{\"final_loss\":0.5}" + }, + "body": { + "protocol_version": 1, + "command": "stage_run_files", + "repository": "policyengine/fixture-staging", + "run_id": "fixture-run", + "prefix": "runs", + "files": [ + { + "relative_path": "progress.json", + "local_path": "{directory}/progress.json" + }, + { + "relative_path": "artifacts/fit_summary.json", + "local_path": "{directory}/fit_summary.json" + } + ], + "latest_pointer": null, + "run_index_entry": null, + "flush": true, + "commit_interval_seconds": 30.0 + }, + "response_keys": [ + "receipt" + ] +} diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/identifiers.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/identifiers.py new file mode 100644 index 0000000..9df1352 --- /dev/null +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/identifiers.py @@ -0,0 +1,72 @@ +"""Identifier and path rules shared by the interface and the emitter. + +The interface refuses to build a request that breaks them, and the emitter +checks every request it receives with the same functions. +""" + +from __future__ import annotations + +import re +from pathlib import PurePosixPath +from typing import Final + +SAFE_IDENTIFIER: Final = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") +REPOSITORY_IDENTIFIER: Final = re.compile( + r"^[A-Za-z0-9][A-Za-z0-9._-]{0,95}/[A-Za-z0-9][A-Za-z0-9._-]{0,95}$" +) +ROLE_IDENTIFIER: Final = re.compile(r"^[a-z][a-z0-9_]{0,63}$") +#: Folders that only the release writer changes; staging never writes there. +RESERVED_PREFIXES: Final = frozenset({"releases"}) + + +def safe_identifier(value: object, *, label: str) -> str: + if not isinstance(value, str) or not SAFE_IDENTIFIER.fullmatch(value): + raise ValueError(f"{label} must match {SAFE_IDENTIFIER.pattern}.") + return value + + +def repository_identifier(value: object) -> str: + """A Hugging Face repository identifier, ``owner/name``.""" + if not isinstance(value, str) or not REPOSITORY_IDENTIFIER.fullmatch(value): + raise ValueError("repository must be a Hugging Face owner/name identifier.") + return value + + +def role_identifier(value: object, *, label: str = "role") -> str: + if not isinstance(value, str) or not ROLE_IDENTIFIER.fullmatch(value): + raise ValueError(f"{label} must match {ROLE_IDENTIFIER.pattern}.") + return value + + +def repository_prefix(value: object) -> str: + """A folder of safe segments that is not reserved for release records.""" + if not isinstance(value, str) or not value: + raise ValueError("prefix must be a nonempty relative folder.") + segments = value.split("/") + for segment in segments: + safe_identifier(segment, label="prefix segment") + if segments[0] in RESERVED_PREFIXES: + raise ValueError(f"prefix may not start with {segments[0]!r}.") + return value + + +def relative_file_path(value: object, *, label: str) -> str: + """A relative POSIX path that cannot leave its folder.""" + if not isinstance(value, str) or not value or "\\" in value: + raise ValueError(f"{label} must be a nonempty POSIX relative path.") + path = PurePosixPath(value) + if ( + path.is_absolute() + or any(part in {"", ".", ".."} for part in value.split("/")) + or path.as_posix() != value + ): + raise ValueError(f"{label} must stay inside its folder.") + return value + + +def bare_file_name(value: object, *, label: str) -> str: + """One path segment: a file name without folders.""" + name = relative_file_path(value, label=label) + if "/" in name: + raise ValueError(f"{label} must be a bare file name.") + return name diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/protocol.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/protocol.py index 532cdc0..3c218d7 100644 --- a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/protocol.py +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/protocol.py @@ -18,12 +18,17 @@ TELEMETRY_MODULE: Final = "telemetry" ORRERY_MODULE: Final = "orrery" -type ModuleName = Literal["telemetry", "orrery"] +HUGGING_FACE_MODULE: Final = "huggingface" +type ModuleName = Literal["telemetry", "orrery", "huggingface"] #: Telemetry requests only wait for the service to commit their events. ACKNOWLEDGEMENT_TIMEOUT_SECONDS: Final = 0.2 #: Orrery requests hash a bounded set of preserved files before answering. FILE_REQUEST_TIMEOUT_SECONDS: Final = 30.0 +#: Staging hand-overs copy a few small files into the emitter's queue. +STAGING_REQUEST_TIMEOUT_SECONDS: Final = 5.0 +#: Read-back downloads a run's committed files from the Hub. +READ_BACK_TIMEOUT_SECONDS: Final = 60.0 STARTUP_TIMEOUT_SECONDS: Final = 15.0 STARTUP_POLL_SECONDS: Final = 0.02 DRAIN_SECONDS: Final = 15.0 diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/staging.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/staging.py new file mode 100644 index 0000000..06c4a19 --- /dev/null +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/staging.py @@ -0,0 +1,192 @@ +"""Typed staging hand-overs: run files under ``//``. + +The build writes and validates its staging documents locally, as before, and +hands each file over by path. The emitter copies the bytes into its queue +before answering, then writes every pending path of a repository in one +Hugging Face commit per cycle. Uploads stay best-effort: no call here raises +because the emitter or the Hub is unavailable. + +Constructing a request that breaks the path rules raises ``ValueError``. +""" + +from __future__ import annotations + +import math +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Literal, NotRequired, TypedDict + +from microcosm_emitter_interface.identifiers import ( + relative_file_path, + repository_identifier, + repository_prefix, + safe_identifier, +) +from microcosm_emitter_interface.protocol import PROTOCOL_VERSION +from microcosm_emitter_interface.serialization import wire_value + +DEFAULT_RUN_PREFIX = "runs" +DEFAULT_COMMIT_INTERVAL_SECONDS = 30.0 + + +@dataclass(frozen=True) +class HuggingFaceSettings: + """Configuration for the emitter's Hugging Face module.""" + + def configuration(self, spool_path: Path) -> dict[str, Any]: + return {"spool_path": str(spool_path)} + + +@dataclass(frozen=True) +class RunFile: + """One local file, written at ``//``.""" + + relative_path: str + local_path: Path + + def __post_init__(self) -> None: + relative_file_path(self.relative_path, label="relative_path") + + +@dataclass(frozen=True) +class StageRunFiles: + """Hand over a staging run's current files. + + - ``files`` land under ``//``, as Microcosm's staging + writers lay them out (``runs//progress.json``, ...). + - ``latest_pointer`` is staging contract version 1's pointer file, written + to ``latest_staging.json`` at the repository root. + - ``run_index_entry`` is version 1's summary of this run; the emitter + merges it into the repository's ``runs.json`` when it commits. + - An unforced hand-over is committed one ``commit_interval_seconds`` after + the repository's last commit; ``flush`` asks for the next cycle (final + states and artifacts). Either way a cycle makes one commit per + repository, carrying the newest bytes of every pending path. + """ + + repository: str + run_id: str + files: Sequence[RunFile] = () + prefix: str = DEFAULT_RUN_PREFIX + latest_pointer: Path | None = None + run_index_entry: Mapping[str, Any] | None = None + flush: bool = False + commit_interval_seconds: float = DEFAULT_COMMIT_INTERVAL_SECONDS + protocol_version: int = field(default=PROTOCOL_VERSION, repr=False) + + def __post_init__(self) -> None: + repository_identifier(self.repository) + safe_identifier(self.run_id, label="run_id") + repository_prefix(self.prefix) + object.__setattr__(self, "files", tuple(self.files)) + if any(not isinstance(item, RunFile) for item in self.files): + raise ValueError("files must be RunFile values.") + if not isinstance(self.flush, bool): + raise ValueError("flush must be a boolean.") + interval = self.commit_interval_seconds + if ( + isinstance(interval, bool) + or not isinstance(interval, (int, float)) + or not math.isfinite(interval) + or interval < 0 + ): + raise ValueError("commit_interval_seconds must be finite and >= 0.") + object.__setattr__(self, "commit_interval_seconds", float(interval)) + if self.run_index_entry is not None and ( + not isinstance(self.run_index_entry, Mapping) + or self.run_index_entry.get("run_id") != self.run_id + ): + raise ValueError("run_index_entry must describe this run.") + + def message(self) -> dict[str, Any]: + return { + "protocol_version": self.protocol_version, + "command": "stage_run_files", + "repository": self.repository, + "run_id": self.run_id, + "prefix": self.prefix, + "files": [ + { + "relative_path": item.relative_path, + "local_path": str(Path(item.local_path).resolve()), + } + for item in self.files + ], + "latest_pointer": ( + None + if self.latest_pointer is None + else str(Path(self.latest_pointer).resolve()) + ), + "run_index_entry": ( + None + if self.run_index_entry is None + else wire_value(dict(self.run_index_entry)) + ), + "flush": self.flush, + "commit_interval_seconds": self.commit_interval_seconds, + } + + +def run_reference_message(command: str, repository: str, run_id: str) -> dict[str, Any]: + return { + "protocol_version": PROTOCOL_VERSION, + "command": command, + "repository": repository_identifier(repository), + "run_id": safe_identifier(run_id, label="run_id"), + } + + +class RunFilesReceipt(TypedDict): + """What the emitter holds for one run, and what reached the Hub. + + ``status`` is ``recorded`` once the emitter answered, ``refused`` for a + request it will not queue (``refusal_code`` says why), and + ``unavailable`` when no emitter answered. File counts include each + file of every attempted commit. + """ + + status: Literal["recorded", "refused", "unavailable"] + repository: str + run_id: str + pending_files: int + files_attempted: int + files_committed: int + commit_attempts: int + commits: int + last_error_code: str | None + last_commit: str | None + refusal_code: NotRequired[str] + + +class ReadBackResult(TypedDict): + """Whether the Hub holds exactly the bytes the emitter committed.""" + + status: Literal["passed", "failed", "pending", "refused", "unavailable"] + files: int + mismatched_paths: list[str] + error_code: str | None + + +def unavailable_run_files_receipt(repository: str, run_id: str) -> RunFilesReceipt: + return { + "status": "unavailable", + "repository": repository, + "run_id": run_id, + "pending_files": 0, + "files_attempted": 0, + "files_committed": 0, + "commit_attempts": 0, + "commits": 0, + "last_error_code": None, + "last_commit": None, + } + + +def unavailable_read_back() -> ReadBackResult: + return { + "status": "unavailable", + "files": 0, + "mismatched_paths": [], + "error_code": None, + } diff --git a/pyproject.toml b/pyproject.toml index 3f3b96b..c7066ab 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,7 @@ Source = "https://github.com/PolicyEngine/microcosm-emitter" [project.entry-points."microcosm.emitter.modules"] telemetry = "microcosm_emitter.telemetry.service:create_module" orrery = "microcosm_emitter.orrery.service:create_module" +huggingface = "microcosm_emitter.huggingface.service:create_module" [project.scripts] microcosm-emitter = "microcosm_emitter.cli:main" diff --git a/src/microcosm_emitter/core/models.py b/src/microcosm_emitter/core/models.py index d116285..5472756 100644 --- a/src/microcosm_emitter/core/models.py +++ b/src/microcosm_emitter/core/models.py @@ -178,7 +178,8 @@ class StagedRunDelivery(SpoolModel): commits: Mapped[int] = mapped_column(Integer, nullable=False) last_error_code: Mapped[str | None] = mapped_column(Text) last_commit: Mapped[str | None] = mapped_column(Text) - run_index_entry: Mapped[dict[str, Any] | None] = mapped_column(JsonObjectText) + # JSON text in the build's key order: the merged runs.json keeps it. + run_index_entry: Mapped[str | None] = mapped_column(Text) run_index_pending: Mapped[bool] = mapped_column(Boolean, nullable=False) run_index_due_at: Mapped[float] = mapped_column(Float, nullable=False) updated_at: Mapped[float] = mapped_column(Float, nullable=False) diff --git a/src/microcosm_emitter/huggingface/service/__init__.py b/src/microcosm_emitter/huggingface/service/__init__.py index cdb589d..04c8be6 100644 --- a/src/microcosm_emitter/huggingface/service/__init__.py +++ b/src/microcosm_emitter/huggingface/service/__init__.py @@ -1 +1,10 @@ """The Hugging Face hand-over module; constructed only by the service host.""" + +from microcosm_emitter.host.contracts import JsonObject, ModuleContext, ServiceModule + + +def create_module(configuration: JsonObject, context: ModuleContext) -> ServiceModule: + """Construct the module without starting workers or performing network I/O.""" + from microcosm_emitter.huggingface.service.module import HuggingFaceModule + + return HuggingFaceModule(configuration, context) diff --git a/src/microcosm_emitter/huggingface/service/delivery.py b/src/microcosm_emitter/huggingface/service/delivery.py new file mode 100644 index 0000000..05245ce --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/delivery.py @@ -0,0 +1,147 @@ +"""The delivery cycle: at most one Hugging Face commit per repository per cycle. + +Each cycle (one module tick) finds the repositories with work due, leases each +one, and makes one commit carrying every pending staging run file of that +repository. Failures keep the work queued and back off; nothing here raises +into the build. +""" + +from __future__ import annotations + +import sys +import threading +import time +from collections.abc import Callable +from datetime import UTC, datetime +from typing import Any + +from microcosm_emitter.huggingface.service.hub import DatasetRepository +from microcosm_emitter.huggingface.service.queue import HandoverQueue +from microcosm_emitter.huggingface.service.run_files import commit_run_files + +#: Long enough for one commit of a large upload; a crashed writer's lease +#: expires and another host on this machine resumes the repository. +REPOSITORY_LEASE_SECONDS = 900.0 +CREDENTIAL_MISSING = "CREDENTIAL_MISSING" +UPLOAD_FAILED = "UPLOAD_FAILED" +CREDENTIAL_MISSING_MESSAGE = ( + "Microcosm Hugging Face uploads are paused: no Hugging Face credential was " + "found (HF_TOKEN, HUGGINGFACE_TOKEN or `hf auth login`). Handed-over files " + "stay queued; the build will continue." +) +UPLOAD_FAILED_MESSAGE = ( + "A Microcosm Hugging Face upload failed ({error_type}); handed-over files " + "stay queued for a later attempt and the build will continue." +) + + +def iso_timestamp(seconds: float) -> str: + """Microcosm ``staging._now()`` format, at ``seconds``.""" + return datetime.fromtimestamp(seconds, UTC).replace(microsecond=0).isoformat() + + +class HuggingFaceDelivery: + """Commit queued hand-overs with the build's credential.""" + + def __init__( + self, + queue: HandoverQueue, + *, + api_factory: Callable[[], Any | None], + clock: Callable[[], float] = time.time, + ) -> None: + self.queue = queue + self.api_factory = api_factory + self.clock = clock + self._warned: set[str] = set() + self._warning_lock = threading.Lock() + + def warn_once(self, key: str, message: str) -> None: + with self._warning_lock: + if key in self._warned: + return + self._warned.add(key) + print(message, file=sys.stderr, flush=True) + + def run_cycle(self) -> int: + """Make at most one commit per due repository; return how many.""" + repositories = self.queue.due_repositories(self.clock()) + if not repositories: + return 0 + api = self.api_factory() + return sum( + self._deliver_repository(repository, api) for repository in repositories + ) + + def drain(self, deadline: float) -> None: + """At shutdown, commit every repository's pending run files once. + + Ignores the commit interval and spacing, but still leases each + repository. ``deadline`` is a ``time.monotonic`` value. + """ + api = self.api_factory() + for repository in self.queue.repositories_with_pending_run_files(): + if time.monotonic() >= deadline: + return + self._deliver_repository(repository, api, run_files_only=True) + + def _deliver_repository( + self, repository: str, api: Any | None, *, run_files_only: bool = False + ) -> int: + lease = self.queue.claim_repository( + repository, now=self.clock(), lease_seconds=REPOSITORY_LEASE_SECONDS + ) + if lease is None: + return 0 + try: + if api is None: + self._credential_missing(repository) + return 0 + hub = DatasetRepository(repository, api) + return int(self._deliver_run_files(hub)) + finally: + self.queue.release_repository(repository, lease) + + def _credential_missing(self, repository: str) -> None: + self.warn_once(CREDENTIAL_MISSING, CREDENTIAL_MISSING_MESSAGE) + entries = self.queue.pending_run_files(repository) + index_entries = self.queue.pending_run_index_entries(repository) + if entries or index_entries: + self.queue.record_run_files_failed( + repository, + entries, + index_entries, + error_code=CREDENTIAL_MISSING, + now=self.clock(), + ) + + def _deliver_run_files(self, hub: DatasetRepository) -> bool: + entries = self.queue.pending_run_files(hub.repo_id) + index_entries = self.queue.pending_run_index_entries(hub.repo_id) + if not entries and not index_entries: + return False + try: + revision = commit_run_files( + hub, + entries, + index_entries, + updated_at=iso_timestamp(self.clock()), + ) + except Exception as error: + # Only the type name: Hub errors can carry URLs and identifiers. + self.warn_once( + UPLOAD_FAILED, + UPLOAD_FAILED_MESSAGE.format(error_type=type(error).__name__), + ) + self.queue.record_run_files_failed( + hub.repo_id, + entries, + index_entries, + error_code=UPLOAD_FAILED, + now=self.clock(), + ) + return False + self.queue.record_run_files_committed( + hub.repo_id, entries, index_entries, revision=revision, now=self.clock() + ) + return True diff --git a/src/microcosm_emitter/huggingface/service/module.py b/src/microcosm_emitter/huggingface/service/module.py new file mode 100644 index 0000000..8739b92 --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/module.py @@ -0,0 +1,182 @@ +"""The Hugging Face hand-over module of the local emitter host. + +Requests are queries: the module checks and durably records a hand-over, then +answers with its receipt. The module's worker makes the Hugging Face commits. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable, Mapping +from pathlib import Path +from typing import Any + +from microcosm_emitter.host.contracts import JsonObject, ModuleContext +from microcosm_emitter.huggingface.service.delivery import ( + CREDENTIAL_MISSING, + HuggingFaceDelivery, +) +from microcosm_emitter.huggingface.service.hub import ( + DatasetRepository, + hugging_face_api, +) +from microcosm_emitter.huggingface.service.queue import HandoverQueue, sha256_bytes +from microcosm_emitter.huggingface.service.requests import ( + RequestRefusedError, + parse_stage_run_files, + run_reference, +) + +WORKER_ERROR_MESSAGE = ( + "Microcosm Hugging Face delivery hit a local queue error; handed-over files " + "stay queued for a later attempt." +) + + +def _refused_run_files_receipt(message: Mapping[str, Any], code: str) -> dict[str, Any]: + repository = message.get("repository") + run_id = message.get("run_id") + return { + "status": "refused", + "repository": repository if isinstance(repository, str) else "", + "run_id": run_id if isinstance(run_id, str) else "", + "pending_files": 0, + "files_attempted": 0, + "files_committed": 0, + "commit_attempts": 0, + "commits": 0, + "last_error_code": None, + "last_commit": None, + "refusal_code": code, + } + + +class HuggingFaceModule: + """Queue hand-overs on request; commit them on the worker.""" + + def __init__( + self, + configuration: JsonObject, + context: ModuleContext, + *, + api_factory: Callable[[], Any | None] = hugging_face_api, + clock: Callable[[], float] = time.time, + ) -> None: + self.configuration, self.context = configuration, context + self.api_factory, self.clock = api_factory, clock + self.queue: HandoverQueue | None = None + self.delivery: HuggingFaceDelivery | None = None + self._closed = False + + def initialize(self) -> None: + self.queue = HandoverQueue(Path(self.configuration["spool_path"])) + self.delivery = HuggingFaceDelivery( + self.queue, api_factory=self.api_factory, clock=self.clock + ) + + def handle_message(self, message: JsonObject) -> None: + raise ValueError("Hugging Face hand-overs are queries with receipts.") + + def query(self, message: JsonObject) -> dict[str, Any]: + assert self.queue is not None, "The host initializes modules first." + command = message.get("command") + if command == "stage_run_files": + return {"receipt": self._stage_run_files(message)} + if command == "run_files_receipt": + try: + repository, run_id = run_reference(message) + except RequestRefusedError as refusal: + return {"receipt": _refused_run_files_receipt(message, refusal.code)} + return { + "receipt": { + "status": "recorded", + **self.queue.run_receipt(repository, run_id), + } + } + if command == "read_back_run_files": + return {"read_back": self._read_back_run_files(message)} + raise ValueError("unsupported Hugging Face hand-over command") + + def _stage_run_files(self, message: JsonObject) -> dict[str, Any]: + assert self.queue is not None + try: + hand_over = parse_stage_run_files(message) + except RequestRefusedError as refusal: + return _refused_run_files_receipt(message, refusal.code) + self.queue.stage_run_files( + repository=hand_over.repository, + run_id=hand_over.run_id, + files=hand_over.files, + run_index_entry=hand_over.run_index_entry, + flush=hand_over.flush, + interval=hand_over.interval, + now=self.clock(), + ) + return { + "status": "recorded", + **self.queue.run_receipt(hand_over.repository, hand_over.run_id), + } + + def _read_back_run_files(self, message: JsonObject) -> dict[str, Any]: + assert self.queue is not None + result: dict[str, Any] = { + "status": "refused", + "files": 0, + "mismatched_paths": [], + "error_code": None, + } + try: + repository, run_id = run_reference(message) + except RequestRefusedError as refusal: + return {**result, "error_code": refusal.code} + if self.queue.run_receipt(repository, run_id)["pending_files"]: + return {**result, "status": "pending"} + api = self.api_factory() + if api is None: + return {**result, "status": "failed", "error_code": CREDENTIAL_MISSING} + committed = self.queue.committed_run_files(repository, run_id) + hub = DatasetRepository(repository, api) + mismatched = [] + try: + for path_in_repo, digest in sorted(committed.items()): + if sha256_bytes(hub.download(path_in_repo)) != digest: + mismatched.append(path_in_repo) + except Exception: + return { + **result, + "status": "failed", + "files": len(committed), + "error_code": "REPOSITORY_UNAVAILABLE", + } + return { + "status": "failed" if mismatched else "passed", + "files": len(committed), + "mismatched_paths": mismatched, + "error_code": None, + } + + def tick(self, now: float) -> None: + assert self.delivery is not None, "The host initializes modules first." + try: + self.delivery.run_cycle() + except Exception: + # The host stops a module whose tick raises; a local queue error + # leaves the work queued and must not end every later delivery. + self.delivery.warn_once("worker", WORKER_ERROR_MESSAGE) + + def parent_exited(self) -> None: + """Queued hand-overs stay eligible for the next emitter host.""" + + def close(self, deadline: float) -> None: + if self._closed: + return + self._closed = True + try: + if self.delivery is not None and time.monotonic() < deadline: + self.delivery.drain(deadline) + except Exception: + if self.delivery is not None: + self.delivery.warn_once("worker", WORKER_ERROR_MESSAGE) + finally: + if self.queue is not None: + self.queue.close() diff --git a/src/microcosm_emitter/huggingface/service/queue.py b/src/microcosm_emitter/huggingface/service/queue.py index ca1cd9f..d3d1458 100644 --- a/src/microcosm_emitter/huggingface/service/queue.py +++ b/src/microcosm_emitter/huggingface/service/queue.py @@ -18,6 +18,7 @@ from __future__ import annotations import hashlib +import json import time from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass @@ -51,6 +52,15 @@ def sha256_bytes(content: bytes) -> str: return hashlib.sha256(content).hexdigest() +def _entry_text(entry: Mapping[str, Any]) -> str: + """Order-preserving JSON, so the merged index keeps the build's key order.""" + return json.dumps(dict(entry), separators=(",", ":"), allow_nan=False) + + +def _entry(text: str | None) -> dict[str, Any]: + return {} if text is None else json.loads(text) + + def retry_delay(failures: int) -> float: return min(MAX_RETRY_SECONDS, float(2 ** min(failures, 8))) @@ -257,7 +267,7 @@ def stage_run_files( row.run_id, row.updated_at = run_id, now delivery = self._delivery(session, repository, run_id, now) if run_index_entry is not None: - entry = dict(run_index_entry) + entry = _entry_text(run_index_entry) if delivery.run_index_entry != entry or delivery.run_index_pending: delivery.run_index_due_at = ( min(delivery.run_index_due_at, due_at) @@ -281,6 +291,25 @@ def stage_run_files( delivery.run_index_due_at = min(delivery.run_index_due_at, now) delivery.updated_at = now + def repositories_with_pending_run_files(self) -> list[str]: + with self._sessions() as session: + return sorted( + set( + session.scalars( + select(StagedRunFile.repository).where( + StagedRunFile.pending.is_(True) + ) + ) + ) + | set( + session.scalars( + select(StagedRunDelivery.repository).where( + StagedRunDelivery.run_index_pending.is_(True) + ) + ) + ) + ) + def pending_run_files(self, repository: str) -> list[PendingRunFile]: with self._sessions() as session: rows = session.scalars( @@ -310,7 +339,7 @@ def pending_run_index_entries( ) .order_by(StagedRunDelivery.run_id) ) - return [(row.run_id, dict(row.run_index_entry or {})) for row in rows] + return [(row.run_id, _entry(row.run_index_entry)) for row in rows] @staticmethod def _files_per_run(entries: Iterable[PendingRunFile]) -> dict[str, int]: @@ -350,7 +379,7 @@ def record_run_files_committed( delivery.updated_at = now for run_id, entry in index_entries: delivery = self._delivery(session, repository, run_id, now) - if delivery.run_index_entry == dict(entry): + if delivery.run_index_entry == _entry_text(entry): delivery.run_index_pending = False schedule = self._schedule(session, repository) schedule.last_commit_at = now diff --git a/src/microcosm_emitter/huggingface/service/requests.py b/src/microcosm_emitter/huggingface/service/requests.py new file mode 100644 index 0000000..4b20631 --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/requests.py @@ -0,0 +1,143 @@ +"""Check hand-over requests with the interface's own rules, then read their files. + +A request that breaks a rule is refused as a whole with a fixed code; nothing +is queued, and no exception text is returned. +""" + +from __future__ import annotations + +import math +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from microcosm_emitter.compatibility import SUPPORTED_PROTOCOL_VERSIONS +from microcosm_emitter_interface.identifiers import ( + relative_file_path, + repository_identifier, + repository_prefix, + safe_identifier, +) + +#: Transport bounds for one staging run file and one hand-over. Microcosm's +#: staging contract keeps its own, smaller limits. +MAX_RUN_FILE_BYTES = 32 * 1024 * 1024 +MAX_HAND_OVER_BYTES = 64 * 1024 * 1024 +#: Staging contract version 1's pointer to the latest run, at the root. +LATEST_STAGING_POINTER_PATH = "latest_staging.json" + + +class RequestRefusedError(ValueError): + """A request the emitter will not queue; ``code`` says why.""" + + def __init__(self, code: str) -> None: + super().__init__(code) + self.code = code + + +@dataclass(frozen=True) +class RunFilesHandOver: + repository: str + run_id: str + files: list[tuple[str, bytes]] + run_index_entry: dict[str, Any] | None + flush: bool + interval: float + + +def require_supported_version(message: Mapping[str, Any]) -> None: + version = message.get("protocol_version") + if ( + isinstance(version, bool) + or not isinstance(version, int) + or version not in SUPPORTED_PROTOCOL_VERSIONS + ): + raise RequestRefusedError("UNSUPPORTED_PROTOCOL_VERSION") + + +def _valid(check, *arguments, **keywords): + try: + return check(*arguments, **keywords) + except ValueError: + raise RequestRefusedError("INVALID_REQUEST") from None + + +def run_reference(message: Mapping[str, Any]) -> tuple[str, str]: + require_supported_version(message) + return ( + _valid(repository_identifier, message.get("repository")), + _valid(safe_identifier, message.get("run_id"), label="run_id"), + ) + + +def _read_local_file(value: object) -> bytes: + if not isinstance(value, str) or not value: + raise RequestRefusedError("INVALID_REQUEST") + path = Path(value) + if not path.is_absolute(): + raise RequestRefusedError("INVALID_REQUEST") + try: + if not path.is_file(): + raise RequestRefusedError("LOCAL_FILE_MISSING") + if path.stat().st_size > MAX_RUN_FILE_BYTES: + raise RequestRefusedError("FILE_TOO_LARGE") + content = path.read_bytes() + except OSError: + raise RequestRefusedError("LOCAL_FILE_MISSING") from None + if len(content) > MAX_RUN_FILE_BYTES: + raise RequestRefusedError("FILE_TOO_LARGE") + return content + + +def parse_stage_run_files(message: Mapping[str, Any]) -> RunFilesHandOver: + """Validate a ``stage_run_files`` request and read every named file.""" + repository, run_id = run_reference(message) + prefix = _valid(repository_prefix, message.get("prefix")) + flush = message.get("flush") + interval = message.get("commit_interval_seconds") + if not isinstance(flush, bool) or ( + isinstance(interval, bool) + or not isinstance(interval, (int, float)) + or not math.isfinite(interval) + or interval < 0 + ): + raise RequestRefusedError("INVALID_REQUEST") + entries = message.get("files") + if not isinstance(entries, list): + raise RequestRefusedError("INVALID_REQUEST") + targets: list[tuple[str, object]] = [] + for entry in entries: + if not isinstance(entry, Mapping): + raise RequestRefusedError("INVALID_REQUEST") + relative = _valid( + relative_file_path, entry.get("relative_path"), label="relative_path" + ) + targets.append((f"{prefix}/{run_id}/{relative}", entry.get("local_path"))) + pointer = message.get("latest_pointer") + if pointer is not None: + targets.append((LATEST_STAGING_POINTER_PATH, pointer)) + if len({path_in_repo for path_in_repo, _ in targets}) != len(targets): + raise RequestRefusedError("INVALID_REQUEST") + run_index_entry = message.get("run_index_entry") + if run_index_entry is not None and ( + not isinstance(run_index_entry, Mapping) + or run_index_entry.get("run_id") != run_id + ): + raise RequestRefusedError("INVALID_REQUEST") + files = [] + total = 0 + for path_in_repo, local_path in targets: + content = _read_local_file(local_path) + total += len(content) + if total > MAX_HAND_OVER_BYTES: + raise RequestRefusedError("FILE_TOO_LARGE") + files.append((path_in_repo, content)) + return RunFilesHandOver( + repository=repository, + run_id=run_id, + files=files, + run_index_entry=None if run_index_entry is None else dict(run_index_entry), + flush=flush, + interval=float(interval), + ) diff --git a/src/microcosm_emitter/huggingface/service/run_files.py b/src/microcosm_emitter/huggingface/service/run_files.py new file mode 100644 index 0000000..31f67f7 --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/run_files.py @@ -0,0 +1,154 @@ +"""One commit carrying every pending staging run file of a repository. + +This is the transport that Microcosm's ``StagingRunBundleWriter`` (staging +contract version 1, ``staging.py``) and ``StagingRunBundleWriterV2`` +(``staging_v2.py``) performed with one ``upload_file`` commit per file. The +paths and bytes are those the build handed over. Version 1's ``runs.json`` +index is merged here, because the merge needs the repository's current copy. +""" + +from __future__ import annotations + +import json +import math +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Any + +from microcosm_emitter.huggingface.service.hub import ( + DatasetRepository, + add_operation, + is_absent, + is_parent_conflict, +) +from microcosm_emitter.huggingface.service.queue import PendingRunFile + +#: Staging contract version 1's index of runs, at the repository root. +RUNS_INDEX_PATH = "runs.json" +RUNS_INDEX_SCHEMA_VERSION = 1 +MAX_PARENT_CONFLICT_ATTEMPTS = 5 + + +def _jsonable(value: Any) -> Any: + """Microcosm ``staging._jsonable``: the same conversion before writing.""" + if isinstance(value, dict): + return {str(key): _jsonable(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_jsonable(item) for item in value] + if isinstance(value, Path): + return str(value) + if isinstance(value, float): + return value if math.isfinite(value) else None + if isinstance(value, (str, int, bool)) or value is None: + return value + item = getattr(value, "item", None) + if callable(item): + return _jsonable(item()) + return str(value) + + +def merge_runs_index( + existing: bytes | None, + entries: Sequence[Mapping[str, Any]], + *, + updated_at: str, +) -> bytes: + """Upsert each run's summary into ``runs.json`` as Microcosm did. + + Other runs are kept, each entry replaces its own run, and runs are sorted + newest first by ``updated_at``, then ``started_at``, then ``run_id``. The + bytes equal ``staging._write_json`` output: one-space indentation, no + trailing newline. + """ + runs: list[dict[str, Any]] = [] + if existing is not None: + try: + document = json.loads(existing) + except ValueError: + document = None + listed = document.get("runs") if isinstance(document, dict) else None + if isinstance(listed, list): + runs = [ + run for run in listed if isinstance(run, dict) and run.get("run_id") + ] + for entry in entries: + runs = [run for run in runs if run.get("run_id") != entry.get("run_id")] + runs.append(dict(entry)) + runs.sort( + key=lambda run: str( + run.get("updated_at") or run.get("started_at") or run.get("run_id") + ), + reverse=True, + ) + return json.dumps( + _jsonable( + { + "schema_version": RUNS_INDEX_SCHEMA_VERSION, + "updated_at": updated_at, + "runs": runs, + } + ), + indent=1, + allow_nan=False, + ).encode() + + +def commit_message( + entries: Sequence[PendingRunFile], + index_entries: Sequence[tuple[str, Mapping[str, Any]]], +) -> str: + run_ids = sorted( + {entry.run_id for entry in entries} | {r for r, _ in index_entries} + ) + return f"Stage run files for {', '.join(run_ids)}" + + +def commit_run_files( + repository: DatasetRepository, + entries: Sequence[PendingRunFile], + index_entries: Sequence[tuple[str, Mapping[str, Any]]], + *, + updated_at: str, +) -> str: + """Write the pending files (and merged index) in one commit; return it. + + With an index to merge, the commit is pinned to the head the index was + read at; when another writer commits first, the index is read again and + merged again, so neither writer's runs are lost. + """ + message = commit_message(entries, index_entries) + for attempt in range(MAX_PARENT_CONFLICT_ATTEMPTS): + operations = [ + add_operation(entry.path_in_repo, entry.content) for entry in entries + ] + parent = None + if index_entries: + parent = repository.head_revision() + try: + existing: bytes | None = repository.download( + RUNS_INDEX_PATH, revision=parent + ) + except Exception as error: + if not is_absent(error): + raise + existing = None + operations.append( + add_operation( + RUNS_INDEX_PATH, + merge_runs_index( + existing, + [entry for _, entry in index_entries], + updated_at=updated_at, + ), + ) + ) + try: + return repository.commit(operations, message=message, parent_commit=parent) + except Exception as error: + if ( + parent is None + or not is_parent_conflict(error) + or attempt + 1 == MAX_PARENT_CONFLICT_ATTEMPTS + ): + raise + raise AssertionError("unreachable") # pragma: no cover diff --git a/tests/huggingface/conftest.py b/tests/huggingface/conftest.py new file mode 100644 index 0000000..63542de --- /dev/null +++ b/tests/huggingface/conftest.py @@ -0,0 +1,61 @@ +"""A Hugging Face module driven by hand, against the in-memory Hub.""" + +import os +from types import SimpleNamespace + +import pytest +from fake_hugging_face import ( + PRIVATE_REPOSITORY, + STAGING_REPOSITORY, + FakeHuggingFaceApi, +) + + +class ManualClock: + def __init__(self, start: float = 1_000_000.0) -> None: + self.now = start + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +@pytest.fixture +def hub(): + api = FakeHuggingFaceApi( + whoami={"auth": {"accessToken": {"role": "write"}}, "name": "fixture"} + ) + api.create_repository(STAGING_REPOSITORY, {"README.md": b"staging fixture"}) + api.create_repository(PRIVATE_REPOSITORY, {"README.md": b"private fixture"}) + return api + + +@pytest.fixture +def hugging_face_module(tmp_path, hub): + """Start the module with an injected Hub; yields controls for the test.""" + from microcosm_emitter.host.contracts import ModuleContext + from microcosm_emitter.huggingface.service.module import HuggingFaceModule + + clock = ManualClock() + available_api = {"value": hub} + module = HuggingFaceModule( + {"spool_path": str(tmp_path / "events.sqlite3")}, + ModuleContext(os.getpid()), + api_factory=lambda: available_api["value"], + clock=clock, + ) + module.initialize() + controls = SimpleNamespace( + module=module, + clock=clock, + hub=hub, + tmp_path=tmp_path, + remove_credential=lambda: available_api.update(value=None), + restore_credential=lambda: available_api.update(value=hub), + ) + try: + yield controls + finally: + module.close(0) diff --git a/tests/huggingface/test_run_files.py b/tests/huggingface/test_run_files.py new file mode 100644 index 0000000..ae9687e --- /dev/null +++ b/tests/huggingface/test_run_files.py @@ -0,0 +1,409 @@ +"""Staging run files: Microcosm's paths and bytes, one commit per cycle.""" + +import json +import time +from datetime import UTC, datetime + +import pytest +from fake_hugging_face import STAGING_REPOSITORY, FakeHubHTTPError + +from microcosm_emitter.huggingface.service import requests as request_module +from microcosm_emitter_interface.staging import RunFile, StageRunFiles + + +def hand_over( + controls, + run_id, + files, + *, + flush=True, + interval=30.0, + latest_pointer=None, + run_index_entry=None, + prefix="runs", +): + directory = controls.tmp_path / "local" / run_id + run_files = [] + for relative_path, content in files.items(): + local = directory / relative_path + local.parent.mkdir(parents=True, exist_ok=True) + local.write_bytes(content) + run_files.append(RunFile(relative_path, local)) + pointer_path = None + if latest_pointer is not None: + pointer_path = directory / "latest_staging.json" + pointer_path.parent.mkdir(parents=True, exist_ok=True) + pointer_path.write_bytes(latest_pointer) + request = StageRunFiles( + repository=STAGING_REPOSITORY, + run_id=run_id, + files=tuple(run_files), + prefix=prefix, + latest_pointer=pointer_path, + run_index_entry=run_index_entry, + flush=flush, + commit_interval_seconds=interval, + ) + return controls.module.query(request.message())["receipt"] + + +def cycle(controls): + controls.module.delivery.run_cycle() + + +def commits(controls): + return controls.hub.commits(STAGING_REPOSITORY) + + +def iso(seconds): + return datetime.fromtimestamp(seconds, UTC).replace(microsecond=0).isoformat() + + +def test_hand_overs_within_one_cycle_reach_the_hub_as_one_commit( + hugging_face_module, +): + controls = hugging_face_module + hand_over(controls, "run-1", {"progress.json": b'{"stage":"created"}'}) + hand_over(controls, "run-1", {"run_manifest.json": b'{"run_id":"run-1"}'}) + hand_over( + controls, + "run-1", + {"progress.json": b'{"stage":"calibrating"}', "events.ndjson": b"{}\n"}, + ) + cycle(controls) + (commit,) = commits(controls) + assert commit.added == { + "runs/run-1/progress.json": b'{"stage":"calibrating"}', + "runs/run-1/run_manifest.json": b'{"run_id":"run-1"}', + "runs/run-1/events.ndjson": b"{}\n", + } + assert commit.message == "Stage run files for run-1" + + +def test_staging_v2_bundle_keeps_its_paths_and_bytes(hugging_face_module): + controls = hugging_face_module + files = { + "run_manifest.json": b'{\n "schema_name": "microcosm.staging.run-manifest"\n}\n', + "progress.json": b'{\n "schema_name": "microcosm.staging.progress"\n}\n', + "events.ndjson": b'{"sequence":1}\n{"sequence":2}\n', + "calibration_progress.json": b'{\n "events": []\n}\n', + "artifacts/fit_summary.json": b'{\n "final_loss": 0.5\n}\n', + } + receipt = hand_over(controls, "uk-frs-calibration-attempt-fixture", files) + assert receipt["status"] == "recorded" + assert receipt["pending_files"] == len(files) + cycle(controls) + (commit,) = commits(controls) + assert commit.added == { + f"runs/uk-frs-calibration-attempt-fixture/{name}": content + for name, content in files.items() + } + receipt = controls.module.query( + { + "protocol_version": 1, + "command": "run_files_receipt", + "repository": STAGING_REPOSITORY, + "run_id": "uk-frs-calibration-attempt-fixture", + } + )["receipt"] + assert receipt["pending_files"] == 0 + assert receipt["files_committed"] == len(files) + assert receipt["commits"] == 1 and receipt["last_commit"] == commit.oid + + +def test_unchanged_files_are_not_committed_again(hugging_face_module): + controls = hugging_face_module + hand_over(controls, "run-1", {"run_manifest.json": b"{}", "progress.json": b"1"}) + cycle(controls) + controls.clock.advance(60) + hand_over(controls, "run-1", {"run_manifest.json": b"{}", "progress.json": b"2"}) + cycle(controls) + assert [commit.added for commit in commits(controls)][-1] == { + "runs/run-1/progress.json": b"2" + } + + +def test_unforced_hand_overs_wait_for_the_interval_and_are_coalesced( + hugging_face_module, +): + controls = hugging_face_module + hand_over(controls, "run-1", {"progress.json": b"0"}) + cycle(controls) + for value in (b"1", b"2", b"3"): + controls.clock.advance(20) + hand_over( + controls, "run-1", {"progress.json": value}, flush=False, interval=300 + ) + cycle(controls) + assert len(commits(controls)) == 1 + controls.clock.advance(300) + cycle(controls) + assert [commit.added for commit in commits(controls)] == [ + {"runs/run-1/progress.json": b"0"}, + {"runs/run-1/progress.json": b"3"}, + ] + + +def test_staging_v1_pointer_and_runs_index_join_the_run_files_commit( + hugging_face_module, +): + controls = hugging_face_module + other_run = { + "run_id": "run-older", + "status": "passed", + "updated_at": "2026-10-01T00:00:00+00:00", + } + controls.hub.write_directly( + STAGING_REPOSITORY, + { + "runs.json": json.dumps( + {"schema_version": 1, "updated_at": "x", "runs": [other_run]}, indent=1 + ).encode() + }, + ) + this_run = { + "run_id": "us-run", + "candidate_release_id": "populace-us-fixture", + "status": "running", + "stage": "created", + "started_at": "2026-10-20T00:00:00+00:00", + "updated_at": "2026-10-20T00:00:00+00:00", + "progress_path": "runs/us-run/progress.json", + "run_manifest_path": "runs/us-run/run_manifest.json", + } + pointer = json.dumps({"schema_version": 1, "run_id": "us-run"}, indent=1).encode() + hand_over( + controls, + "us-run", + {"progress.json": b"{}"}, + latest_pointer=pointer, + run_index_entry=this_run, + ) + cycle(controls) + commit = commits(controls)[-1] + assert commit.added["runs/us-run/progress.json"] == b"{}" + assert commit.added["latest_staging.json"] == pointer + expected_index = json.dumps( + { + "schema_version": 1, + "updated_at": iso(controls.clock.now), + "runs": [this_run, other_run], + }, + indent=1, + allow_nan=False, + ).encode() + assert commit.added["runs.json"] == expected_index + (create_commit,) = [ + arguments + for method, arguments in controls.hub.calls + if method == "create_commit" + ] + assert create_commit["parent_commit"] is not None + + +def test_staging_v1_runs_index_starts_empty_without_a_repository_copy( + hugging_face_module, +): + controls = hugging_face_module + entry = {"run_id": "us-run", "updated_at": "2026-10-20T00:00:00+00:00"} + hand_over(controls, "us-run", {}, run_index_entry=entry) + cycle(controls) + (commit,) = commits(controls) + assert json.loads(commit.added["runs.json"])["runs"] == [entry] + + +def test_runs_index_written_by_another_build_meanwhile_is_kept(hugging_face_module): + controls = hugging_face_module + concurrent_run = {"run_id": "concurrent", "updated_at": "2026-10-19T00:00:00Z"} + controls.hub.run_before_next_commit( + lambda: controls.hub.write_directly( + STAGING_REPOSITORY, + { + "runs.json": json.dumps( + {"schema_version": 1, "runs": [concurrent_run]} + ).encode() + }, + message="Another build's index", + ) + ) + entry = {"run_id": "us-run", "updated_at": "2026-10-20T00:00:00+00:00"} + hand_over(controls, "us-run", {"progress.json": b"{}"}, run_index_entry=entry) + cycle(controls) + ours = [ + commit + for commit in commits(controls) + if commit.message == "Stage run files for us-run" + ] + assert len(ours) == 1 + assert json.loads(ours[0].added["runs.json"])["runs"] == [entry, concurrent_run] + + +def test_a_failed_commit_keeps_the_files_queued_and_retries_after_backoff( + hugging_face_module, +): + controls = hugging_face_module + controls.hub.fail_next("create_commit", FakeHubHTTPError(503)) + hand_over(controls, "run-1", {"progress.json": b"1"}) + cycle(controls) + assert commits(controls) == [] + receipt = controls.module.queue.run_receipt(STAGING_REPOSITORY, "run-1") + assert receipt["last_error_code"] == "UPLOAD_FAILED" + assert receipt["pending_files"] == 1 + cycle(controls) + assert commits(controls) == [] + controls.clock.advance(2) + cycle(controls) + assert len(commits(controls)) == 1 + receipt = controls.module.queue.run_receipt(STAGING_REPOSITORY, "run-1") + assert receipt["last_error_code"] is None + assert receipt["commit_attempts"] == 2 and receipt["commits"] == 1 + + +def test_worker_errors_never_stop_the_module(hugging_face_module, capsys): + controls = hugging_face_module + controls.hub.fail_next("create_commit", RuntimeError("synthetic")) + hand_over(controls, "run-1", {"progress.json": b"1"}) + controls.module.tick(time.monotonic()) + controls.module.tick(time.monotonic()) + assert commits(controls) == [] + assert "synthetic" not in capsys.readouterr().err + + +def test_without_a_credential_files_stay_queued_and_one_warning_is_printed( + hugging_face_module, capsys +): + controls = hugging_face_module + controls.remove_credential() + hand_over(controls, "run-1", {"progress.json": b"1"}) + cycle(controls) + controls.clock.advance(5) + cycle(controls) + assert commits(controls) == [] + receipt = controls.module.queue.run_receipt(STAGING_REPOSITORY, "run-1") + assert receipt["last_error_code"] == "CREDENTIAL_MISSING" + assert capsys.readouterr().err.count("no Hugging Face credential") == 1 + controls.restore_credential() + controls.clock.advance(300) + cycle(controls) + assert len(commits(controls)) == 1 + + +def test_close_commits_what_is_pending_without_waiting_for_the_interval( + hugging_face_module, +): + controls = hugging_face_module + hand_over(controls, "run-1", {"progress.json": b"1"}) + cycle(controls) + controls.clock.advance(1) + hand_over(controls, "run-1", {"progress.json": b"2"}, flush=False, interval=300) + controls.module.close(time.monotonic() + 5) + assert [commit.added for commit in commits(controls)][-1] == { + "runs/run-1/progress.json": b"2" + } + + +def test_read_back_compares_the_hub_with_the_committed_bytes(hugging_face_module): + controls = hugging_face_module + query = { + "protocol_version": 1, + "command": "read_back_run_files", + "repository": STAGING_REPOSITORY, + "run_id": "run-1", + } + hand_over(controls, "run-1", {"progress.json": b"1", "run_manifest.json": b"{}"}) + assert controls.module.query(query)["read_back"]["status"] == "pending" + cycle(controls) + assert controls.module.query(query)["read_back"] == { + "status": "passed", + "files": 2, + "mismatched_paths": [], + "error_code": None, + } + controls.hub.write_directly( + STAGING_REPOSITORY, {"runs/run-1/progress.json": b"changed"} + ) + assert controls.module.query(query)["read_back"] == { + "status": "failed", + "files": 2, + "mismatched_paths": ["runs/run-1/progress.json"], + "error_code": None, + } + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("relative_path", "../escape.json"), + ("relative_path", "/absolute.json"), + ("relative_path", "nested/../../escape.json"), + ("relative_path", "back\\slash.json"), + ("relative_path", ""), + ("prefix", "releases"), + ("prefix", "../runs"), + ("prefix", ""), + ("run_id", "../run"), + ("repository", "no-owner"), + ("commit_interval_seconds", -1), + ("flush", "yes"), + ], +) +def test_requests_that_could_leave_the_run_folder_are_refused( + hugging_face_module, field, value +): + controls = hugging_face_module + local = controls.tmp_path / "progress.json" + local.write_bytes(b"{}") + message = StageRunFiles( + repository=STAGING_REPOSITORY, + run_id="run-1", + files=(RunFile("progress.json", local),), + ).message() + if field == "relative_path": + message["files"][0]["relative_path"] = value + else: + message[field] = value + receipt = controls.module.query(message)["receipt"] + assert receipt["status"] == "refused" + assert receipt["refusal_code"] == "INVALID_REQUEST" + assert controls.module.queue.pending_run_files(STAGING_REPOSITORY) == [] + + +def test_a_missing_local_file_refuses_the_whole_hand_over(hugging_face_module): + controls = hugging_face_module + present = controls.tmp_path / "present.json" + present.write_bytes(b"{}") + message = StageRunFiles( + repository=STAGING_REPOSITORY, + run_id="run-1", + files=( + RunFile("present.json", present), + RunFile("missing.json", controls.tmp_path / "missing.json"), + ), + ).message() + receipt = controls.module.query(message)["receipt"] + assert (receipt["status"], receipt["refusal_code"]) == ( + "refused", + "LOCAL_FILE_MISSING", + ) + assert controls.module.queue.pending_run_files(STAGING_REPOSITORY) == [] + + +def test_files_above_the_transport_limit_are_refused(hugging_face_module, monkeypatch): + controls = hugging_face_module + monkeypatch.setattr(request_module, "MAX_RUN_FILE_BYTES", 4) + receipt = hand_over(controls, "run-1", {"progress.json": b"12345"}) + assert (receipt["status"], receipt["refusal_code"]) == ("refused", "FILE_TOO_LARGE") + + +def test_unsupported_protocol_versions_are_refused(hugging_face_module): + controls = hugging_face_module + local = controls.tmp_path / "progress.json" + local.write_bytes(b"{}") + message = StageRunFiles( + repository=STAGING_REPOSITORY, + run_id="run-1", + files=(RunFile("progress.json", local),), + ).message() + message["protocol_version"] = 99 + receipt = controls.module.query(message)["receipt"] + assert receipt["refusal_code"] == "UNSUPPORTED_PROTOCOL_VERSION" diff --git a/tests/interface/test_client_messages.py b/tests/interface/test_client_messages.py index 71680e3..afe3d52 100644 --- a/tests/interface/test_client_messages.py +++ b/tests/interface/test_client_messages.py @@ -14,8 +14,15 @@ materialize, protocol_fixtures, ) +from microcosm_emitter_interface.staging import ( + HuggingFaceSettings, + RunFile, + StageRunFiles, +) from microcosm_emitter_interface.telemetry import TelemetryRun, TelemetrySettings +FIXTURE_STAGING_REPOSITORY = "policyengine/fixture-staging" + class RecordingTransport: """Records requests and answers queries like an emitter that queued nothing.""" @@ -28,6 +35,16 @@ def send(self, module, message, *, timeout=None): def query(self, module, message, *, timeout=None): self.requests.append(("query", module, dict(message))) + if module == "huggingface": + return { + "receipt": {"status": "recorded", "pending_files": 0}, + "read_back": { + "status": "passed", + "files": 0, + "mismatched_paths": [], + "error_code": None, + }, + } if message.get("command") == "inventory": return {"inventory": None} return { @@ -52,7 +69,10 @@ def recording_client(): process=None, transport=transport, runtime_directory=Path("/nonexistent"), - description={"modules": ["orrery", "telemetry"], "protocol_versions": [1]}, + description={ + "modules": ["huggingface", "orrery", "telemetry"], + "protocol_versions": [1], + }, ) return EmitterClient(process), transport @@ -130,6 +150,37 @@ def fixtures_by_name() -> dict[str, ProtocolFixture]: "orrery-retry": lambda client, directory: client.retry_orrery_inputs( "fixture-publication" ), + "huggingface-stage-run-files": lambda client, directory: client.stage_run_files( + StageRunFiles( + repository=FIXTURE_STAGING_REPOSITORY, + run_id="fixture-run", + files=( + RunFile("progress.json", directory / "progress.json"), + RunFile("artifacts/fit_summary.json", directory / "fit_summary.json"), + ), + flush=True, + ) + ), + "huggingface-stage-run-files-with-index": lambda client, directory: ( + client.stage_run_files( + StageRunFiles( + repository=FIXTURE_STAGING_REPOSITORY, + run_id="fixture-run", + files=(RunFile("progress.json", directory / "progress.json"),), + latest_pointer=directory / "latest_staging.json", + run_index_entry={"run_id": "fixture-run", "status": "running"}, + commit_interval_seconds=300, + ) + ) + ), + "huggingface-run-files-receipt": lambda client, directory: client.run_files_receipt( + FIXTURE_STAGING_REPOSITORY, "fixture-run" + ), + "huggingface-read-back-run-files": lambda client, directory: ( + client.read_back_run_files( + FIXTURE_STAGING_REPOSITORY, "fixture-run", wait_seconds=0 + ) + ), } CONFIGURED_BY: dict[str, Callable[[Path], dict[str, Any]]] = { @@ -142,6 +193,9 @@ def fixtures_by_name() -> dict[str, ProtocolFixture]: development_collector_url="http://127.0.0.1:1", publication_origin="http://127.0.0.1:1", ).configuration(spool_path), + "huggingface-configuration": lambda spool_path: HuggingFaceSettings().configuration( + spool_path + ), } @@ -169,7 +223,9 @@ def test_typed_call_sends_the_recorded_message(recording_client, tmp_path, name) directory = tmp_path / "inputs" SENT_BY[name](client, directory) expected = materialize(fixture, directory=directory.resolve(), spool_path=tmp_path) - assert transport.requests[0] == (fixture.kind, fixture.module, expected) + # A call may send other requests first (read-back first commits pending + # files), but it must send the recorded message. + assert (fixture.kind, fixture.module, expected) in transport.requests @pytest.mark.parametrize("name", sorted(CONFIGURED_BY)) diff --git a/tests/interface/test_protocol_fixtures.py b/tests/interface/test_protocol_fixtures.py index 01480ff..a863a00 100644 --- a/tests/interface/test_protocol_fixtures.py +++ b/tests/interface/test_protocol_fixtures.py @@ -16,10 +16,12 @@ from pathlib import Path import pytest +from fake_hugging_face import FakeHuggingFaceApi from microcosm_emitter.compatibility import SUPPORTED_PROTOCOL_VERSIONS from microcosm_emitter.host.contracts import ModuleContext from microcosm_emitter.host.runtime import ServiceRuntime +from microcosm_emitter.huggingface.service.module import HuggingFaceModule from microcosm_emitter.orrery.service.module import OrreryModule from microcosm_emitter.telemetry.service.module import TelemetryModule from microcosm_emitter_interface.launcher import DEFAULT_COMMAND, start_emitter @@ -35,6 +37,21 @@ "protocol-v1/host-close.json": ( "68fa81b4116852f91f3c0dab3716a1974a8150039fb2b238f79efc094d9b0f6f" ), + "protocol-v1/huggingface-configuration.json": ( + "5a5ccd639953b24c774418e6e23735f515e46a673a68624e26d640839ef4b1c2" + ), + "protocol-v1/huggingface-read-back-run-files.json": ( + "ab643a81520cbf35fada4362d6bf08f63c84d45ed41d9ddf1a0a592855c91313" + ), + "protocol-v1/huggingface-run-files-receipt.json": ( + "acc37a2747fff6dc5ec134410f3c068b0cfb830e8703cdb273a83c690971315e" + ), + "protocol-v1/huggingface-stage-run-files-with-index.json": ( + "efe20fc318c6a576a37684884aff3497ac9aad093f4df81ca9b34a4ecbc0c0e7" + ), + "protocol-v1/huggingface-stage-run-files.json": ( + "e29d312aed140076789ab4f4c38bf7a9d11eadf1b446ea96d46a6f88ca7e4bd0" + ), "protocol-v1/host-hello.json": ( "36fcac377b69a9103b7c1cfd884d2ff667e030463edc9da349a379bb4868b927" ), @@ -175,6 +192,33 @@ def test_orrery_module_answers_each_query_fixture_in_order(tmp_path): module.close(time.monotonic()) +def test_hugging_face_module_answers_each_query_fixture_in_order(tmp_path): + hub = FakeHuggingFaceApi() + hub.create_repository("policyengine/fixture-staging") + module = HuggingFaceModule( + configuration_of("huggingface", tmp_path), + ModuleContext(os.getpid()), + api_factory=lambda: hub, + ) + module.initialize() + directory = tmp_path / "inputs" + directory.mkdir() + try: + for fixture in fixtures_of("huggingface", "query"): + for name, text in fixture.files.items(): + (directory / name).write_text(text) + response = module.query( + materialize(fixture, directory=directory, spool_path=tmp_path) + ) + assert set(fixture.response_keys) <= set(response), fixture.name + if "receipt" in response: + assert response["receipt"]["status"] == "recorded", fixture.name + if "read_back" in response: + assert response["read_back"]["status"] in {"pending", "passed"} + finally: + module.close(time.monotonic()) + + def test_installed_command_answers_the_hello_fixture(tmp_path): """The answer has at least the recorded keys; a newer emitter may add more.""" (fixture,) = fixtures_of("host", "hello") diff --git a/tests/interface/test_staging_client.py b/tests/interface/test_staging_client.py new file mode 100644 index 0000000..99024a8 --- /dev/null +++ b/tests/interface/test_staging_client.py @@ -0,0 +1,205 @@ +"""Typed staging-run-file hand-overs, through a real socket to the module.""" + +import os +import shutil +import tempfile +import threading +import time +from dataclasses import dataclass, field +from pathlib import Path + +import pytest +from fake_hugging_face import STAGING_REPOSITORY, FakeHuggingFaceApi + +from microcosm_emitter.host.contracts import ModuleContext +from microcosm_emitter.host.runtime import ServiceRuntime +from microcosm_emitter.huggingface.service.module import HuggingFaceModule +from microcosm_emitter_interface.client import EmitterClient +from microcosm_emitter_interface.staging import RunFile, StageRunFiles +from microcosm_emitter_interface.transport import SocketTransport + + +@dataclass +class InThreadEmitter: + """An emitter host running in a thread of the test process.""" + + transport: SocketTransport + description: dict = field(default_factory=dict) + + def close(self): + self.transport.close() + + +@pytest.fixture +def emitter_with_hub(tmp_path): + hub = FakeHuggingFaceApi() + hub.create_repository(STAGING_REPOSITORY) + runtime_directory = Path(tempfile.mkdtemp(prefix="pe-stage-", dir="/tmp")) + socket_path = runtime_directory / "service.sock" + module = HuggingFaceModule( + {"spool_path": str(tmp_path / "events.sqlite3")}, + ModuleContext(os.getpid()), + api_factory=lambda: hub, + ) + runtime = ServiceRuntime( + socket_path, + modules={"huggingface": module}, + is_parent_alive=lambda: True, + tick_seconds=0.01, + ) + thread = threading.Thread(target=runtime.run) + thread.start() + transport = SocketTransport(socket_path) + deadline = time.monotonic() + 5 + while True: + try: + transport.ping() + break + except Exception: + if time.monotonic() > deadline: + raise + time.sleep(0.01) + client = EmitterClient( + InThreadEmitter( + transport, {"modules": ["huggingface"], "protocol_versions": [1]} + ) + ) + try: + yield client, hub + finally: + client.close() + thread.join(timeout=20) + shutil.rmtree(runtime_directory, ignore_errors=True) + + +@pytest.mark.parametrize( + "build_request", + [ + lambda local: StageRunFiles( + repository=STAGING_REPOSITORY, + run_id="run-1", + files=(RunFile("../escape.json", local),), + ), + lambda local: StageRunFiles( + repository=STAGING_REPOSITORY, + run_id="run-1", + files=(RunFile("/absolute.json", local),), + ), + lambda local: StageRunFiles( + repository=STAGING_REPOSITORY, run_id="run-1", prefix="releases" + ), + lambda local: StageRunFiles(repository=STAGING_REPOSITORY, run_id="../run"), + lambda local: StageRunFiles(repository="no-owner", run_id="run-1"), + lambda local: StageRunFiles( + repository=STAGING_REPOSITORY, run_id="run-1", commit_interval_seconds=-1 + ), + lambda local: StageRunFiles( + repository=STAGING_REPOSITORY, + run_id="run-1", + run_index_entry={"run_id": "another-run"}, + ), + ], + ids=[ + "parent-directory", + "absolute", + "release-records", + "run-identifier", + "repository", + "interval", + "index-entry-of-another-run", + ], +) +def test_requests_that_could_leave_the_run_folder_cannot_be_built( + tmp_path, build_request +): + with pytest.raises(ValueError): + build_request(tmp_path / "progress.json") + + +def test_message_names_each_file_by_its_absolute_local_path(tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + request = StageRunFiles( + repository=STAGING_REPOSITORY, + run_id="run-1", + files=(RunFile("artifacts/fit_summary.json", Path("fit.json")),), + flush=True, + commit_interval_seconds=300, + ) + assert request.message() == { + "protocol_version": 1, + "command": "stage_run_files", + "repository": STAGING_REPOSITORY, + "run_id": "run-1", + "prefix": "runs", + "files": [ + { + "relative_path": "artifacts/fit_summary.json", + "local_path": str((tmp_path / "fit.json").resolve()), + } + ], + "latest_pointer": None, + "run_index_entry": None, + "flush": True, + "commit_interval_seconds": 300.0, + } + + +def test_client_without_the_hugging_face_module_reports_unavailable(tmp_path): + client = EmitterClient( + InThreadEmitter( + SocketTransport(tmp_path / "unused.sock"), + {"modules": ["telemetry"], "protocol_versions": [1]}, + ) + ) + receipt = client.stage_run_files( + StageRunFiles(repository=STAGING_REPOSITORY, run_id="run-1") + ) + assert receipt["status"] == "unavailable" + assert client.read_back_run_files(STAGING_REPOSITORY, "run-1", wait_seconds=0)[ + "status" + ] == ("unavailable") + + +def test_unreachable_emitter_never_raises(tmp_path): + client = EmitterClient( + InThreadEmitter( + SocketTransport(tmp_path / "nothing-listens.sock"), + {"modules": ["huggingface"], "protocol_versions": [1]}, + ) + ) + receipt = client.stage_run_files( + StageRunFiles(repository=STAGING_REPOSITORY, run_id="run-1", flush=True) + ) + assert receipt["status"] == "unavailable" + + +def test_hand_over_reaches_the_hub_in_one_commit_and_reads_back( + tmp_path, emitter_with_hub +): + client, hub = emitter_with_hub + progress = tmp_path / "progress.json" + progress.write_bytes(b'{"stage":"complete"}') + manifest = tmp_path / "run_manifest.json" + manifest.write_bytes(b'{"run_id":"run-1"}') + receipt = client.stage_run_files( + StageRunFiles( + repository=STAGING_REPOSITORY, + run_id="run-1", + files=( + RunFile("progress.json", progress), + RunFile("run_manifest.json", manifest), + ), + flush=True, + ) + ) + assert receipt["status"] == "recorded" + read_back = client.read_back_run_files(STAGING_REPOSITORY, "run-1", wait_seconds=20) + assert read_back["status"] == "passed", read_back + assert read_back["files"] == 2 + (commit,) = hub.commits(STAGING_REPOSITORY) + assert commit.added == { + "runs/run-1/progress.json": b'{"stage":"complete"}', + "runs/run-1/run_manifest.json": b'{"run_id":"run-1"}', + } + receipt = client.run_files_receipt(STAGING_REPOSITORY, "run-1") + assert receipt["commits"] == 1 and receipt["last_commit"] == commit.oid diff --git a/tests/support/fake_hugging_face.py b/tests/support/fake_hugging_face.py index 14c0af4..a80ad07 100644 --- a/tests/support/fake_hugging_face.py +++ b/tests/support/fake_hugging_face.py @@ -19,6 +19,9 @@ from types import SimpleNamespace from typing import Any +STAGING_REPOSITORY = "policyengine/fixture-staging" +PRIVATE_REPOSITORY = "policyengine/fixture-private" + class FakeHubHTTPError(OSError): """Carries ``response.status_code`` like ``huggingface_hub.HfHubHTTPError``.""" diff --git a/tests/test_package_names.py b/tests/test_package_names.py index 7cbe974..015af11 100644 --- a/tests/test_package_names.py +++ b/tests/test_package_names.py @@ -68,6 +68,13 @@ def test_installed_distribution_includes_host_telemetry_and_migrations(): graphs = list(entry_points(group="microcosm.emitter.modules", name="orrery")) assert len(graphs) == 1 assert graphs[0].value == "microcosm_emitter.orrery.service:create_module" + hugging_face = list( + entry_points(group="microcosm.emitter.modules", name="huggingface") + ) + assert len(hugging_face) == 1 + assert ( + hugging_face[0].value == "microcosm_emitter.huggingface.service:create_module" + ) assert find_spec("microcosm_emitter.orrery.client") is not None assert find_spec("microcosm_emitter.host.__main__") is not None assert find_spec("microcosm_emitter.telemetry.client") is not None diff --git a/tools/installed-interface-smoke.py b/tools/installed-interface-smoke.py index 0e2b9ad..edb5ffd 100644 --- a/tools/installed-interface-smoke.py +++ b/tools/installed-interface-smoke.py @@ -12,7 +12,10 @@ def main(emitter_command: str) -> None: assert not interface.requires, interface.requires from microcosm_emitter_interface import ( EmitterClient, + HuggingFaceSettings, OrrerySettings, + RunFile, + StageRunFiles, TelemetryRun, TelemetrySettings, ) @@ -31,6 +34,9 @@ def main(emitter_command: str) -> None: development_collector_url="http://127.0.0.1:1", publication_origin="http://127.0.0.1:1", ), + # Without a credential (removed by check-installed.py) the module + # queues hand-overs and never contacts the Hub. + hugging_face=HuggingFaceSettings(), spool_path=Path.cwd() / "interface.sqlite3", command=emitter_command, # A clean install compiles dependency bytecode on first use. @@ -38,7 +44,7 @@ def main(emitter_command: str) -> None: ) assert client.available, client.startup_error try: - assert client.modules == {"telemetry", "orrery"} + assert client.modules == {"telemetry", "orrery", "huggingface"} assert client.emitter_version == interface.version assert client.telemetry.transition_stage("compile") assert client.telemetry.complete() @@ -46,6 +52,17 @@ def main(emitter_command: str) -> None: assert inventory is not None receipt = client.hand_over_orrery_inputs(directory, inventory, wait_seconds=0) assert receipt["status"] == "pending", receipt + progress = Path.cwd() / "progress.json" + progress.write_text('{"stage":"created"}') + staged = client.stage_run_files( + StageRunFiles( + repository="policyengine/installed-check", + run_id="interface-wheel", + files=(RunFile("progress.json", progress),), + flush=True, + ) + ) + assert staged["status"] == "recorded" and staged["pending_files"] == 1 finally: process = client._process client.close() From 4d5204e581b324c63cda5a4857fdc06802c4efad Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 01:04:06 +0400 Subject: [PATCH 03/23] Stage dataset bundles with each file's role through the queue The transport of Microcosm's staging_dataset.stage_bundle (used by uk_runtime/rowwise_staging.py and the Orrery evidence bundle) moves to the emitter as a durable job: - `StageDatasetBundle` names each file with its role (dataset, calibration_diagnostics, build_manifest, staged_manifest, checksums, ...), SHA-256 and size; exactly one staged_manifest is required. - The worker re-hashes preserved files and refuses changed bytes, then writes every file under `//` in one commit pinned to the head it checked, with Microcosm's commit message. - A matching repository inventory is `already_staged` with its own commit; a different one is `REMOTE_DIFFERS` and is never overwritten. - REPOSITORY_UNAVAILABLE and UPLOAD_FAILED are retried with backoff; a commit that loses a race is checked and tried again. - `check_write_access` ports the staged-dataset preflight. - Jobs are filtered by the kinds this emitter delivers, so a job queued by a newer emitter sharing the database waits for it. Co-Authored-By: Claude Opus 5.5 --- README.md | 9 + .../huggingface-dataset-bundles.added.md | 1 + .../microcosm-emitter-interface/README.md | 11 +- .../microcosm_emitter_interface/__init__.py | 3 + .../microcosm_emitter_interface/bundles.py | 216 +++++++++ .../src/microcosm_emitter_interface/client.py | 81 ++++ .../huggingface-check-write-access.json | 14 + .../huggingface-handover-receipt.json | 14 + .../huggingface-stage-dataset-bundle.json | 61 +++ .../microcosm_emitter_interface/protocol.py | 6 +- .../huggingface/service/bundles.py | 136 ++++++ .../huggingface/service/delivery.py | 85 +++- .../huggingface/service/module.py | 77 ++++ .../huggingface/service/queue.py | 76 +++- .../huggingface/service/requests.py | 64 +++ tests/huggingface/test_bundles.py | 409 ++++++++++++++++++ tests/interface/conftest.py | 76 ++++ tests/interface/test_bundle_client.py | 156 +++++++ tests/interface/test_client_messages.py | 44 +- tests/interface/test_protocol_fixtures.py | 20 +- tests/interface/test_staging_client.py | 74 +--- 21 files changed, 1535 insertions(+), 98 deletions(-) create mode 100644 changelog.d/huggingface-dataset-bundles.added.md create mode 100644 packages/microcosm-emitter-interface/src/microcosm_emitter_interface/bundles.py create mode 100644 packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-check-write-access.json create mode 100644 packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-handover-receipt.json create mode 100644 packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-stage-dataset-bundle.json create mode 100644 src/microcosm_emitter/huggingface/service/bundles.py create mode 100644 tests/huggingface/test_bundles.py create mode 100644 tests/interface/test_bundle_client.py diff --git a/README.md b/README.md index 3bab13a..b79f8cb 100644 --- a/README.md +++ b/README.md @@ -67,6 +67,15 @@ The `huggingface` module takes over the uploads Microcosm made itself. Microcosm - A commit that merges `runs.json` is pinned to the head it read (`parent_commit`); if another build commits first, the module reads the index again and merges again. - `read_back_run_files` commits pending files, then compares the Hub's copy of each committed file with the committed bytes. +**Staged dataset bundles** (Microcosm's `staging_dataset.stage_bundle`, used by `uk_runtime/rowwise_staging.py` and the Orrery evidence bundle). A build verifies its bundle and writes the sidecars as before, then hands the files over by path, each with its role (`dataset`, `calibration_diagnostics`, `build_manifest`, `staged_manifest`, `checksums`, ...), SHA-256 and size. The hand-over is a durable job: + +- The worker re-hashes each preserved file and never uploads changed bytes (`PRESERVED_BYTES_CHANGED`); a missing file waits for an operator to restore it (`PRESERVED_FILES_MISSING`). +- As in Microcosm, it writes every file under `//` (by default `staged//`) in one commit pinned to the head it checked, with the message `Stage under /`. A repository copy of the bundle's `staged_manifest` whose `files` equal the outputs' digests is `already_staged` with that copy's commit; a different one is `REMOTE_DIFFERS` and is never overwritten. +- Unlike Microcosm's single attempt, `REPOSITORY_UNAVAILABLE` and `UPLOAD_FAILED` are retried with backoff, and a commit that loses a race to another writer is checked and tried again. The receipt keeps the last error code while the job is pending. +- `check_write_access` is Microcosm's staged-dataset preflight: whether a credential is present, the repository is reachable, and the Hub reports a write scope for it. + +A repository with both a due bundle and due run files gets the bundle in one cycle and the run files in a later one, so each cycle still makes one commit per repository. + ## Development Use Python 3.13 or 3.14 on Linux or macOS. diff --git a/changelog.d/huggingface-dataset-bundles.added.md b/changelog.d/huggingface-dataset-bundles.added.md new file mode 100644 index 0000000..723e519 --- /dev/null +++ b/changelog.d/huggingface-dataset-bundles.added.md @@ -0,0 +1 @@ +Add `stage_dataset_bundle`, `dataset_bundle_receipt` and `check_write_access`. A build hands over its staged dataset bundle with each file's role; the emitter writes it under `staged//` in one commit with Microcosm's message, paths and idempotency (`already_staged`, `REMOTE_DIFFERS`), retries Hub failures from the durable queue, and never uploads bytes that changed after the hand-over. diff --git a/packages/microcosm-emitter-interface/README.md b/packages/microcosm-emitter-interface/README.md index 210901f..6c4f027 100644 --- a/packages/microcosm-emitter-interface/README.md +++ b/packages/microcosm-emitter-interface/README.md @@ -23,7 +23,16 @@ client.close() `EmitterClient.start` runs the `microcosm-emitter` command as the build's own child process. The command comes from the `command` argument, then the `MICROCOSM_EMITTER_COMMAND` environment variable, then `PATH`. The child inherits the build's environment variables and user, so it finds the same Hugging Face credential: `HF_TOKEN`, then `HUGGINGFACE_TOKEN`, then the token saved by `hf auth login`. -No call raises into the build: not because the emitter is missing, slow or refusing a request, and not because an argument cannot be converted to JSON (for example a multi-element NumPy array or a self-referencing mapping). Each call converts its arguments inside its error handling and reports whether the emitter accepted it; after a refusal, `client.last_failure` says why (`emitter_unavailable`, `invalid_arguments` or `not_accepted`). Startup settings that cannot be converted start an unavailable client (`invalid_configuration`). The build continues in every case. +No call raises into the build: not because the emitter is missing, slow or refusing a request, and not because an argument cannot be converted to JSON (for example a multi-element NumPy array or a self-referencing mapping). Each call converts its arguments inside its error handling and reports whether the emitter accepted it; after a refusal, `client.last_failure` says why (`emitter_unavailable`, `invalid_arguments` or `not_accepted`). Startup settings that cannot be converted start an unavailable client (`invalid_configuration`). The build continues in every case. Constructing a request object (for example `StageRunFiles(...)`) that breaks the shared path and identifier rules raises `ValueError` in the build's own code, before any call. + +## Calls + +| Hand-over | Call | Emitter module | +|---|---|---| +| Telemetry event | `client.telemetry.stage(...)`, `transition_stage`, `progress`, `calibration_progress`, `fail`, `refuse`, `complete`, `emit` (each one `emit_telemetry`) | `telemetry` | +| Staging run files under `runs//` | `stage_run_files(StageRunFiles(...))`, `run_files_receipt`, `read_back_run_files` | `huggingface` | +| Staged dataset bundle, each file with its role | `stage_dataset_bundle(StageDatasetBundle(...))`, `dataset_bundle_receipt`, `check_write_access` | `huggingface` | +| Orrery inputs | `orrery_inventory`, `hand_over_orrery_inputs`, `orrery_receipt`, `retry_orrery_inputs` | `orrery` | ## Compatibility diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/__init__.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/__init__.py index ccb71f1..f9e7526 100644 --- a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/__init__.py +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/__init__.py @@ -3,6 +3,7 @@ Importing this package starts no process and opens no socket or file. """ +from microcosm_emitter_interface.bundles import BundleFile, StageDatasetBundle from microcosm_emitter_interface.client import EmitterClient, RequestFailure from microcosm_emitter_interface.launcher import ( COMMAND_ENVIRONMENT_VARIABLE, @@ -20,6 +21,7 @@ from microcosm_emitter_interface.telemetry import TelemetryRun, TelemetrySettings __all__ = [ + "BundleFile", "COMMAND_ENVIRONMENT_VARIABLE", "DEFAULT_COMMAND", "PROTOCOL_VERSION", @@ -29,6 +31,7 @@ "OrrerySettings", "RequestFailure", "RunFile", + "StageDatasetBundle", "StageRunFiles", "TelemetryRun", "TelemetrySettings", diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/bundles.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/bundles.py new file mode 100644 index 0000000..f679593 --- /dev/null +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/bundles.py @@ -0,0 +1,216 @@ +"""Typed staged-dataset-bundle hand-overs, with each file's role named. + +The build verifies its bundle against its manifest and writes the sidecars +(``staged_manifest.json``, ``sha256sums.txt``) as before, then hands the files +over by path with their roles and digests. The emitter writes them under +``//`` in one commit, and is idempotent on the outputs' +digests: a repository copy of the bundle inventory (role ``staged_manifest``) +whose ``files`` equal the outputs' digests is ``already_staged``, and a +different one is refused as ``REMOTE_DIFFERS``. + +Constructing a request that breaks the rules raises ``ValueError``. +""" + +from __future__ import annotations + +import hashlib +import json +import re +from collections.abc import Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Final, Literal, TypedDict + +from microcosm_emitter_interface.identifiers import ( + bare_file_name, + repository_identifier, + repository_prefix, + role_identifier, + safe_identifier, +) +from microcosm_emitter_interface.protocol import PROTOCOL_VERSION + +DEFAULT_BUNDLE_PREFIX: Final = "staged" +#: The bundle's own inventory, compared with the repository copy. +BUNDLE_INVENTORY_ROLE: Final = "staged_manifest" +#: The build's manifest; it gains evidence after staging, so it is not an output. +BUILD_MANIFEST_ROLE: Final = "build_manifest" +#: The ``sha256sum`` listing written beside the bundle. +CHECKSUMS_ROLE: Final = "checksums" +SIDECAR_ROLES: Final = frozenset( + {BUNDLE_INVENTORY_ROLE, BUILD_MANIFEST_ROLE, CHECKSUMS_ROLE} +) +#: Roles of outputs readers look for by role rather than by file name. +DATASET_ROLE: Final = "dataset" +CALIBRATION_DIAGNOSTICS_ROLE: Final = "calibration_diagnostics" +_SHA256: Final = re.compile(r"^[0-9a-f]{64}$") + + +@dataclass(frozen=True) +class BundleFile: + """One preserved file and what it is. + + Every role other than the sidecar roles (``staged_manifest``, + ``build_manifest``, ``checksums``) marks an output, for example + ``dataset`` or ``calibration_diagnostics``. + """ + + name: str + role: str + local_path: Path + sha256: str + bytes: int + + def __post_init__(self) -> None: + bare_file_name(self.name, label="bundle file name") + role_identifier(self.role) + if not isinstance(self.sha256, str) or not _SHA256.fullmatch(self.sha256): + raise ValueError("sha256 must be a lowercase hexadecimal digest.") + if ( + isinstance(self.bytes, bool) + or not isinstance(self.bytes, int) + or self.bytes < 0 + ): + raise ValueError("bytes must be a nonnegative integer.") + + +def check_bundle_files(files: Sequence[BundleFile]) -> None: + names = [item.name for item in files] + if len(set(names)) != len(names): + raise ValueError("A bundle names each file once.") + roles = [item.role for item in files] + if roles.count(BUNDLE_INVENTORY_ROLE) != 1: + raise ValueError("A bundle has exactly one staged_manifest file.") + for role in (BUILD_MANIFEST_ROLE, CHECKSUMS_ROLE): + if roles.count(role) > 1: + raise ValueError(f"A bundle has at most one {role} file.") + if all(role in SIDECAR_ROLES for role in roles): + raise ValueError("A bundle has at least one output file.") + + +@dataclass(frozen=True) +class StageDatasetBundle: + """Hand over a finished build's bundle for ``//``. + + The files are written in the order given, in one commit whose message + defaults to Microcosm's ``Stage under /``. + """ + + repository: str + run_id: str + files: Sequence[BundleFile] + prefix: str = DEFAULT_BUNDLE_PREFIX + commit_message: str | None = None + protocol_version: int = field(default=PROTOCOL_VERSION, repr=False) + + def __post_init__(self) -> None: + repository_identifier(self.repository) + safe_identifier(self.run_id, label="run_id") + repository_prefix(self.prefix) + object.__setattr__(self, "files", tuple(self.files)) + if any(not isinstance(item, BundleFile) for item in self.files): + raise ValueError("files must be BundleFile values.") + check_bundle_files(self.files) + if self.commit_message is not None and ( + not isinstance(self.commit_message, str) or not self.commit_message.strip() + ): + raise ValueError("commit_message must be nonempty when given.") + + @property + def remote_prefix(self) -> str: + return f"{self.prefix}/{self.run_id}" + + def message(self) -> dict[str, Any]: + return { + "protocol_version": self.protocol_version, + "command": "stage_dataset_bundle", + "repository": self.repository, + "run_id": self.run_id, + "prefix": self.prefix, + "files": [ + { + "name": item.name, + "role": item.role, + "local_path": str(Path(item.local_path).resolve()), + "sha256": item.sha256, + "bytes": item.bytes, + } + for item in self.files + ], + "commit_message": self.commit_message, + } + + def handover_id(self) -> str: + return handover_identifier("dataset_bundle", self.message()) + + +def handover_identifier(kind: str, message: dict[str, Any]) -> str: + """A stable identifier of a hand-over's content, shared with the emitter.""" + content = { + key: value for key, value in message.items() if key != "protocol_version" + } + canonical = json.dumps( + {"kind": kind, "request": content}, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ) + return hashlib.sha256(canonical.encode()).hexdigest()[:32] + + +class DatasetBundleReceipt(TypedDict): + """The emitter's latest result for one bundle hand-over. + + ``uploaded``, ``already_staged`` and ``failed`` are final. While + ``pending``, ``error_code`` holds the last attempt's failure: + ``REPOSITORY_UNAVAILABLE``, ``UPLOAD_FAILED``, ``CREDENTIAL_MISSING`` or + ``PRESERVED_FILES_MISSING``. A final ``failed`` carries ``REMOTE_DIFFERS`` + or ``PRESERVED_BYTES_CHANGED``. ``refused`` requests were never queued. + """ + + handover_id: str + status: Literal[ + "pending", "uploaded", "already_staged", "failed", "refused", "unavailable" + ] + error_code: str | None + revision: str | None + attempts: int + repository: str + prefix: str + + +class WriteAccess(TypedDict): + """Whether a remote stage could work, checked before a long build. + + ``credential`` is ``missing`` without ``HF_TOKEN``, ``HUGGINGFACE_TOKEN`` + or a saved login. ``can_write`` is ``None`` when the Hub does not describe + the token's scope; then only the upload itself proves it. + """ + + status: Literal["checked", "unavailable"] + credential: Literal["present", "missing"] | None + reachable: bool | None + can_write: bool | None + error_type: str | None + + +def unavailable_bundle_receipt(request: StageDatasetBundle) -> DatasetBundleReceipt: + return { + "handover_id": request.handover_id(), + "status": "unavailable", + "error_code": None, + "revision": None, + "attempts": 0, + "repository": request.repository, + "prefix": request.remote_prefix, + } + + +def unavailable_write_access() -> WriteAccess: + return { + "status": "unavailable", + "credential": None, + "reachable": None, + "can_write": None, + "error_type": None, + } diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py index cddf20e..b335fa3 100644 --- a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py @@ -19,6 +19,14 @@ from pathlib import Path from typing import Any, Literal, cast +from microcosm_emitter_interface.bundles import ( + DatasetBundleReceipt, + StageDatasetBundle, + WriteAccess, + unavailable_bundle_receipt, + unavailable_write_access, +) +from microcosm_emitter_interface.identifiers import repository_identifier from microcosm_emitter_interface.launcher import ( EmitterProcess, EmitterStartupError, @@ -43,7 +51,9 @@ FILE_REQUEST_TIMEOUT_SECONDS, HUGGING_FACE_MODULE, ORRERY_MODULE, + PROTOCOL_VERSION, READ_BACK_TIMEOUT_SECONDS, + SETTLED_STATUSES, STAGING_REQUEST_TIMEOUT_SECONDS, STARTUP_TIMEOUT_SECONDS, TELEMETRY_MODULE, @@ -402,6 +412,77 @@ def read_back_run_files( return cast(ReadBackResult, value) return unavailable_read_back() + def _handover_receipt(self, handover_id: str) -> dict[str, Any] | None: + response = self._query( + HUGGING_FACE_MODULE, + lambda: { + "protocol_version": PROTOCOL_VERSION, + "command": "handover_receipt", + "handover_id": handover_id, + }, + kind="staging", + timeout=STAGING_REQUEST_TIMEOUT_SECONDS, + ) + value = None if response is None else response.get("receipt") + return value if isinstance(value, dict) else None + + def _wait_until_settled( + self, receipt: dict[str, Any], wait_seconds: float + ) -> dict[str, Any]: + deadline = time.monotonic() + max(0, wait_seconds) + while ( + receipt.get("status") not in SETTLED_STATUSES + and receipt.get("status") != "unavailable" + and time.monotonic() < deadline + ): + time.sleep(0.1) + latest = self._handover_receipt(receipt["handover_id"]) + if latest is None: + break + receipt = latest + return receipt + + def stage_dataset_bundle( + self, request: StageDatasetBundle, *, wait_seconds: float = 0 + ) -> DatasetBundleReceipt: + """Queue a staged dataset bundle; optionally wait for its upload. + + Returns once the hand-over is durably queued (``pending``), or after + it settles if that happens within ``wait_seconds``. + """ + message = request.message() + response = self._query( + HUGGING_FACE_MODULE, + lambda: message, + kind="staging", + timeout=FILE_REQUEST_TIMEOUT_SECONDS, + ) + value = None if response is None else response.get("receipt") + if not isinstance(value, dict): + return unavailable_bundle_receipt(request) + return cast(DatasetBundleReceipt, self._wait_until_settled(value, wait_seconds)) + + def dataset_bundle_receipt(self, handover_id: str) -> DatasetBundleReceipt | None: + return cast(DatasetBundleReceipt | None, self._handover_receipt(handover_id)) + + def check_write_access(self, repository: str) -> WriteAccess: + """Whether the build's credential can reach and write ``repository``.""" + message = { + "protocol_version": PROTOCOL_VERSION, + "command": "check_write_access", + "repository": repository_identifier(repository), + } + response = self._query( + HUGGING_FACE_MODULE, + lambda: message, + kind="staging", + timeout=READ_BACK_TIMEOUT_SECONDS, + ) + value = None if response is None else response.get("access") + if isinstance(value, dict): + return cast(WriteAccess, value) + return unavailable_write_access() + def close(self) -> None: """Ask the emitter to drain and stop; repeated calls do nothing.""" if self._closed: diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-check-write-access.json b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-check-write-access.json new file mode 100644 index 0000000..c2ab9d7 --- /dev/null +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-check-write-access.json @@ -0,0 +1,14 @@ +{ + "protocol_version": 1, + "module": "huggingface", + "kind": "query", + "sequence": 7, + "body": { + "protocol_version": 1, + "command": "check_write_access", + "repository": "policyengine/fixture-private" + }, + "response_keys": [ + "access" + ] +} diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-handover-receipt.json b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-handover-receipt.json new file mode 100644 index 0000000..89f3464 --- /dev/null +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-handover-receipt.json @@ -0,0 +1,14 @@ +{ + "protocol_version": 1, + "module": "huggingface", + "kind": "query", + "sequence": 6, + "body": { + "protocol_version": 1, + "command": "handover_receipt", + "handover_id": "fixture-handover" + }, + "response_keys": [ + "receipt" + ] +} diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-stage-dataset-bundle.json b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-stage-dataset-bundle.json new file mode 100644 index 0000000..24d4506 --- /dev/null +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-stage-dataset-bundle.json @@ -0,0 +1,61 @@ +{ + "protocol_version": 1, + "module": "huggingface", + "kind": "query", + "sequence": 5, + "files": { + "dataset.h5": "synthetic rows", + "calibration_diagnostics.json": "{\"schema_version\":8}", + "build_manifest.json": "{\"outputs\":{}}", + "staged_manifest.json": "{\"files\":{}}", + "sha256sums.txt": "fixture dataset.h5\n" + }, + "body": { + "protocol_version": 1, + "command": "stage_dataset_bundle", + "repository": "policyengine/fixture-private", + "run_id": "fixture-run", + "prefix": "staged", + "files": [ + { + "name": "dataset.h5", + "role": "dataset", + "local_path": "{directory}/dataset.h5", + "sha256": "5cc590aa15e97595e13963c72aea8d3fbe65fbf5dd5f8fc690ee479f73d45548", + "bytes": 14 + }, + { + "name": "calibration_diagnostics.json", + "role": "calibration_diagnostics", + "local_path": "{directory}/calibration_diagnostics.json", + "sha256": "2300b80de59210bc04c40b1aade68e7b5fc265901dbc9af4e9f506c0bcc36e7c", + "bytes": 20 + }, + { + "name": "build_manifest.json", + "role": "build_manifest", + "local_path": "{directory}/build_manifest.json", + "sha256": "5dbd51d32a488390bb3bfcea54726682fd3d80e073ddc7caf7df9d40404276b0", + "bytes": 14 + }, + { + "name": "staged_manifest.json", + "role": "staged_manifest", + "local_path": "{directory}/staged_manifest.json", + "sha256": "aae5a71db7cd42382ef749f87ca847684d9d4a517cc8235f53ea31bd492c3577", + "bytes": 12 + }, + { + "name": "sha256sums.txt", + "role": "checksums", + "local_path": "{directory}/sha256sums.txt", + "sha256": "167666035837f1b2e262cca1a2caef817e4fd111ea3fc9283328a2d79b4be4ed", + "bytes": 20 + } + ], + "commit_message": null + }, + "response_keys": [ + "receipt" + ] +} diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/protocol.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/protocol.py index 3c218d7..2c8022e 100644 --- a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/protocol.py +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/protocol.py @@ -27,8 +27,12 @@ FILE_REQUEST_TIMEOUT_SECONDS: Final = 30.0 #: Staging hand-overs copy a few small files into the emitter's queue. STAGING_REQUEST_TIMEOUT_SECONDS: Final = 5.0 -#: Read-back downloads a run's committed files from the Hub. +#: Read-back and the write-access check call the Hub while answering. READ_BACK_TIMEOUT_SECONDS: Final = 60.0 +#: Statuses after which a queued hand-over never changes. +SETTLED_STATUSES: Final = frozenset( + {"uploaded", "already_staged", "failed", "published", "refused"} +) STARTUP_TIMEOUT_SECONDS: Final = 15.0 STARTUP_POLL_SECONDS: Final = 0.02 DRAIN_SECONDS: Final = 15.0 diff --git a/src/microcosm_emitter/huggingface/service/bundles.py b/src/microcosm_emitter/huggingface/service/bundles.py new file mode 100644 index 0000000..b6c0047 --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/bundles.py @@ -0,0 +1,136 @@ +"""Stage one dataset bundle: the transport of Microcosm's ``stage_bundle``. + +The behaviour follows ``staging_dataset.stage_bundle``: one commit of every +file under ``//``, pinned to the head it checked, and +idempotent on the outputs' digests. Two things change because the work is now +queued: transient Hub failures are retried instead of ending the stage, and a +commit that loses a race to another writer is checked and tried again. +""" + +from __future__ import annotations + +import hashlib +import json +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from microcosm_emitter.huggingface.service.hub import ( + DatasetRepository, + add_operation, + is_parent_conflict, +) +from microcosm_emitter_interface.bundles import BUNDLE_INVENTORY_ROLE, SIDECAR_ROLES + +MAX_PARENT_CONFLICT_ATTEMPTS = 5 +REPOSITORY_UNAVAILABLE = "REPOSITORY_UNAVAILABLE" +UPLOAD_FAILED = "UPLOAD_FAILED" +REMOTE_DIFFERS = "REMOTE_DIFFERS" +PRESERVED_FILES_MISSING = "PRESERVED_FILES_MISSING" +PRESERVED_BYTES_CHANGED = "PRESERVED_BYTES_CHANGED" + + +@dataclass(frozen=True) +class BundleOutcome: + """One attempt's result: final unless ``retry``.""" + + status: str + error_code: str | None = None + revision: str | None = None + retry: bool = False + committed: bool = False + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for block in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _same_digests(remote: Any, local: Mapping[str, str]) -> bool: + """Microcosm ``_same_digests``: compare the outputs' digests only.""" + if not isinstance(remote, Mapping) or set(remote) != set(local): + return False + for name, digest in local.items(): + other = remote.get(name) + if not isinstance(other, Mapping) or other.get("sha256") != digest: + return False + return True + + +def _local_files_outcome(files: list[Mapping[str, Any]]) -> BundleOutcome | None: + for entry in files: + path = Path(entry["local_path"]) + try: + if not path.is_file(): + return BundleOutcome("pending", PRESERVED_FILES_MISSING, retry=True) + if ( + path.stat().st_size != entry["bytes"] + or _sha256_file(path) != entry["sha256"] + ): + return BundleOutcome("failed", PRESERVED_BYTES_CHANGED) + except OSError: + return BundleOutcome("pending", PRESERVED_FILES_MISSING, retry=True) + return None + + +def stage_bundle( + repository: DatasetRepository, request: Mapping[str, Any] +) -> BundleOutcome: + """Write the bundle once, or report why not; never raises for the Hub.""" + files = list(request["files"]) + local = _local_files_outcome(files) + if local is not None: + return local + remote_prefix = f"{request['prefix']}/{request['run_id']}" + inventory = next(entry for entry in files if entry["role"] == BUNDLE_INVENTORY_ROLE) + remote_inventory = f"{remote_prefix}/{inventory['name']}" + outputs = { + entry["name"]: entry["sha256"] + for entry in files + if entry["role"] not in SIDECAR_ROLES + } + message = request.get("commit_message") or ( + f"Stage {request['run_id']} under {remote_prefix}" + ) + for attempt in range(MAX_PARENT_CONFLICT_ATTEMPTS): + try: + exists = repository.file_exists(remote_inventory) + except Exception: + return BundleOutcome("pending", REPOSITORY_UNAVAILABLE, retry=True) + if exists: + try: + remote = json.loads(repository.download(remote_inventory)) + head = repository.head_revision() + except Exception: + return BundleOutcome("pending", REPOSITORY_UNAVAILABLE, retry=True) + remote_files = remote.get("files") if isinstance(remote, Mapping) else None + if not _same_digests(remote_files, outputs): + return BundleOutcome("failed", REMOTE_DIFFERS) + # The bundle's own commit, not the repository head at re-run time. + try: + revision = repository.last_commit(remote_inventory) or head + except Exception: + revision = head + return BundleOutcome("already_staged", revision=revision) + try: + parent = repository.head_revision() + revision = repository.commit( + [ + add_operation( + f"{remote_prefix}/{entry['name']}", Path(entry["local_path"]) + ) + for entry in files + ], + message=message, + parent_commit=parent, + ) + except Exception as error: + if is_parent_conflict(error) and attempt + 1 < MAX_PARENT_CONFLICT_ATTEMPTS: + continue + return BundleOutcome("pending", UPLOAD_FAILED, retry=True) + return BundleOutcome("uploaded", revision=revision, committed=True) + return BundleOutcome("pending", UPLOAD_FAILED, retry=True) # pragma: no cover diff --git a/src/microcosm_emitter/huggingface/service/delivery.py b/src/microcosm_emitter/huggingface/service/delivery.py index 05245ce..8c98395 100644 --- a/src/microcosm_emitter/huggingface/service/delivery.py +++ b/src/microcosm_emitter/huggingface/service/delivery.py @@ -15,13 +15,17 @@ from datetime import UTC, datetime from typing import Any +from microcosm_emitter.huggingface.service.bundles import stage_bundle from microcosm_emitter.huggingface.service.hub import DatasetRepository -from microcosm_emitter.huggingface.service.queue import HandoverQueue +from microcosm_emitter.huggingface.service.queue import ClaimedJob, HandoverQueue from microcosm_emitter.huggingface.service.run_files import commit_run_files #: Long enough for one commit of a large upload; a crashed writer's lease #: expires and another host on this machine resumes the repository. REPOSITORY_LEASE_SECONDS = 900.0 +#: The job kinds this emitter delivers. +DATASET_BUNDLE = "dataset_bundle" +SUPPORTED_JOB_KINDS = frozenset({DATASET_BUNDLE}) CREDENTIAL_MISSING = "CREDENTIAL_MISSING" UPLOAD_FAILED = "UPLOAD_FAILED" CREDENTIAL_MISSING_MESSAGE = ( @@ -65,7 +69,9 @@ def warn_once(self, key: str, message: str) -> None: def run_cycle(self) -> int: """Make at most one commit per due repository; return how many.""" - repositories = self.queue.due_repositories(self.clock()) + repositories = self.queue.due_repositories( + self.clock(), kinds=SUPPORTED_JOB_KINDS + ) if not repositories: return 0 api = self.api_factory() @@ -95,15 +101,86 @@ def _deliver_repository( return 0 try: if api is None: - self._credential_missing(repository) + self._credential_missing(repository, run_files_only=run_files_only) return 0 hub = DatasetRepository(repository, api) + if not run_files_only: + job = self._claim_job(repository) + if job is not None: + return int(self._deliver_job(hub, job)) + if not self.queue.run_files_due(repository, self.clock()): + return 0 return int(self._deliver_run_files(hub)) finally: self.queue.release_repository(repository, lease) - def _credential_missing(self, repository: str) -> None: + def _claim_job(self, repository: str) -> ClaimedJob | None: + return self.queue.claim_job( + repository, + now=self.clock(), + lease_seconds=REPOSITORY_LEASE_SECONDS, + kinds=SUPPORTED_JOB_KINDS, + ) + + def _finish_job( + self, + job: ClaimedJob, + *, + status: str, + error_code: str | None, + revision: str | None, + retry: bool, + committed: bool = False, + ) -> None: + self.queue.finish_job( + job, + receipt={ + "handover_id": job.handover_id, + "status": status, + "error_code": error_code, + "revision": revision, + "attempts": job.attempts + 1, + "repository": job.target, + "prefix": f"{job.request['prefix']}/{job.request['run_id']}", + }, + retry=retry, + now=self.clock(), + committed=committed, + ) + + def _deliver_job(self, hub: DatasetRepository, job: ClaimedJob) -> bool: + outcome = stage_bundle(hub, job.request) + if outcome.retry and outcome.error_code is not None: + self.warn_once( + outcome.error_code, + UPLOAD_FAILED_MESSAGE.format(error_type=outcome.error_code), + ) + self._finish_job( + job, + status=outcome.status, + error_code=outcome.error_code, + revision=outcome.revision, + retry=outcome.retry, + committed=outcome.committed, + ) + return outcome.committed + + def _credential_missing(self, repository: str, *, run_files_only: bool) -> None: self.warn_once(CREDENTIAL_MISSING, CREDENTIAL_MISSING_MESSAGE) + job = None if run_files_only else self._claim_job(repository) + if job is not None: + self._finish_job( + job, + status="pending", + error_code=CREDENTIAL_MISSING, + revision=None, + retry=True, + ) + return + if not run_files_only and not self.queue.run_files_due( + repository, self.clock() + ): + return entries = self.queue.pending_run_files(repository) index_entries = self.queue.pending_run_index_entries(repository) if entries or index_entries: diff --git a/src/microcosm_emitter/huggingface/service/module.py b/src/microcosm_emitter/huggingface/service/module.py index 8739b92..30ef45d 100644 --- a/src/microcosm_emitter/huggingface/service/module.py +++ b/src/microcosm_emitter/huggingface/service/module.py @@ -14,6 +14,7 @@ from microcosm_emitter.host.contracts import JsonObject, ModuleContext from microcosm_emitter.huggingface.service.delivery import ( CREDENTIAL_MISSING, + DATASET_BUNDLE, HuggingFaceDelivery, ) from microcosm_emitter.huggingface.service.hub import ( @@ -23,9 +24,12 @@ from microcosm_emitter.huggingface.service.queue import HandoverQueue, sha256_bytes from microcosm_emitter.huggingface.service.requests import ( RequestRefusedError, + parse_stage_dataset_bundle, parse_stage_run_files, + require_supported_version, run_reference, ) +from microcosm_emitter_interface.identifiers import repository_identifier WORKER_ERROR_MESSAGE = ( "Microcosm Hugging Face delivery hit a local queue error; handed-over files " @@ -95,8 +99,81 @@ def query(self, message: JsonObject) -> dict[str, Any]: } if command == "read_back_run_files": return {"read_back": self._read_back_run_files(message)} + if command == "stage_dataset_bundle": + return {"receipt": self._stage_dataset_bundle(message)} + if command == "handover_receipt": + return {"receipt": self._handover_receipt(message)} + if command == "check_write_access": + return {"access": self._check_write_access(message)} raise ValueError("unsupported Hugging Face hand-over command") + def _stage_dataset_bundle(self, message: JsonObject) -> dict[str, Any]: + assert self.queue is not None + try: + hand_over = parse_stage_dataset_bundle(message) + except RequestRefusedError as refusal: + repository = message.get("repository") + return { + "handover_id": "", + "status": "refused", + "error_code": refusal.code, + "revision": None, + "attempts": 0, + "repository": repository if isinstance(repository, str) else "", + "prefix": "", + } + return self.queue.enqueue_job( + handover_id=hand_over.handover_id, + kind=DATASET_BUNDLE, + target=hand_over.repository, + request=hand_over.request, + receipt={ + "handover_id": hand_over.handover_id, + "status": "pending", + "error_code": None, + "revision": None, + "attempts": 0, + "repository": hand_over.repository, + "prefix": hand_over.remote_prefix, + }, + now=self.clock(), + ) + + def _handover_receipt(self, message: JsonObject) -> dict[str, Any] | None: + assert self.queue is not None + require_supported_version(message) + handover_id = message.get("handover_id") + if not isinstance(handover_id, str) or not handover_id: + raise ValueError("handover_id must be a nonempty string") + return self.queue.job_receipt(handover_id) + + def _check_write_access(self, message: JsonObject) -> dict[str, Any]: + """Microcosm's staged-dataset preflight: credential, reach and scope.""" + require_supported_version(message) + repository = repository_identifier(message.get("repository")) + access: dict[str, Any] = { + "status": "checked", + "credential": "missing", + "reachable": None, + "can_write": None, + "error_type": None, + } + api = self.api_factory() + if api is None: + return access + hub = DatasetRepository(repository, api) + access["credential"] = "present" + try: + hub.head_revision() + except Exception as error: + # Only the type name: the Hub's message can carry request details. + return {**access, "reachable": False, "error_type": type(error).__name__} + try: + can_write = hub.credential_can_write() + except Exception: + can_write = None + return {**access, "reachable": True, "can_write": can_write} + def _stage_run_files(self, message: JsonObject) -> dict[str, Any]: assert self.queue is not None try: diff --git a/src/microcosm_emitter/huggingface/service/queue.py b/src/microcosm_emitter/huggingface/service/queue.py index d3d1458..4797821 100644 --- a/src/microcosm_emitter/huggingface/service/queue.py +++ b/src/microcosm_emitter/huggingface/service/queue.py @@ -152,8 +152,14 @@ def record_commit(self, repository: str, *, now: float) -> None: with self._sessions.begin() as session: self._schedule(session, repository).last_commit_at = now - def due_repositories(self, now: float) -> list[str]: - """Repositories with work due now that may receive a commit now.""" + def due_repositories( + self, now: float, *, kinds: Iterable[str] | None = None + ) -> list[str]: + """Repositories with work due now that may receive a commit now. + + ``kinds`` limits jobs to those this emitter can deliver, so a job + queued by a newer emitter sharing the database waits for it. + """ with self._sessions() as session: run_file_repositories = set( session.scalars( @@ -169,16 +175,15 @@ def due_repositories(self, now: float) -> list[str]: ) ) ) - job_repositories = set( - session.scalars( - select(HandoverJob.target).where( - HandoverJob.destination == DESTINATION_HUGGING_FACE, - HandoverJob.status == JOB_PENDING, - HandoverJob.next_attempt_at <= now, - HandoverJob.lease_until <= now, - ) - ) + job_query = select(HandoverJob.target).where( + HandoverJob.destination == DESTINATION_HUGGING_FACE, + HandoverJob.status == JOB_PENDING, + HandoverJob.next_attempt_at <= now, + HandoverJob.lease_until <= now, ) + if kinds is not None: + job_query = job_query.where(HandoverJob.kind.in_(list(kinds))) + job_repositories = set(session.scalars(job_query)) due = [] for repository in sorted(run_file_repositories | job_repositories): schedule = session.get(RepositoryCommitSchedule, repository) @@ -291,6 +296,32 @@ def stage_run_files( delivery.run_index_due_at = min(delivery.run_index_due_at, now) delivery.updated_at = now + def run_files_due(self, repository: str, now: float) -> bool: + """Whether pending run files are due and past any failure backoff.""" + with self._sessions() as session: + schedule = session.get(RepositoryCommitSchedule, repository) + if schedule is not None and schedule.retry_after > now: + return False + due_file = session.scalar( + select(StagedRunFile.path_in_repo) + .where( + StagedRunFile.repository == repository, + StagedRunFile.pending.is_(True), + StagedRunFile.due_at <= now, + ) + .limit(1) + ) + due_index = session.scalar( + select(StagedRunDelivery.run_id) + .where( + StagedRunDelivery.repository == repository, + StagedRunDelivery.run_index_pending.is_(True), + StagedRunDelivery.run_index_due_at <= now, + ) + .limit(1) + ) + return due_file is not None or due_index is not None + def repositories_with_pending_run_files(self) -> list[str]: with self._sessions() as session: return sorted( @@ -499,6 +530,7 @@ def claim_job( now: float, lease_seconds: float, destination: str = DESTINATION_HUGGING_FACE, + kinds: Iterable[str] | None = None, ) -> ClaimedJob | None: """Lease the earliest-due pending job for ``target``. @@ -506,17 +538,19 @@ def claim_job( leaves its job eligible again once the lease expires. """ with self._sessions.begin() as session: + statement = select(HandoverJob).where( + HandoverJob.destination == destination, + HandoverJob.target == target, + HandoverJob.status == JOB_PENDING, + HandoverJob.next_attempt_at <= now, + HandoverJob.lease_until <= now, + ) + if kinds is not None: + statement = statement.where(HandoverJob.kind.in_(list(kinds))) candidate = session.scalar( - select(HandoverJob) - .where( - HandoverJob.destination == destination, - HandoverJob.target == target, - HandoverJob.status == JOB_PENDING, - HandoverJob.next_attempt_at <= now, - HandoverJob.lease_until <= now, - ) - .order_by(HandoverJob.next_attempt_at, HandoverJob.created_at) - .limit(1) + statement.order_by( + HandoverJob.next_attempt_at, HandoverJob.created_at + ).limit(1) ) if candidate is None: return None diff --git a/src/microcosm_emitter/huggingface/service/requests.py b/src/microcosm_emitter/huggingface/service/requests.py index 4b20631..5b7e13b 100644 --- a/src/microcosm_emitter/huggingface/service/requests.py +++ b/src/microcosm_emitter/huggingface/service/requests.py @@ -13,6 +13,11 @@ from typing import Any from microcosm_emitter.compatibility import SUPPORTED_PROTOCOL_VERSIONS +from microcosm_emitter_interface.bundles import ( + BundleFile, + StageDatasetBundle, + handover_identifier, +) from microcosm_emitter_interface.identifiers import ( relative_file_path, repository_identifier, @@ -141,3 +146,62 @@ def parse_stage_run_files(message: Mapping[str, Any]) -> RunFilesHandOver: flush=flush, interval=float(interval), ) + + +@dataclass(frozen=True) +class BundleHandOver: + handover_id: str + repository: str + remote_prefix: str + request: dict[str, Any] + + +def parse_stage_dataset_bundle(message: Mapping[str, Any]) -> BundleHandOver: + """Check a bundle with the interface's rules and the files' sizes. + + Hashing large files waits for the worker; here only existence and size + are checked, so the request answers quickly. + """ + require_supported_version(message) + entries = message.get("files") + if not isinstance(entries, list) or any( + not isinstance(entry, Mapping) for entry in entries + ): + raise RequestRefusedError("INVALID_REQUEST") + try: + request = StageDatasetBundle( + repository=message.get("repository"), + run_id=message.get("run_id"), + prefix=message.get("prefix"), + commit_message=message.get("commit_message"), + files=tuple( + BundleFile( + name=entry.get("name"), + role=entry.get("role"), + local_path=Path(str(entry.get("local_path"))), + sha256=entry.get("sha256"), + bytes=entry.get("bytes"), + ) + for entry in entries + ), + ) + except (ValueError, TypeError): + raise RequestRefusedError("INVALID_REQUEST") from None + for entry, item in zip(entries, request.files, strict=True): + path = Path(str(entry.get("local_path"))) + if not path.is_absolute(): + raise RequestRefusedError("INVALID_REQUEST") + try: + if not path.is_file(): + raise RequestRefusedError("LOCAL_FILE_MISSING") + if path.stat().st_size != item.bytes: + raise RequestRefusedError("FILE_SIZE_DIFFERS") + except OSError: + raise RequestRefusedError("LOCAL_FILE_MISSING") from None + canonical = dict(message) + return BundleHandOver( + handover_id=handover_identifier("dataset_bundle", canonical), + repository=request.repository, + remote_prefix=request.remote_prefix, + request={key: value for key, value in canonical.items()}, + ) diff --git a/tests/huggingface/test_bundles.py b/tests/huggingface/test_bundles.py new file mode 100644 index 0000000..9bc01d2 --- /dev/null +++ b/tests/huggingface/test_bundles.py @@ -0,0 +1,409 @@ +"""Staged dataset bundles: Microcosm's paths, idempotency and error codes.""" + +import hashlib +import json + +import pytest +from fake_hugging_face import ( + PRIVATE_REPOSITORY, + FakeHubHTTPError, + FakeHuggingFaceApi, +) + +from microcosm_emitter.host.contracts import ModuleContext +from microcosm_emitter.huggingface.service.module import HuggingFaceModule +from microcosm_emitter_interface.bundles import BundleFile, StageDatasetBundle + +DATASET = b"synthetic h5 bytes" +DIAGNOSTICS = b'{"schema_version":8}' + + +def bundle_files(directory, *, dataset=DATASET, manifest_note="first"): + """A bundle laid out as Microcosm's write_sidecars leaves it.""" + directory.mkdir(parents=True, exist_ok=True) + outputs = { + "microcosm_uk_2024_25.h5": ("dataset", dataset), + "microcosm_uk_2024_25.diagnostics.json": ( + "calibration_diagnostics", + DIAGNOSTICS, + ), + } + digests = { + name: {"sha256": hashlib.sha256(content).hexdigest(), "bytes": len(content)} + for name, (_, content) in outputs.items() + } + manifest = json.dumps({"outputs": digests, "note": manifest_note}).encode() + staged_manifest = json.dumps( + {"schema_name": "microcosm.staged-dataset.manifest", "files": digests}, + indent=2, + sort_keys=True, + ).encode() + sums = "".join( + f"{entry['sha256']} {name}\n" for name, entry in sorted(digests.items()) + ).encode() + files = { + **outputs, + "rowwise_candidate_manifest.json": ("build_manifest", manifest), + "staged_manifest.json": ("staged_manifest", staged_manifest), + "sha256sums.txt": ("checksums", sums), + } + bundle = [] + for name, (role, content) in files.items(): + (directory / name).write_bytes(content) + bundle.append( + BundleFile( + name=name, + role=role, + local_path=directory / name, + sha256=hashlib.sha256(content).hexdigest(), + bytes=len(content), + ) + ) + return tuple(bundle), {name: content for name, (_, content) in files.items()} + + +def stage_bundle(controls, run_id="uk-run-1", **options): + directory = controls.tmp_path / options.pop("directory", run_id) + files, contents = bundle_files(directory, **options) + request = StageDatasetBundle( + repository=PRIVATE_REPOSITORY, run_id=run_id, files=files + ) + receipt = controls.module.query(request.message())["receipt"] + return request, receipt, contents + + +def receipt_of(controls, handover_id): + return controls.module.query( + { + "protocol_version": 1, + "command": "handover_receipt", + "handover_id": handover_id, + } + )["receipt"] + + +def cycle(controls): + controls.module.delivery.run_cycle() + + +def private_commits(controls): + return controls.hub.commits(PRIVATE_REPOSITORY) + + +def test_bundle_is_one_commit_at_microcosm_paths_with_its_message( + hugging_face_module, +): + controls = hugging_face_module + head_before = controls.hub.repositories[PRIVATE_REPOSITORY].head.oid + request, receipt, contents = stage_bundle(controls) + assert receipt["status"] == "pending" + assert receipt["prefix"] == "staged/uk-run-1" + cycle(controls) + (commit,) = private_commits(controls) + assert commit.added == { + f"staged/uk-run-1/{name}": content for name, content in contents.items() + } + assert list(commit.added) == [ + f"staged/uk-run-1/{item.name}" for item in request.files + ] + assert commit.message == "Stage uk-run-1 under staged/uk-run-1" + assert commit.parent == head_before + (create_commit,) = [ + arguments + for method, arguments in controls.hub.calls + if method == "create_commit" + ] + assert create_commit["parent_commit"] == head_before + receipt = receipt_of(controls, request.handover_id()) + assert receipt["status"] == "uploaded" + assert receipt["revision"] == commit.oid + assert receipt["error_code"] is None + + +def test_restaging_the_same_outputs_keeps_the_original_revision(hugging_face_module): + controls = hugging_face_module + first, _, _ = stage_bundle(controls) + cycle(controls) + (commit,) = private_commits(controls) + controls.clock.advance(60) + # The build appended evidence to its manifest; the outputs are unchanged. + again, _, _ = stage_bundle(controls, directory="again", manifest_note="evidence") + assert again.handover_id() != first.handover_id() + cycle(controls) + assert private_commits(controls) == [commit] + receipt = receipt_of(controls, again.handover_id()) + assert (receipt["status"], receipt["revision"]) == ("already_staged", commit.oid) + + +def test_a_different_bundle_under_the_same_run_is_never_overwritten( + hugging_face_module, +): + controls = hugging_face_module + stage_bundle(controls) + cycle(controls) + controls.clock.advance(60) + other, _, _ = stage_bundle(controls, directory="other", dataset=b"other bytes") + cycle(controls) + assert len(private_commits(controls)) == 1 + receipt = receipt_of(controls, other.handover_id()) + assert (receipt["status"], receipt["error_code"]) == ("failed", "REMOTE_DIFFERS") + controls.clock.advance(600) + cycle(controls) + assert receipt_of(controls, other.handover_id())["attempts"] == 1 + + +@pytest.mark.parametrize( + ("method", "error_code"), + [("file_exists", "REPOSITORY_UNAVAILABLE"), ("create_commit", "UPLOAD_FAILED")], +) +def test_hub_failures_are_retried_with_microcosm_error_codes( + hugging_face_module, method, error_code +): + controls = hugging_face_module + controls.hub.fail_next(method, FakeHubHTTPError(503)) + request, _, _ = stage_bundle(controls) + cycle(controls) + receipt = receipt_of(controls, request.handover_id()) + assert (receipt["status"], receipt["error_code"], receipt["attempts"]) == ( + "pending", + error_code, + 1, + ) + assert private_commits(controls) == [] + controls.clock.advance(2) + cycle(controls) + assert receipt_of(controls, request.handover_id())["status"] == "uploaded" + + +def test_a_commit_that_loses_a_race_is_retried_on_the_new_head(hugging_face_module): + controls = hugging_face_module + controls.hub.run_before_next_commit( + lambda: controls.hub.write_directly(PRIVATE_REPOSITORY, {"other.json": b"{}"}) + ) + request, _, _ = stage_bundle(controls) + cycle(controls) + assert receipt_of(controls, request.handover_id())["status"] == "uploaded" + ours = [ + commit + for commit in private_commits(controls) + if commit.message.startswith("Stage uk-run-1") + ] + assert len(ours) == 1 + + +def test_local_bytes_changed_after_the_hand_over_are_never_uploaded( + hugging_face_module, +): + controls = hugging_face_module + request, _, _ = stage_bundle(controls) + dataset = next(item for item in request.files if item.role == "dataset") + dataset.local_path.write_bytes(b"X" * dataset.bytes) + cycle(controls) + assert private_commits(controls) == [] + receipt = receipt_of(controls, request.handover_id()) + assert (receipt["status"], receipt["error_code"]) == ( + "failed", + "PRESERVED_BYTES_CHANGED", + ) + + +def test_missing_local_files_stay_queued_until_restored(hugging_face_module): + controls = hugging_face_module + request, _, contents = stage_bundle(controls) + dataset = next(item for item in request.files if item.role == "dataset") + dataset.local_path.unlink() + cycle(controls) + receipt = receipt_of(controls, request.handover_id()) + assert (receipt["status"], receipt["error_code"]) == ( + "pending", + "PRESERVED_FILES_MISSING", + ) + dataset.local_path.write_bytes(contents[dataset.name]) + controls.clock.advance(2) + cycle(controls) + assert receipt_of(controls, request.handover_id())["status"] == "uploaded" + + +def test_repeating_a_hand_over_queues_it_once(hugging_face_module): + controls = hugging_face_module + request, first, _ = stage_bundle(controls) + second = controls.module.query(request.message())["receipt"] + assert first == second + cycle(controls) + controls.clock.advance(60) + cycle(controls) + assert len(private_commits(controls)) == 1 + + +def test_bundles_and_run_files_for_one_repository_take_separate_cycles( + hugging_face_module, +): + from microcosm_emitter_interface.staging import RunFile, StageRunFiles + + controls = hugging_face_module + request, _, _ = stage_bundle(controls) + progress = controls.tmp_path / "progress.json" + progress.write_bytes(b"{}") + controls.module.query( + StageRunFiles( + repository=PRIVATE_REPOSITORY, + run_id="uk-run-1", + files=(RunFile("progress.json", progress),), + flush=True, + ).message() + ) + cycle(controls) + assert len(private_commits(controls)) == 1 + assert receipt_of(controls, request.handover_id())["status"] == "uploaded" + cycle(controls) + assert len(private_commits(controls)) == 1 + controls.clock.advance(10) + cycle(controls) + assert [list(commit.added) for commit in private_commits(controls)][-1] == [ + "runs/uk-run-1/progress.json" + ] + + +def test_without_a_credential_bundles_stay_queued(hugging_face_module): + controls = hugging_face_module + controls.remove_credential() + request, _, _ = stage_bundle(controls) + cycle(controls) + receipt = receipt_of(controls, request.handover_id()) + assert (receipt["status"], receipt["error_code"]) == ( + "pending", + "CREDENTIAL_MISSING", + ) + controls.restore_credential() + controls.clock.advance(2) + cycle(controls) + assert receipt_of(controls, request.handover_id())["status"] == "uploaded" + + +@pytest.mark.parametrize( + ("edit", "refusal_code"), + [ + (lambda message: message["files"].pop(3), "INVALID_REQUEST"), + ( + lambda message: message["files"][0].update(role="Not A Role"), + "INVALID_REQUEST", + ), + ( + lambda message: message["files"][1].update( + name=message["files"][0]["name"] + ), + "INVALID_REQUEST", + ), + (lambda message: message["files"][0].update(sha256="abc"), "INVALID_REQUEST"), + (lambda message: message["files"][0].update(bytes=1), "FILE_SIZE_DIFFERS"), + ( + lambda message: message["files"][0].update( + local_path="/nonexistent/fixture.h5" + ), + "LOCAL_FILE_MISSING", + ), + (lambda message: message.update(prefix="releases"), "INVALID_REQUEST"), + ( + lambda message: message.update(protocol_version=99), + "UNSUPPORTED_PROTOCOL_VERSION", + ), + ], + ids=[ + "no-inventory", + "role", + "duplicate-name", + "digest", + "size", + "missing-file", + "release-records", + "protocol-version", + ], +) +def test_malformed_bundle_hand_overs_are_refused_and_not_queued( + hugging_face_module, edit, refusal_code +): + controls = hugging_face_module + files, _ = bundle_files(controls.tmp_path / "bundle") + message = StageDatasetBundle( + repository=PRIVATE_REPOSITORY, run_id="uk-run-1", files=files + ).message() + edit(message) + receipt = controls.module.query(message)["receipt"] + assert (receipt["status"], receipt["error_code"]) == ("refused", refusal_code) + cycle(controls) + assert private_commits(controls) == [] + + +@pytest.mark.parametrize( + ("whoami", "fail_head", "expected"), + [ + ( + {"auth": {"accessToken": {"role": "write"}}}, + False, + {"credential": "present", "reachable": True, "can_write": True}, + ), + ( + {"auth": {"accessToken": {"role": "read"}}}, + False, + {"credential": "present", "reachable": True, "can_write": False}, + ), + ( + {"name": "fixture"}, + False, + {"credential": "present", "reachable": True, "can_write": None}, + ), + ( + {"auth": {"accessToken": {"role": "write"}}}, + True, + {"credential": "present", "reachable": False, "can_write": None}, + ), + ], + ids=["write", "read", "unknown-scope", "unreachable"], +) +def test_write_access_check_reports_reach_and_scope( + tmp_path, whoami, fail_head, expected +): + import os + + hub = FakeHuggingFaceApi(whoami=whoami) + hub.create_repository(PRIVATE_REPOSITORY) + if fail_head: + hub.fail_next("repo_info", FakeHubHTTPError(503)) + module = HuggingFaceModule( + {"spool_path": str(tmp_path / "events.sqlite3")}, + ModuleContext(os.getpid()), + api_factory=lambda: hub, + ) + module.initialize() + try: + access = module.query( + { + "protocol_version": 1, + "command": "check_write_access", + "repository": PRIVATE_REPOSITORY, + } + )["access"] + finally: + module.close(0) + assert {key: access[key] for key in expected} == expected + assert access["status"] == "checked" + assert access["error_type"] == ("FakeHubHTTPError" if fail_head else None) + + +def test_write_access_check_without_a_credential(hugging_face_module): + controls = hugging_face_module + controls.remove_credential() + access = controls.module.query( + { + "protocol_version": 1, + "command": "check_write_access", + "repository": PRIVATE_REPOSITORY, + } + )["access"] + assert access == { + "status": "checked", + "credential": "missing", + "reachable": None, + "can_write": None, + "error_type": None, + } diff --git a/tests/interface/conftest.py b/tests/interface/conftest.py index 8d1f57a..39caf58 100644 --- a/tests/interface/conftest.py +++ b/tests/interface/conftest.py @@ -1,13 +1,23 @@ """Loopback services shared by the interface tests; nothing leaves this machine.""" import json +import os +import shutil import stat import sys +import tempfile import threading +import time +from dataclasses import dataclass, field from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path import pytest +from fake_hugging_face import ( + PRIVATE_REPOSITORY, + STAGING_REPOSITORY, + FakeHuggingFaceApi, +) @pytest.fixture @@ -74,3 +84,69 @@ def isolated_path(monkeypatch, tmp_path): directory.mkdir() monkeypatch.setenv("PATH", str(directory)) return directory + + +@dataclass +class InThreadEmitter: + """An emitter host running in a thread of the test process.""" + + transport: object + description: dict = field(default_factory=dict) + + def close(self): + self.transport.close() + + +@pytest.fixture +def in_thread_emitter(): + return InThreadEmitter + + +@pytest.fixture +def emitter_with_hub(tmp_path): + """A client talking over a real socket to the module and an in-memory Hub.""" + from microcosm_emitter.host.contracts import ModuleContext + from microcosm_emitter.host.runtime import ServiceRuntime + from microcosm_emitter.huggingface.service.module import HuggingFaceModule + from microcosm_emitter_interface.client import EmitterClient + from microcosm_emitter_interface.transport import SocketTransport + + hub = FakeHuggingFaceApi(whoami={"auth": {"accessToken": {"role": "write"}}}) + hub.create_repository(STAGING_REPOSITORY) + hub.create_repository(PRIVATE_REPOSITORY) + runtime_directory = Path(tempfile.mkdtemp(prefix="pe-stage-", dir="/tmp")) + socket_path = runtime_directory / "service.sock" + module = HuggingFaceModule( + {"spool_path": str(tmp_path / "events.sqlite3")}, + ModuleContext(os.getpid()), + api_factory=lambda: hub, + ) + runtime = ServiceRuntime( + socket_path, + modules={"huggingface": module}, + is_parent_alive=lambda: True, + tick_seconds=0.01, + ) + thread = threading.Thread(target=runtime.run) + thread.start() + transport = SocketTransport(socket_path) + deadline = time.monotonic() + 5 + while True: + try: + transport.ping() + break + except Exception: + if time.monotonic() > deadline: + raise + time.sleep(0.01) + client = EmitterClient( + InThreadEmitter( + transport, {"modules": ["huggingface"], "protocol_versions": [1]} + ) + ) + try: + yield client, hub + finally: + client.close() + thread.join(timeout=20) + shutil.rmtree(runtime_directory, ignore_errors=True) diff --git a/tests/interface/test_bundle_client.py b/tests/interface/test_bundle_client.py new file mode 100644 index 0000000..72ef59c --- /dev/null +++ b/tests/interface/test_bundle_client.py @@ -0,0 +1,156 @@ +"""Typed staged-dataset-bundle hand-overs, each file named with its role.""" + +import hashlib +from pathlib import Path + +import pytest +from fake_hugging_face import PRIVATE_REPOSITORY + +from microcosm_emitter_interface.bundles import BundleFile, StageDatasetBundle +from microcosm_emitter_interface.client import EmitterClient +from microcosm_emitter_interface.transport import SocketTransport + + +def bundle_file(directory: Path, name: str, role: str, content: bytes) -> BundleFile: + path = directory / name + path.write_bytes(content) + return BundleFile( + name=name, + role=role, + local_path=path, + sha256=hashlib.sha256(content).hexdigest(), + bytes=len(content), + ) + + +def minimal_bundle(directory: Path) -> tuple[BundleFile, ...]: + return ( + bundle_file(directory, "dataset.h5", "dataset", b"rows"), + bundle_file( + directory, "calibration_diagnostics.json", "calibration_diagnostics", b"{}" + ), + bundle_file( + directory, + "staged_manifest.json", + "staged_manifest", + b'{"files": {"calibration_diagnostics.json": {"sha256": "x"}}}', + ), + ) + + +def with_replaced(files, index, **changes): + original = files[index] + replaced = BundleFile( + name=changes.get("name", original.name), + role=changes.get("role", original.role), + local_path=original.local_path, + sha256=changes.get("sha256", original.sha256), + bytes=changes.get("bytes", original.bytes), + ) + return files[:index] + (replaced,) + files[index + 1 :] + + +@pytest.mark.parametrize( + "build_files", + [ + lambda files: files[:2], + lambda files: files + (files[2],), + lambda files: with_replaced(files, 0, name="folder/dataset.h5"), + lambda files: with_replaced(files, 0, role="Dataset"), + lambda files: with_replaced(files, 0, sha256="not-a-digest"), + lambda files: with_replaced(files, 0, bytes=-1), + lambda files: files[2:], + lambda files: with_replaced(files, 1, name="dataset.h5"), + ], + ids=[ + "no-inventory", + "two-inventories", + "folder-in-name", + "role", + "digest", + "size", + "no-output", + "duplicate-name", + ], +) +def test_bundles_that_break_the_rules_cannot_be_built(tmp_path, build_files): + with pytest.raises(ValueError): + StageDatasetBundle( + repository=PRIVATE_REPOSITORY, + run_id="uk-run-1", + files=build_files(minimal_bundle(tmp_path)), + ) + + +def test_hand_over_identity_follows_the_request_content(tmp_path): + files = minimal_bundle(tmp_path) + request = StageDatasetBundle( + repository=PRIVATE_REPOSITORY, run_id="uk-run-1", files=files + ) + same = StageDatasetBundle( + repository=PRIVATE_REPOSITORY, run_id="uk-run-1", files=files + ) + changed = StageDatasetBundle( + repository=PRIVATE_REPOSITORY, + run_id="uk-run-1", + files=with_replaced(files, 0, sha256="0" * 64), + ) + assert request.handover_id() == same.handover_id() + assert request.handover_id() != changed.handover_id() + assert len(request.handover_id()) == 32 + + +def test_message_names_each_file_role_and_digest(tmp_path): + files = minimal_bundle(tmp_path) + message = StageDatasetBundle( + repository=PRIVATE_REPOSITORY, + run_id="uk-run-1", + files=files, + commit_message="Stage the fixture", + ).message() + assert message["command"] == "stage_dataset_bundle" + assert message["prefix"] == "staged" + assert message["commit_message"] == "Stage the fixture" + assert message["files"][1] == { + "name": "calibration_diagnostics.json", + "role": "calibration_diagnostics", + "local_path": str(files[1].local_path.resolve()), + "sha256": files[1].sha256, + "bytes": 2, + } + + +def test_client_stages_a_bundle_and_waits_for_its_upload(tmp_path, emitter_with_hub): + client, hub = emitter_with_hub + files = minimal_bundle(tmp_path) + request = StageDatasetBundle( + repository=PRIVATE_REPOSITORY, run_id="uk-run-1", files=files + ) + receipt = client.stage_dataset_bundle(request, wait_seconds=20) + assert receipt["status"] == "uploaded", receipt + (commit,) = hub.commits(PRIVATE_REPOSITORY) + assert receipt["revision"] == commit.oid + assert set(commit.added) == {f"staged/uk-run-1/{item.name}" for item in files} + assert client.dataset_bundle_receipt(request.handover_id()) == receipt + access = client.check_write_access(PRIVATE_REPOSITORY) + assert (access["status"], access["can_write"]) == ("checked", True) + + +def test_client_without_the_module_reports_unavailable(tmp_path, in_thread_emitter): + client = EmitterClient( + in_thread_emitter( + SocketTransport(tmp_path / "unused.sock"), + {"modules": ["telemetry"], "protocol_versions": [1]}, + ) + ) + request = StageDatasetBundle( + repository=PRIVATE_REPOSITORY, + run_id="uk-run-1", + files=minimal_bundle(tmp_path), + ) + receipt = client.stage_dataset_bundle(request, wait_seconds=1) + assert (receipt["status"], receipt["handover_id"]) == ( + "unavailable", + request.handover_id(), + ) + assert client.check_write_access(PRIVATE_REPOSITORY)["status"] == "unavailable" diff --git a/tests/interface/test_client_messages.py b/tests/interface/test_client_messages.py index afe3d52..37c1c0e 100644 --- a/tests/interface/test_client_messages.py +++ b/tests/interface/test_client_messages.py @@ -6,6 +6,7 @@ import pytest +from microcosm_emitter_interface.bundles import BundleFile, StageDatasetBundle from microcosm_emitter_interface.client import EmitterClient from microcosm_emitter_interface.launcher import EmitterProcess from microcosm_emitter_interface.orrery import OrrerySettings @@ -22,6 +23,27 @@ from microcosm_emitter_interface.telemetry import TelemetryRun, TelemetrySettings FIXTURE_STAGING_REPOSITORY = "policyengine/fixture-staging" +FIXTURE_PRIVATE_REPOSITORY = "policyengine/fixture-private" + + +def fixture_bundle(directory: Path) -> StageDatasetBundle: + """The bundle recorded in huggingface-stage-dataset-bundle.json.""" + fixture = fixtures_by_name()["huggingface-stage-dataset-bundle"] + recorded = materialize(fixture, directory=directory, spool_path=Path("unused")) + return StageDatasetBundle( + repository=FIXTURE_PRIVATE_REPOSITORY, + run_id="fixture-run", + files=tuple( + BundleFile( + name=entry["name"], + role=entry["role"], + local_path=Path(entry["local_path"]), + sha256=entry["sha256"], + bytes=entry["bytes"], + ) + for entry in recorded["files"] + ), + ) class RecordingTransport: @@ -37,7 +59,18 @@ def query(self, module, message, *, timeout=None): self.requests.append(("query", module, dict(message))) if module == "huggingface": return { - "receipt": {"status": "recorded", "pending_files": 0}, + "receipt": { + "status": "recorded", + "pending_files": 0, + "handover_id": "fixture-handover", + }, + "access": { + "status": "checked", + "credential": "present", + "reachable": True, + "can_write": True, + "error_type": None, + }, "read_back": { "status": "passed", "files": 0, @@ -176,6 +209,15 @@ def fixtures_by_name() -> dict[str, ProtocolFixture]: "huggingface-run-files-receipt": lambda client, directory: client.run_files_receipt( FIXTURE_STAGING_REPOSITORY, "fixture-run" ), + "huggingface-stage-dataset-bundle": lambda client, directory: ( + client.stage_dataset_bundle(fixture_bundle(directory)) + ), + "huggingface-handover-receipt": lambda client, directory: ( + client.dataset_bundle_receipt("fixture-handover") + ), + "huggingface-check-write-access": lambda client, directory: ( + client.check_write_access(FIXTURE_PRIVATE_REPOSITORY) + ), "huggingface-read-back-run-files": lambda client, directory: ( client.read_back_run_files( FIXTURE_STAGING_REPOSITORY, "fixture-run", wait_seconds=0 diff --git a/tests/interface/test_protocol_fixtures.py b/tests/interface/test_protocol_fixtures.py index a863a00..fcc9ff6 100644 --- a/tests/interface/test_protocol_fixtures.py +++ b/tests/interface/test_protocol_fixtures.py @@ -37,6 +37,15 @@ "protocol-v1/host-close.json": ( "68fa81b4116852f91f3c0dab3716a1974a8150039fb2b238f79efc094d9b0f6f" ), + "protocol-v1/huggingface-check-write-access.json": ( + "750b7eb8c7c9ea4112bb60560522d3b1591c6c9838239a5a54d72f90c095a631" + ), + "protocol-v1/huggingface-handover-receipt.json": ( + "e3d6735361da9326ce9c1aeb893d1c644a749e574ded8c1f2f2071f12e00ca09" + ), + "protocol-v1/huggingface-stage-dataset-bundle.json": ( + "6f576285b7a87607c47755ebee00ca051694b39455812a0068ebee2c27277728" + ), "protocol-v1/huggingface-configuration.json": ( "5a5ccd639953b24c774418e6e23735f515e46a673a68624e26d640839ef4b1c2" ), @@ -193,8 +202,9 @@ def test_orrery_module_answers_each_query_fixture_in_order(tmp_path): def test_hugging_face_module_answers_each_query_fixture_in_order(tmp_path): - hub = FakeHuggingFaceApi() + hub = FakeHuggingFaceApi(whoami={"auth": {"accessToken": {"role": "write"}}}) hub.create_repository("policyengine/fixture-staging") + hub.create_repository("policyengine/fixture-private") module = HuggingFaceModule( configuration_of("huggingface", tmp_path), ModuleContext(os.getpid()), @@ -211,10 +221,14 @@ def test_hugging_face_module_answers_each_query_fixture_in_order(tmp_path): materialize(fixture, directory=directory, spool_path=tmp_path) ) assert set(fixture.response_keys) <= set(response), fixture.name - if "receipt" in response: - assert response["receipt"]["status"] == "recorded", fixture.name + if response.get("receipt") is not None: + assert response["receipt"]["status"] in {"recorded", "pending"}, ( + fixture.name + ) if "read_back" in response: assert response["read_back"]["status"] in {"pending", "passed"} + if "access" in response: + assert response["access"]["status"] == "checked" finally: module.close(time.monotonic()) diff --git a/tests/interface/test_staging_client.py b/tests/interface/test_staging_client.py index 99024a8..6a985ab 100644 --- a/tests/interface/test_staging_client.py +++ b/tests/interface/test_staging_client.py @@ -1,77 +1,15 @@ """Typed staging-run-file hand-overs, through a real socket to the module.""" -import os -import shutil -import tempfile -import threading -import time -from dataclasses import dataclass, field from pathlib import Path import pytest -from fake_hugging_face import STAGING_REPOSITORY, FakeHuggingFaceApi +from fake_hugging_face import STAGING_REPOSITORY -from microcosm_emitter.host.contracts import ModuleContext -from microcosm_emitter.host.runtime import ServiceRuntime -from microcosm_emitter.huggingface.service.module import HuggingFaceModule from microcosm_emitter_interface.client import EmitterClient from microcosm_emitter_interface.staging import RunFile, StageRunFiles from microcosm_emitter_interface.transport import SocketTransport -@dataclass -class InThreadEmitter: - """An emitter host running in a thread of the test process.""" - - transport: SocketTransport - description: dict = field(default_factory=dict) - - def close(self): - self.transport.close() - - -@pytest.fixture -def emitter_with_hub(tmp_path): - hub = FakeHuggingFaceApi() - hub.create_repository(STAGING_REPOSITORY) - runtime_directory = Path(tempfile.mkdtemp(prefix="pe-stage-", dir="/tmp")) - socket_path = runtime_directory / "service.sock" - module = HuggingFaceModule( - {"spool_path": str(tmp_path / "events.sqlite3")}, - ModuleContext(os.getpid()), - api_factory=lambda: hub, - ) - runtime = ServiceRuntime( - socket_path, - modules={"huggingface": module}, - is_parent_alive=lambda: True, - tick_seconds=0.01, - ) - thread = threading.Thread(target=runtime.run) - thread.start() - transport = SocketTransport(socket_path) - deadline = time.monotonic() + 5 - while True: - try: - transport.ping() - break - except Exception: - if time.monotonic() > deadline: - raise - time.sleep(0.01) - client = EmitterClient( - InThreadEmitter( - transport, {"modules": ["huggingface"], "protocol_versions": [1]} - ) - ) - try: - yield client, hub - finally: - client.close() - thread.join(timeout=20) - shutil.rmtree(runtime_directory, ignore_errors=True) - - @pytest.mark.parametrize( "build_request", [ @@ -144,9 +82,11 @@ def test_message_names_each_file_by_its_absolute_local_path(tmp_path, monkeypatc } -def test_client_without_the_hugging_face_module_reports_unavailable(tmp_path): +def test_client_without_the_hugging_face_module_reports_unavailable( + tmp_path, in_thread_emitter +): client = EmitterClient( - InThreadEmitter( + in_thread_emitter( SocketTransport(tmp_path / "unused.sock"), {"modules": ["telemetry"], "protocol_versions": [1]}, ) @@ -160,9 +100,9 @@ def test_client_without_the_hugging_face_module_reports_unavailable(tmp_path): ] == ("unavailable") -def test_unreachable_emitter_never_raises(tmp_path): +def test_unreachable_emitter_never_raises(tmp_path, in_thread_emitter): client = EmitterClient( - InThreadEmitter( + in_thread_emitter( SocketTransport(tmp_path / "nothing-listens.sock"), {"modules": ["huggingface"], "protocol_versions": [1]}, ) From 6e5716ea17fe5f13129ec4b89f35f53d4ff3bc77 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:05:50 +0400 Subject: [PATCH 04/23] Build staging and bundle messages inside the client's error handling The staging calls still built their messages before the client's error handling: an unconvertible `run_index_entry`, a bundle file whose path cannot be resolved, or an invalid repository or run identifier given to `run_files_receipt`, `read_back_run_files` or `check_write_access` raised into the build. `stage_dataset_bundle`'s unavailable receipt also rebuilt the message, so it raised again. Each call now passes a message builder that the client runs inside its error handling, as the telemetry and Orrery calls do since the previous pull request; a failure returns the call's unavailable result and records `invalid_arguments` in `last_failure`. The unavailable bundle receipt has an empty `handover_id` when the files cannot be converted. Co-Authored-By: Claude Opus 5.5 --- .../microcosm_emitter_interface/bundles.py | 10 ++- .../src/microcosm_emitter_interface/client.py | 37 +++++----- tests/interface/test_client_errors.py | 69 +++++++++++++++++++ 3 files changed, 96 insertions(+), 20 deletions(-) diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/bundles.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/bundles.py index f679593..eb7c8e8 100644 --- a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/bundles.py +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/bundles.py @@ -195,8 +195,16 @@ class WriteAccess(TypedDict): def unavailable_bundle_receipt(request: StageDatasetBundle) -> DatasetBundleReceipt: + """No emitter accepted the bundle; built without converting its files. + + ``handover_id`` is empty when the request's files cannot be converted. + """ + try: + handover_id = request.handover_id() + except Exception: + handover_id = "" return { - "handover_id": request.handover_id(), + "handover_id": handover_id, "status": "unavailable", "error_code": None, "revision": None, diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py index b335fa3..ceed801 100644 --- a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py @@ -354,10 +354,9 @@ def stage_run_files(self, request: StageRunFiles) -> RunFilesReceipt: The emitter copies the files' bytes before answering, so the build may rewrite them immediately afterwards. """ - message = request.message() response = self._query( HUGGING_FACE_MODULE, - lambda: message, + request.message, kind="staging", timeout=STAGING_REQUEST_TIMEOUT_SECONDS, ) @@ -367,10 +366,9 @@ def stage_run_files(self, request: StageRunFiles) -> RunFilesReceipt: return unavailable_run_files_receipt(request.repository, request.run_id) def run_files_receipt(self, repository: str, run_id: str) -> RunFilesReceipt: - message = run_reference_message("run_files_receipt", repository, run_id) response = self._query( HUGGING_FACE_MODULE, - lambda: message, + lambda: run_reference_message("run_files_receipt", repository, run_id), kind="staging", timeout=STAGING_REQUEST_TIMEOUT_SECONDS, ) @@ -388,11 +386,15 @@ def read_back_run_files( whose files are still pending then reads back as ``pending``. """ if HUGGING_FACE_MODULE not in self.modules: + self._fail("staging", "emitter_unavailable") return unavailable_read_back() - receipt = self.stage_run_files( - StageRunFiles(repository=repository, run_id=run_id, flush=True) - ) - deadline = time.monotonic() + max(0, wait_seconds) + try: + flush = StageRunFiles(repository=repository, run_id=run_id, flush=True) + except Exception as error: + self._fail("staging", "invalid_arguments", error) + return unavailable_read_back() + receipt = self.stage_run_files(flush) + deadline = time.monotonic() + _wait_seconds(wait_seconds) while ( receipt["status"] == "recorded" and receipt["pending_files"] @@ -400,10 +402,9 @@ def read_back_run_files( ): time.sleep(0.1) receipt = self.run_files_receipt(repository, run_id) - message = run_reference_message("read_back_run_files", repository, run_id) response = self._query( HUGGING_FACE_MODULE, - lambda: message, + lambda: run_reference_message("read_back_run_files", repository, run_id), kind="staging", timeout=READ_BACK_TIMEOUT_SECONDS, ) @@ -429,7 +430,7 @@ def _handover_receipt(self, handover_id: str) -> dict[str, Any] | None: def _wait_until_settled( self, receipt: dict[str, Any], wait_seconds: float ) -> dict[str, Any]: - deadline = time.monotonic() + max(0, wait_seconds) + deadline = time.monotonic() + _wait_seconds(wait_seconds) while ( receipt.get("status") not in SETTLED_STATUSES and receipt.get("status") != "unavailable" @@ -450,10 +451,9 @@ def stage_dataset_bundle( Returns once the hand-over is durably queued (``pending``), or after it settles if that happens within ``wait_seconds``. """ - message = request.message() response = self._query( HUGGING_FACE_MODULE, - lambda: message, + request.message, kind="staging", timeout=FILE_REQUEST_TIMEOUT_SECONDS, ) @@ -467,14 +467,13 @@ def dataset_bundle_receipt(self, handover_id: str) -> DatasetBundleReceipt | Non def check_write_access(self, repository: str) -> WriteAccess: """Whether the build's credential can reach and write ``repository``.""" - message = { - "protocol_version": PROTOCOL_VERSION, - "command": "check_write_access", - "repository": repository_identifier(repository), - } response = self._query( HUGGING_FACE_MODULE, - lambda: message, + lambda: { + "protocol_version": PROTOCOL_VERSION, + "command": "check_write_access", + "repository": repository_identifier(repository), + }, kind="staging", timeout=READ_BACK_TIMEOUT_SECONDS, ) diff --git a/tests/interface/test_client_errors.py b/tests/interface/test_client_errors.py index d0ec57e..d297715 100644 --- a/tests/interface/test_client_errors.py +++ b/tests/interface/test_client_errors.py @@ -252,3 +252,72 @@ def test_unconvertible_spool_path_starts_an_unavailable_client(): ) assert not client.available assert client.startup_error.reason == "invalid_configuration" + + +def hugging_face_client() -> tuple[EmitterClient, RecordingTransport]: + transport = RecordingTransport() + process = EmitterProcess( + process=None, + transport=transport, + runtime_directory=Path("/nonexistent"), + description={"modules": ["huggingface"], "protocol_versions": [1]}, + ) + return EmitterClient(process), transport + + +@pytest.mark.parametrize("running", [True, False], ids=["running", "no-emitter"]) +def test_unconvertible_staging_arguments_report_failure(running): + from microcosm_emitter_interface.staging import StageRunFiles + + client, transport = hugging_face_client() if running else (EmitterClient(), None) + request = StageRunFiles( + repository="policyengine/fixture-staging", + run_id="run-1", + run_index_entry={"run_id": "run-1", "loss": MultiElementArray()}, + ) + receipt = client.stage_run_files(request) + assert receipt["status"] == "unavailable" + assert receipt["run_id"] == "run-1" + invalid_run = client.run_files_receipt("policyengine/fixture-staging", "../run") + assert invalid_run["status"] == "unavailable" + read_back = client.read_back_run_files("no-owner", "run-1", wait_seconds=0) + assert read_back["status"] == "unavailable" + access = client.check_write_access("no-owner") + assert access["status"] == "unavailable" + if running: + assert transport.requests == [] + assert client.last_failure.reason == "invalid_arguments" + + +@pytest.mark.parametrize("running", [True, False], ids=["running", "no-emitter"]) +def test_unconvertible_bundle_arguments_report_failure(running, tmp_path): + import hashlib + + from microcosm_emitter_interface.bundles import BundleFile, StageDatasetBundle + + client, transport = hugging_face_client() if running else (EmitterClient(), None) + + def bundle_file(name, role, local_path): + return BundleFile( + name=name, + role=role, + local_path=local_path, + sha256=hashlib.sha256(b"").hexdigest(), + bytes=0, + ) + + request = StageDatasetBundle( + repository="policyengine/fixture-staging", + run_id="run-1", + files=( + bundle_file("dataset.h5", "dataset", UnresolvablePath()), + bundle_file("staged_manifest.json", "staged_manifest", tmp_path / "m"), + ), + ) + receipt = client.stage_dataset_bundle(request, wait_seconds=0) + assert receipt["status"] == "unavailable" + assert receipt["repository"] == "policyengine/fixture-staging" + assert receipt["prefix"] == "staged/run-1" + if running: + assert transport.requests == [] + assert client.last_failure.reason == "invalid_arguments" From 788416106acf72385266bdc7c81da062323e6bd1 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:06:46 +0400 Subject: [PATCH 05/23] Start runs.json empty only when the Hub reports the file missing `commit_run_files` started `runs.json` empty whenever `hub.is_absent` matched the download error, and `is_absent` matched any error class named `EntryNotFoundError` or `FileNotFoundError` in the error's hierarchy. huggingface_hub's `LocalEntryNotFoundError` subclasses both, and `hf_hub_download` raises it for connection errors, timeouts and 429 or 503 responses when the file is not cached. A transient error therefore committed an index without every other run. `is_absent` now holds only for an entry-not-found error that carries the Hub's HTTP 404 response, and never for `LocalEntryNotFoundError`, a missing revision or repository, or a local `FileNotFoundError`. Any other error fails the commit, keeps the run files and index entry queued, and retries after the backoff. Tested with real huggingface_hub error classes. Co-Authored-By: Claude Opus 5.5 --- .../huggingface/service/hub.py | 28 ++++++++++----- .../huggingface/service/run_files.py | 3 ++ tests/huggingface/test_hub.py | 29 ++++++++++++++- tests/huggingface/test_run_files.py | 35 +++++++++++++++++++ 4 files changed, 85 insertions(+), 10 deletions(-) diff --git a/src/microcosm_emitter/huggingface/service/hub.py b/src/microcosm_emitter/huggingface/service/hub.py index 29382d5..6db29e6 100644 --- a/src/microcosm_emitter/huggingface/service/hub.py +++ b/src/microcosm_emitter/huggingface/service/hub.py @@ -17,14 +17,15 @@ #: HTTP statuses the Hub returns when ``parent_commit`` is no longer the head. PARENT_CONFLICT_STATUSES = frozenset({409, 412}) -_ABSENT_ERROR_NAMES = frozenset( - { - "EntryNotFoundError", - "RemoteEntryNotFoundError", - "RevisionNotFoundError", - "FileNotFoundError", - } +#: Errors huggingface_hub raises when the Hub answered that a file is missing +#: (``RemoteEntryNotFoundError`` from 1.0, ``EntryNotFoundError`` before). +_MISSING_FILE_ERROR_NAMES = frozenset( + {"EntryNotFoundError", "RemoteEntryNotFoundError"} ) +#: Raised when the Hub did not answer (a connection error, a timeout, or a 429 +#: or 5xx response for an uncached file). It subclasses ``EntryNotFoundError`` +#: and ``FileNotFoundError``, so a check by class alone mistakes it for absence. +_NO_ANSWER_ERROR_NAME = "LocalEntryNotFoundError" def hugging_face_api() -> Any | None: @@ -64,8 +65,17 @@ def is_parent_conflict(error: BaseException) -> bool: def is_absent(error: BaseException) -> bool: - """The Hub answered that the file or revision does not exist, nothing else.""" - return bool({cls.__name__ for cls in type(error).__mro__} & _ABSENT_ERROR_NAMES) + """The Hub answered HTTP 404 because the file does not exist, nothing else. + + Every other error, including a missing revision or repository and every + error raised without the Hub's answer, means the file's state is unknown. + """ + names = {cls.__name__ for cls in type(error).__mro__} + return ( + _NO_ANSWER_ERROR_NAME not in names + and bool(names & _MISSING_FILE_ERROR_NAMES) + and status_code(error) == 404 + ) class DatasetRepository: diff --git a/src/microcosm_emitter/huggingface/service/run_files.py b/src/microcosm_emitter/huggingface/service/run_files.py index 31f67f7..3204a97 100644 --- a/src/microcosm_emitter/huggingface/service/run_files.py +++ b/src/microcosm_emitter/huggingface/service/run_files.py @@ -129,6 +129,9 @@ def commit_run_files( RUNS_INDEX_PATH, revision=parent ) except Exception as error: + # Only a repository without runs.json starts a new index. Any + # other error retries the cycle: starting empty would drop + # every other run from the index. if not is_absent(error): raise existing = None diff --git a/tests/huggingface/test_hub.py b/tests/huggingface/test_hub.py index a1e7d0f..c60ba71 100644 --- a/tests/huggingface/test_hub.py +++ b/tests/huggingface/test_hub.py @@ -73,10 +73,37 @@ def test_downloads_read_exact_bytes_at_a_revision(hub): with pytest.raises(EntryNotFoundError) as missing: repository.download("missing.json") assert is_absent(missing.value) - assert is_absent(RevisionNotFoundError()) + # A missing revision is not a missing file: the cycle is retried. + assert not is_absent(RevisionNotFoundError()) assert not is_absent(FakeHubHTTPError(503)) +def real_http_error(error_class, status_code): + """A huggingface_hub HTTP error with a response; nothing is sent.""" + import httpx + + request = httpx.Request("GET", "https://hub.invalid/runs.json") + return error_class( + "synthetic", response=httpx.Response(status_code, request=request) + ) + + +def test_only_the_hub_answering_404_for_the_file_counts_as_absent(): + """``LocalEntryNotFoundError`` subclasses ``EntryNotFoundError`` and + ``FileNotFoundError``, but huggingface_hub raises it when the Hub did not + answer: a connection error, a timeout or a 429 or 5xx response for a file + that is not in the local cache.""" + from huggingface_hub import errors + + assert is_absent(real_http_error(errors.RemoteEntryNotFoundError, 404)) + assert not is_absent(errors.LocalEntryNotFoundError("connection error")) + assert not is_absent(real_http_error(errors.RevisionNotFoundError, 404)) + assert not is_absent(real_http_error(errors.RepositoryNotFoundError, 404)) + assert not is_absent(real_http_error(errors.HfHubHTTPError, 503)) + assert not is_absent(real_http_error(errors.HfHubHTTPError, 429)) + assert not is_absent(FileNotFoundError("local cache file vanished")) + + def test_last_commit_names_the_commit_that_last_changed_a_path(hub): repository = DatasetRepository(REPOSITORY, hub) changed = hub.write_directly(REPOSITORY, {"a.json": b"1"}) diff --git a/tests/huggingface/test_run_files.py b/tests/huggingface/test_run_files.py index ae9687e..6075dbd 100644 --- a/tests/huggingface/test_run_files.py +++ b/tests/huggingface/test_run_files.py @@ -212,6 +212,41 @@ def test_staging_v1_runs_index_starts_empty_without_a_repository_copy( assert json.loads(commit.added["runs.json"])["runs"] == [entry] +@pytest.mark.parametrize( + "transient_error", + ["connection error", "rate limited", "service unavailable"], +) +def test_a_transient_download_error_never_starts_the_runs_index_empty( + hugging_face_module, transient_error +): + """huggingface_hub reports each of these as ``LocalEntryNotFoundError``.""" + from huggingface_hub.errors import LocalEntryNotFoundError + + controls = hugging_face_module + other_run = {"run_id": "run-older", "updated_at": "2026-10-01T00:00:00+00:00"} + controls.hub.write_directly( + STAGING_REPOSITORY, + { + "runs.json": json.dumps( + {"schema_version": 1, "updated_at": "x", "runs": [other_run]}, indent=1 + ).encode() + }, + ) + before = len(commits(controls)) + controls.hub.fail_next("hf_hub_download", LocalEntryNotFoundError(transient_error)) + entry = {"run_id": "us-run", "updated_at": "2026-10-20T00:00:00+00:00"} + hand_over(controls, "us-run", {"progress.json": b"{}"}, run_index_entry=entry) + cycle(controls) + assert len(commits(controls)) == before + receipt = controls.module.queue.run_receipt(STAGING_REPOSITORY, "us-run") + assert receipt["last_error_code"] == "UPLOAD_FAILED" + assert receipt["pending_files"] == 2 + controls.clock.advance(2) + cycle(controls) + index = json.loads(controls.hub.files(STAGING_REPOSITORY)["runs.json"]) + assert index["runs"] == [entry, other_run] + + def test_runs_index_written_by_another_build_meanwhile_is_kept(hugging_face_module): controls = hugging_face_module concurrent_run = {"run_id": "concurrent", "updated_at": "2026-10-19T00:00:00Z"} From d0ec0db0733b3a9bb854f03188f677f75ba477ec Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:10:42 +0400 Subject: [PATCH 06/23] Copy a bundle's small files into the queue when it is handed over Microcosm's `full_build_cli` (at 41209b38c) rewrites three bundle files right after `stage_dataset` returns: the build manifest (`replace_manifest`), the checksum listing (`refresh_sha256sums_entry`) and, in national builds, the build record. With the queued upload the worker re-hashed those files and failed the hand-over permanently with `PRESERVED_BYTES_CHANGED`. - Before answering, the module copies the bundle's files smallest first, each of at most 32 MiB and together at most 64 MiB, after checking that each holds its declared size and SHA-256 (`FILE_SIZE_DIFFERS`, `FILE_DIGEST_DIFFERS`). The copies are stored with the job in a new `handover_files` table (added to the unreleased revision `20261020_03`) and named in the receipt's `copied_files`. - The worker uploads the copied bytes as they were handed over, and still re-hashes and refuses changed bytes of the larger files it uploads from their paths (the datasets, which Microcosm does not rewrite). - The copies are deleted once every destination of the hand-over has settled. A test follows Microcosm's order: hand over, rewrite the manifest atomically, append to the checksum listing and rewrite the build record, then run the cycle; the commit carries the handed-over bytes. Co-Authored-By: Claude Opus 5.5 --- README.md | 3 +- .../huggingface-dataset-bundles.added.md | 2 +- .../microcosm_emitter_interface/bundles.py | 14 ++- .../versions/v20261020_03_handover_queue.py | 8 ++ src/microcosm_emitter/core/models.py | 16 +++ .../huggingface/service/bundles.py | 19 ++- .../huggingface/service/delivery.py | 3 +- .../huggingface/service/module.py | 3 + .../huggingface/service/queue.py | 53 +++++++- .../huggingface/service/requests.py | 43 ++++++- tests/huggingface/test_bundles.py | 118 +++++++++++++++++- tests/telemetry/test_persistence.py | 1 + 12 files changed, 264 insertions(+), 19 deletions(-) diff --git a/README.md b/README.md index b79f8cb..d3c59ed 100644 --- a/README.md +++ b/README.md @@ -69,7 +69,8 @@ The `huggingface` module takes over the uploads Microcosm made itself. Microcosm **Staged dataset bundles** (Microcosm's `staging_dataset.stage_bundle`, used by `uk_runtime/rowwise_staging.py` and the Orrery evidence bundle). A build verifies its bundle and writes the sidecars as before, then hands the files over by path, each with its role (`dataset`, `calibration_diagnostics`, `build_manifest`, `staged_manifest`, `checksums`, ...), SHA-256 and size. The hand-over is a durable job: -- The worker re-hashes each preserved file and never uploads changed bytes (`PRESERVED_BYTES_CHANGED`); a missing file waits for an operator to restore it (`PRESERVED_FILES_MISSING`). +- Before answering, the module copies the bundle's small files into the spool database, smallest first, each of at most 32 MiB and together at most 64 MiB, after checking each one's size and SHA-256 against the hand-over (`FILE_SIZE_DIFFERS`, `FILE_DIGEST_DIFFERS`). The receipt's `copied_files` names them, and the build may rewrite them at once: Microcosm rewrites its build manifest, its `sha256sums.txt` and, in national builds, its build record right after staging, and the commit carries the bytes as they were handed over. The copies are deleted once the hand-over settles. +- Larger files (the datasets) are uploaded from their paths. The worker re-hashes each one first and never uploads changed bytes (`PRESERVED_BYTES_CHANGED`); a missing file waits for an operator to restore it (`PRESERVED_FILES_MISSING`). - As in Microcosm, it writes every file under `//` (by default `staged//`) in one commit pinned to the head it checked, with the message `Stage under /`. A repository copy of the bundle's `staged_manifest` whose `files` equal the outputs' digests is `already_staged` with that copy's commit; a different one is `REMOTE_DIFFERS` and is never overwritten. - Unlike Microcosm's single attempt, `REPOSITORY_UNAVAILABLE` and `UPLOAD_FAILED` are retried with backoff, and a commit that loses a race to another writer is checked and tried again. The receipt keeps the last error code while the job is pending. - `check_write_access` is Microcosm's staged-dataset preflight: whether a credential is present, the repository is reachable, and the Hub reports a write scope for it. diff --git a/changelog.d/huggingface-dataset-bundles.added.md b/changelog.d/huggingface-dataset-bundles.added.md index 723e519..fb557cf 100644 --- a/changelog.d/huggingface-dataset-bundles.added.md +++ b/changelog.d/huggingface-dataset-bundles.added.md @@ -1 +1 @@ -Add `stage_dataset_bundle`, `dataset_bundle_receipt` and `check_write_access`. A build hands over its staged dataset bundle with each file's role; the emitter writes it under `staged//` in one commit with Microcosm's message, paths and idempotency (`already_staged`, `REMOTE_DIFFERS`), retries Hub failures from the durable queue, and never uploads bytes that changed after the hand-over. +Add `stage_dataset_bundle`, `dataset_bundle_receipt` and `check_write_access`. A build hands over its staged dataset bundle with each file's role; the emitter writes it under `staged//` in one commit with Microcosm's message, paths and idempotency (`already_staged`, `REMOTE_DIFFERS`), and retries Hub failures from the durable queue. The emitter copies the bundle's small files (each up to 32 MiB, up to 64 MiB per hand-over) into its queue before answering, so the build may rewrite them at once, as Microcosm does with its build manifest, checksum listing and build record; it uploads larger files from their paths and never uploads one whose bytes changed after the hand-over. diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/bundles.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/bundles.py index eb7c8e8..66e7cb9 100644 --- a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/bundles.py +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/bundles.py @@ -8,6 +8,15 @@ whose ``files`` equal the outputs' digests is ``already_staged``, and a different one is refused as ``REMOTE_DIFFERS``. +Before answering, the emitter copies the bundle's small files into its queue, +smallest first, each of at most 32 MiB and together at most 64 MiB, after +checking that each holds exactly its declared size and SHA-256 +(``FILE_SIZE_DIFFERS`` or ``FILE_DIGEST_DIFFERS`` otherwise). The receipt's +``copied_files`` names them. Once the call has returned, the build may +rewrite a copied file: the emitter uploads the bytes it copied. Every other +file is uploaded from its path, re-hashed first, and must stay unchanged +until the hand-over settles (``PRESERVED_BYTES_CHANGED`` otherwise). + Constructing a request that breaks the rules raises ``ValueError``. """ @@ -19,7 +28,7 @@ from collections.abc import Sequence from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Final, Literal, TypedDict +from typing import Any, Final, Literal, NotRequired, TypedDict from microcosm_emitter_interface.identifiers import ( bare_file_name, @@ -177,6 +186,8 @@ class DatasetBundleReceipt(TypedDict): attempts: int repository: str prefix: str + #: The files the emitter copied at hand-over; the build may rewrite them. + copied_files: NotRequired[list[str]] class WriteAccess(TypedDict): @@ -211,6 +222,7 @@ def unavailable_bundle_receipt(request: StageDatasetBundle) -> DatasetBundleRece "attempts": 0, "repository": request.repository, "prefix": request.remote_prefix, + "copied_files": [], } diff --git a/src/microcosm_emitter/core/alembic/versions/v20261020_03_handover_queue.py b/src/microcosm_emitter/core/alembic/versions/v20261020_03_handover_queue.py index 6abfa3f..b6bd30f 100644 --- a/src/microcosm_emitter/core/alembic/versions/v20261020_03_handover_queue.py +++ b/src/microcosm_emitter/core/alembic/versions/v20261020_03_handover_queue.py @@ -59,9 +59,17 @@ def upgrade() -> None: sa.Column("receipt", sa.Text(), nullable=False), sa.Column("created_at", sa.Float(), nullable=False), ) + op.create_table( + "handover_files", + sa.Column("handover_id", sa.Text(), primary_key=True), + sa.Column("name", sa.Text(), primary_key=True), + sa.Column("content", sa.LargeBinary(), nullable=False), + sa.Column("content_sha256", sa.Text(), nullable=False), + ) def downgrade() -> None: + op.drop_table("handover_files") op.drop_table("handover_jobs") op.drop_table("repository_commit_schedules") op.drop_table("staged_run_deliveries") diff --git a/src/microcosm_emitter/core/models.py b/src/microcosm_emitter/core/models.py index 5472756..f4380c5 100644 --- a/src/microcosm_emitter/core/models.py +++ b/src/microcosm_emitter/core/models.py @@ -215,3 +215,19 @@ class HandoverJob(SpoolModel): lease_until: Mapped[float] = mapped_column(Float, nullable=False) receipt: Mapped[dict[str, Any]] = mapped_column(JsonObjectText, nullable=False) created_at: Mapped[float] = mapped_column(Float, nullable=False) + + +class HandoverFile(SpoolModel): + """The bytes of one small hand-over file, copied when it was handed over. + + The build may rewrite such a file at once (Microcosm rewrites its build + manifest, checksum listing and build record right after staging), so the + delivery uploads these bytes instead of reading the file again. The rows + are deleted once every destination of the hand-over has settled. + """ + + __tablename__ = "handover_files" + handover_id: Mapped[str] = mapped_column(Text, primary_key=True) + name: Mapped[str] = mapped_column(Text, primary_key=True) + content: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + content_sha256: Mapped[str] = mapped_column(Text, nullable=False) diff --git a/src/microcosm_emitter/huggingface/service/bundles.py b/src/microcosm_emitter/huggingface/service/bundles.py index b6c0047..1242259 100644 --- a/src/microcosm_emitter/huggingface/service/bundles.py +++ b/src/microcosm_emitter/huggingface/service/bundles.py @@ -78,11 +78,21 @@ def _local_files_outcome(files: list[Mapping[str, Any]]) -> BundleOutcome | None def stage_bundle( - repository: DatasetRepository, request: Mapping[str, Any] + repository: DatasetRepository, + request: Mapping[str, Any], + copied: Mapping[str, bytes] | None = None, ) -> BundleOutcome: - """Write the bundle once, or report why not; never raises for the Hub.""" + """Write the bundle once, or report why not; never raises for the Hub. + + ``copied`` holds the bytes of the files copied at hand-over; they are + uploaded as they were then. Every other file is re-hashed from its path + first and never uploaded if its bytes changed. + """ + copied = {} if copied is None else dict(copied) files = list(request["files"]) - local = _local_files_outcome(files) + local = _local_files_outcome( + [entry for entry in files if entry["name"] not in copied] + ) if local is not None: return local remote_prefix = f"{request['prefix']}/{request['run_id']}" @@ -121,7 +131,8 @@ def stage_bundle( revision = repository.commit( [ add_operation( - f"{remote_prefix}/{entry['name']}", Path(entry["local_path"]) + f"{remote_prefix}/{entry['name']}", + copied.get(entry["name"], Path(entry["local_path"])), ) for entry in files ], diff --git a/src/microcosm_emitter/huggingface/service/delivery.py b/src/microcosm_emitter/huggingface/service/delivery.py index 8c98395..57b5f27 100644 --- a/src/microcosm_emitter/huggingface/service/delivery.py +++ b/src/microcosm_emitter/huggingface/service/delivery.py @@ -142,6 +142,7 @@ def _finish_job( "attempts": job.attempts + 1, "repository": job.target, "prefix": f"{job.request['prefix']}/{job.request['run_id']}", + "copied_files": sorted(job.files), }, retry=retry, now=self.clock(), @@ -149,7 +150,7 @@ def _finish_job( ) def _deliver_job(self, hub: DatasetRepository, job: ClaimedJob) -> bool: - outcome = stage_bundle(hub, job.request) + outcome = stage_bundle(hub, job.request, job.files) if outcome.retry and outcome.error_code is not None: self.warn_once( outcome.error_code, diff --git a/src/microcosm_emitter/huggingface/service/module.py b/src/microcosm_emitter/huggingface/service/module.py index 30ef45d..a8edae8 100644 --- a/src/microcosm_emitter/huggingface/service/module.py +++ b/src/microcosm_emitter/huggingface/service/module.py @@ -121,6 +121,7 @@ def _stage_dataset_bundle(self, message: JsonObject) -> dict[str, Any]: "attempts": 0, "repository": repository if isinstance(repository, str) else "", "prefix": "", + "copied_files": [], } return self.queue.enqueue_job( handover_id=hand_over.handover_id, @@ -135,8 +136,10 @@ def _stage_dataset_bundle(self, message: JsonObject) -> dict[str, Any]: "attempts": 0, "repository": hand_over.repository, "prefix": hand_over.remote_prefix, + "copied_files": sorted(hand_over.copied), }, now=self.clock(), + files=hand_over.copied, ) def _handover_receipt(self, message: JsonObject) -> dict[str, Any] | None: diff --git a/src/microcosm_emitter/huggingface/service/queue.py b/src/microcosm_emitter/huggingface/service/queue.py index 4797821..0d7bf3a 100644 --- a/src/microcosm_emitter/huggingface/service/queue.py +++ b/src/microcosm_emitter/huggingface/service/queue.py @@ -21,11 +21,11 @@ import json import time from collections.abc import Iterable, Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import Any -from sqlalchemy import select, update +from sqlalchemy import delete, select, update from sqlalchemy.orm import Session from microcosm_emitter.core.database import ( @@ -34,6 +34,7 @@ ) from microcosm_emitter.core.migrations import upgrade_spool_database from microcosm_emitter.core.models import ( + HandoverFile, HandoverJob, RepositoryCommitSchedule, StagedRunDelivery, @@ -86,6 +87,8 @@ class ClaimedJob: request: dict[str, Any] attempts: int lease_until: float + #: The bytes of the files copied at hand-over, by name. + files: dict[str, bytes] = field(default_factory=dict) class HandoverQueue: @@ -492,8 +495,13 @@ def enqueue_job( receipt: Mapping[str, Any], now: float, destination: str = DESTINATION_HUGGING_FACE, + files: Mapping[str, bytes] | None = None, ) -> dict[str, Any]: - """Durably record a hand-over once; repeating it returns its receipt.""" + """Durably record a hand-over once; repeating it returns its receipt. + + ``files`` are the bytes copied at hand-over, stored with the job in + one transaction and kept until every destination has settled. + """ with self._sessions.begin() as session: existing = session.get(HandoverJob, (handover_id, destination)) if existing is not None: @@ -521,8 +529,26 @@ def enqueue_job( created_at=now, ) ) + for name, content in sorted((files or {}).items()): + if session.get(HandoverFile, (handover_id, name)) is None: + session.add( + HandoverFile( + handover_id=handover_id, + name=name, + content=content, + content_sha256=sha256_bytes(content), + ) + ) return dict(receipt) + def copied_files(self, handover_id: str) -> dict[str, bytes]: + """The bytes copied when ``handover_id`` was handed over, by name.""" + with self._sessions() as session: + rows = session.scalars( + select(HandoverFile).where(HandoverFile.handover_id == handover_id) + ) + return {row.name: row.content for row in rows} + def claim_job( self, target: str, @@ -569,6 +595,11 @@ def claim_job( if claimed is None: return None session.refresh(candidate) + copied = session.scalars( + select(HandoverFile).where( + HandoverFile.handover_id == candidate.handover_id + ) + ) return ClaimedJob( handover_id=candidate.handover_id, destination=destination, @@ -577,6 +608,7 @@ def claim_job( request=dict(candidate.request), attempts=candidate.attempts, lease_until=lease, + files={row.name: row.content for row in copied}, ) def finish_job( @@ -608,6 +640,21 @@ def finish_job( row.next_attempt_at = now + retry_delay(row.attempts) else: row.status = str(receipt["status"]) + still_pending = session.scalar( + select(HandoverJob.destination) + .where( + HandoverJob.handover_id == job.handover_id, + HandoverJob.status == JOB_PENDING, + ) + .limit(1) + ) + if still_pending is None: + # Every destination has settled: the copies are not needed. + session.execute( + delete(HandoverFile).where( + HandoverFile.handover_id == job.handover_id + ) + ) if committed: self._schedule(session, job.target).last_commit_at = now diff --git a/src/microcosm_emitter/huggingface/service/requests.py b/src/microcosm_emitter/huggingface/service/requests.py index 5b7e13b..8799ccf 100644 --- a/src/microcosm_emitter/huggingface/service/requests.py +++ b/src/microcosm_emitter/huggingface/service/requests.py @@ -6,9 +6,10 @@ from __future__ import annotations +import hashlib import math -from collections.abc import Mapping -from dataclasses import dataclass +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -31,6 +32,11 @@ MAX_HAND_OVER_BYTES = 64 * 1024 * 1024 #: Staging contract version 1's pointer to the latest run, at the root. LATEST_STAGING_POINTER_PATH = "latest_staging.json" +#: Bundle files are copied into the queue when they are handed over, smallest +#: first, each of at most ``MAX_COPIED_FILE_BYTES`` and together at most +#: ``MAX_COPIED_BYTES``; larger files are uploaded from their paths. +MAX_COPIED_FILE_BYTES = 32 * 1024 * 1024 +MAX_COPIED_BYTES = 64 * 1024 * 1024 class RequestRefusedError(ValueError): @@ -154,13 +160,40 @@ class BundleHandOver: repository: str remote_prefix: str request: dict[str, Any] + #: The bytes of the files copied at hand-over, by name. + copied: dict[str, bytes] = field(default_factory=dict) + + +def copy_small_files(files: Sequence[BundleFile]) -> dict[str, bytes]: + """Read the smallest files within the copy limits; check each digest. + + A copied file must hold exactly its declared bytes now, since the build + may rewrite it as soon as the hand-over is answered. + """ + copied: dict[str, bytes] = {} + total = 0 + for item in sorted(files, key=lambda item: (item.bytes, item.name)): + if item.bytes > MAX_COPIED_FILE_BYTES or total + item.bytes > MAX_COPIED_BYTES: + break # Every later file is at least as large. + try: + content = Path(item.local_path).read_bytes() + except OSError: + raise RequestRefusedError("LOCAL_FILE_MISSING") from None + if len(content) != item.bytes: + raise RequestRefusedError("FILE_SIZE_DIFFERS") + if hashlib.sha256(content).hexdigest() != item.sha256: + raise RequestRefusedError("FILE_DIGEST_DIFFERS") + copied[item.name] = content + total += item.bytes + return copied def parse_stage_dataset_bundle(message: Mapping[str, Any]) -> BundleHandOver: """Check a bundle with the interface's rules and the files' sizes. - Hashing large files waits for the worker; here only existence and size - are checked, so the request answers quickly. + Small files are copied and their digests checked now (``copy_small_files``); + hashing the large ones waits for the worker, so the request answers + quickly. """ require_supported_version(message) entries = message.get("files") @@ -198,10 +231,12 @@ def parse_stage_dataset_bundle(message: Mapping[str, Any]) -> BundleHandOver: raise RequestRefusedError("FILE_SIZE_DIFFERS") except OSError: raise RequestRefusedError("LOCAL_FILE_MISSING") from None + copied = copy_small_files(request.files) canonical = dict(message) return BundleHandOver( handover_id=handover_identifier("dataset_bundle", canonical), repository=request.repository, remote_prefix=request.remote_prefix, request={key: value for key, value in canonical.items()}, + copied=copied, ) diff --git a/tests/huggingface/test_bundles.py b/tests/huggingface/test_bundles.py index 9bc01d2..5822942 100644 --- a/tests/huggingface/test_bundles.py +++ b/tests/huggingface/test_bundles.py @@ -2,6 +2,7 @@ import hashlib import json +from pathlib import Path import pytest from fake_hugging_face import ( @@ -11,6 +12,7 @@ ) from microcosm_emitter.host.contracts import ModuleContext +from microcosm_emitter.huggingface.service import requests as request_module from microcosm_emitter.huggingface.service.module import HuggingFaceModule from microcosm_emitter_interface.bundles import BundleFile, StageDatasetBundle @@ -18,8 +20,13 @@ DIAGNOSTICS = b'{"schema_version":8}' -def bundle_files(directory, *, dataset=DATASET, manifest_note="first"): - """A bundle laid out as Microcosm's write_sidecars leaves it.""" +def bundle_files( + directory, *, dataset=DATASET, manifest_note="first", build_record=None +): + """A bundle laid out as Microcosm's write_sidecars leaves it. + + ``build_record`` adds the national build's record as an output. + """ directory.mkdir(parents=True, exist_ok=True) outputs = { "microcosm_uk_2024_25.h5": ("dataset", dataset), @@ -28,6 +35,8 @@ def bundle_files(directory, *, dataset=DATASET, manifest_note="first"): DIAGNOSTICS, ), } + if build_record is not None: + outputs["build_record.json"] = ("build_record", build_record) digests = { name: {"sha256": hashlib.sha256(content).hexdigest(), "bytes": len(content)} for name, (_, content) in outputs.items() @@ -191,8 +200,14 @@ def test_a_commit_that_loses_a_race_is_retried_on_the_new_head(hugging_face_modu assert len(ours) == 1 +@pytest.fixture +def nothing_copied(monkeypatch): + """Every file is too large to copy, so each is uploaded from its path.""" + monkeypatch.setattr(request_module, "MAX_COPIED_FILE_BYTES", 8) + + def test_local_bytes_changed_after_the_hand_over_are_never_uploaded( - hugging_face_module, + hugging_face_module, nothing_copied ): controls = hugging_face_module request, _, _ = stage_bundle(controls) @@ -207,7 +222,9 @@ def test_local_bytes_changed_after_the_hand_over_are_never_uploaded( ) -def test_missing_local_files_stay_queued_until_restored(hugging_face_module): +def test_missing_local_files_stay_queued_until_restored( + hugging_face_module, nothing_copied +): controls = hugging_face_module request, _, contents = stage_bundle(controls) dataset = next(item for item in request.files if item.role == "dataset") @@ -224,6 +241,99 @@ def test_missing_local_files_stay_queued_until_restored(hugging_face_module): assert receipt_of(controls, request.handover_id())["status"] == "uploaded" +LARGE_DATASET = bytes(range(256)) * 16 # 4,096 bytes: above the test copy limit + + +def rewrite_atomically(path, content): + """Microcosm's ``replace_manifest``: a new file renamed over the old one.""" + temporary = path.with_name(f".{path.name}.rewrite") + temporary.write_bytes(content) + temporary.replace(path) + + +def test_files_microcosm_rewrites_after_staging_are_uploaded_as_handed_over( + hugging_face_module, monkeypatch +): + """Microcosm's ``full_build_cli`` order (at ``41209b38c``). + + ``stage_dataset`` hands the bundle over and returns; the close step then + rewrites the build manifest (``replace_manifest``), the checksum listing + (``refresh_sha256sums_entry``) and, in national builds, the build record, + all before the worker's next cycle. The dataset stays unchanged. + """ + monkeypatch.setattr(request_module, "MAX_COPIED_FILE_BYTES", 1024) + controls = hugging_face_module + request, receipt, contents = stage_bundle( + controls, dataset=LARGE_DATASET, build_record=b'{"build":"record"}' + ) + assert receipt["status"] == "pending" + assert receipt["copied_files"] == sorted( + name for name, content in contents.items() if len(content) <= 1024 + ) + assert "microcosm_uk_2024_25.h5" not in receipt["copied_files"] + directory = controls.tmp_path / "uk-run-1" + rewrite_atomically( + directory / "rowwise_candidate_manifest.json", b'{"staging_delivery":{}}' + ) + rewrite_atomically(directory / "build_record.json", b'{"staging_delivery":{}}') + with (directory / "sha256sums.txt").open("ab") as listing: + listing.write(b"0" * 64 + b" build_record.json\n") + cycle(controls) + (commit,) = private_commits(controls) + assert commit.added == { + f"staged/uk-run-1/{name}": content for name, content in contents.items() + } + settled = receipt_of(controls, request.handover_id()) + assert (settled["status"], settled["error_code"]) == ("uploaded", None) + assert settled["copied_files"] == receipt["copied_files"] + + +def test_copied_bytes_are_deleted_once_the_hand_over_settles(hugging_face_module): + controls = hugging_face_module + request, receipt, contents = stage_bundle(controls) + queue = controls.module.queue + assert queue.copied_files(request.handover_id()) == { + name: contents[name] for name in receipt["copied_files"] + } + cycle(controls) + assert receipt_of(controls, request.handover_id())["status"] == "uploaded" + assert queue.copied_files(request.handover_id()) == {} + + +def test_files_are_copied_smallest_first_within_both_limits( + hugging_face_module, monkeypatch +): + controls = hugging_face_module + files, contents = bundle_files(controls.tmp_path / "limits") + sizes = sorted((len(content), name) for name, content in contents.items()) + # Room for the two smallest files only. + monkeypatch.setattr(request_module, "MAX_COPIED_BYTES", sizes[0][0] + sizes[1][0]) + request = StageDatasetBundle( + repository=PRIVATE_REPOSITORY, run_id="uk-run-1", files=files + ) + receipt = controls.module.query(request.message())["receipt"] + assert receipt["copied_files"] == sorted(name for _, name in sizes[:2]) + + +def test_a_copied_file_unlike_its_declared_digest_is_refused(hugging_face_module): + controls = hugging_face_module + files, _ = bundle_files(controls.tmp_path / "bundle") + message = StageDatasetBundle( + repository=PRIVATE_REPOSITORY, run_id="uk-run-1", files=files + ).message() + checksums = next( + entry for entry in message["files"] if entry["role"] == "checksums" + ) + Path(checksums["local_path"]).write_bytes(b"X" * checksums["bytes"]) + receipt = controls.module.query(message)["receipt"] + assert (receipt["status"], receipt["error_code"]) == ( + "refused", + "FILE_DIGEST_DIFFERS", + ) + cycle(controls) + assert private_commits(controls) == [] + + def test_repeating_a_hand_over_queues_it_once(hugging_face_module): controls = hugging_face_module request, first, _ = stage_bundle(controls) diff --git a/tests/telemetry/test_persistence.py b/tests/telemetry/test_persistence.py index a850847..5161166 100644 --- a/tests/telemetry/test_persistence.py +++ b/tests/telemetry/test_persistence.py @@ -87,5 +87,6 @@ def test_multiple_producers_and_local_only_survive_reopening(tmp_path): "staged_run_deliveries", "repository_commit_schedules", "handover_jobs", + "handover_files", } spool.close() From 26face1075fbd20df59d832dbc2cf76f3e73a918 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:11:26 +0400 Subject: [PATCH 07/23] Never pass a read-back of a run with no committed files `_read_back_run_files` returned `passed` with `files: 0` for a run without committed rows, for example after a refused hand-over or when the run was staged through another spool database. Microcosm's `validate_staging_delivery` rejects a passed read-back with zero uploads. Such a read-back is now `failed` with the error code `NO_COMMITTED_FILES`. Co-Authored-By: Claude Opus 5.5 --- README.md | 2 +- .../microcosm_emitter_interface/staging.py | 7 ++- .../huggingface/service/module.py | 8 ++- tests/huggingface/test_run_files.py | 50 +++++++++++++++++++ 4 files changed, 64 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index d3c59ed..91dc56c 100644 --- a/README.md +++ b/README.md @@ -65,7 +65,7 @@ The `huggingface` module takes over the uploads Microcosm made itself. Microcosm - **One commit per repository per cycle.** Each module tick commits, for each repository with work due, every pending path in one commit. An unforced hand-over is due one interval after the repository's last commit (Microcosm's upload interval); a forced one (final states, artifacts) at the next cycle, but never sooner than 10 seconds after the previous commit. Microcosm's writers made one commit per file, up to eight per upload cycle, and each commit triggered the dataset repository's webhook. - Uploads stay best-effort. Requests answer within seconds or report that they were not accepted; a failed commit keeps the files queued and retries with backoff (up to five minutes); without a credential the files stay queued and one warning is printed. On shutdown the module commits pending files once more. Nothing raises into the build. - A commit that merges `runs.json` is pinned to the head it read (`parent_commit`); if another build commits first, the module reads the index again and merges again. -- `read_back_run_files` commits pending files, then compares the Hub's copy of each committed file with the committed bytes. +- `read_back_run_files` commits pending files, then compares the Hub's copy of each committed file with the committed bytes. A run with no file this emitter committed (its hand-overs were refused, or another machine's emitter staged it) reads back as `failed` with `NO_COMMITTED_FILES`, never `passed`. **Staged dataset bundles** (Microcosm's `staging_dataset.stage_bundle`, used by `uk_runtime/rowwise_staging.py` and the Orrery evidence bundle). A build verifies its bundle and writes the sidecars as before, then hands the files over by path, each with its role (`dataset`, `calibration_diagnostics`, `build_manifest`, `staged_manifest`, `checksums`, ...), SHA-256 and size. The hand-over is a durable job: diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/staging.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/staging.py index 06c4a19..121a502 100644 --- a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/staging.py +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/staging.py @@ -160,7 +160,12 @@ class RunFilesReceipt(TypedDict): class ReadBackResult(TypedDict): - """Whether the Hub holds exactly the bytes the emitter committed.""" + """Whether the Hub holds exactly the bytes the emitter committed. + + ``passed`` needs at least one committed file. A run with none (its + hand-overs were refused, or another machine's emitter staged it) reads + back as ``failed`` with ``error_code`` ``NO_COMMITTED_FILES``. + """ status: Literal["passed", "failed", "pending", "refused", "unavailable"] files: int diff --git a/src/microcosm_emitter/huggingface/service/module.py b/src/microcosm_emitter/huggingface/service/module.py index a8edae8..73b3197 100644 --- a/src/microcosm_emitter/huggingface/service/module.py +++ b/src/microcosm_emitter/huggingface/service/module.py @@ -31,6 +31,8 @@ ) from microcosm_emitter_interface.identifiers import repository_identifier +#: A read-back of a run with no file this emitter committed: nothing to compare. +NO_COMMITTED_FILES = "NO_COMMITTED_FILES" WORKER_ERROR_MESSAGE = ( "Microcosm Hugging Face delivery hit a local queue error; handed-over files " "stay queued for a later attempt." @@ -211,10 +213,14 @@ def _read_back_run_files(self, message: JsonObject) -> dict[str, Any]: return {**result, "error_code": refusal.code} if self.queue.run_receipt(repository, run_id)["pending_files"]: return {**result, "status": "pending"} + committed = self.queue.committed_run_files(repository, run_id) + if not committed: + # A refused hand-over, or a run this database never held: passing + # would claim a delivery that this emitter cannot show. + return {**result, "status": "failed", "error_code": NO_COMMITTED_FILES} api = self.api_factory() if api is None: return {**result, "status": "failed", "error_code": CREDENTIAL_MISSING} - committed = self.queue.committed_run_files(repository, run_id) hub = DatasetRepository(repository, api) mismatched = [] try: diff --git a/tests/huggingface/test_run_files.py b/tests/huggingface/test_run_files.py index 6075dbd..6b540b0 100644 --- a/tests/huggingface/test_run_files.py +++ b/tests/huggingface/test_run_files.py @@ -365,6 +365,56 @@ def test_read_back_compares_the_hub_with_the_committed_bytes(hugging_face_module } +def read_back(controls, run_id): + return controls.module.query( + { + "protocol_version": 1, + "command": "read_back_run_files", + "repository": STAGING_REPOSITORY, + "run_id": run_id, + } + )["read_back"] + + +NOTHING_COMMITTED = { + "status": "failed", + "files": 0, + "mismatched_paths": [], + "error_code": "NO_COMMITTED_FILES", +} + + +def test_read_back_of_a_run_this_emitter_never_committed_does_not_pass( + hugging_face_module, +): + """For example a run handed over to another machine's emitter.""" + controls = hugging_face_module + assert read_back(controls, "unknown-run") == NOTHING_COMMITTED + + +def test_read_back_after_a_refused_hand_over_does_not_pass(hugging_face_module): + controls = hugging_face_module + refused = controls.module.query( + { + "protocol_version": 1, + "command": "stage_run_files", + "repository": STAGING_REPOSITORY, + "run_id": "run-1", + "prefix": "runs", + "files": [ + {"relative_path": "progress.json", "local_path": "/nonexistent/p"} + ], + "latest_pointer": None, + "run_index_entry": None, + "flush": True, + "commit_interval_seconds": 30.0, + } + )["receipt"] + assert refused["status"] == "refused" + cycle(controls) + assert read_back(controls, "run-1") == NOTHING_COMMITTED + + @pytest.mark.parametrize( ("field", "value"), [ From 9181bb771beb12fa8de81e74426952d04abdc4f6 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:13:10 +0400 Subject: [PATCH 08/23] Renew repository and job leases during deliveries; shorten them to 120 s Every repository and job lease lasted 900 seconds and was never renewed, so an emitter that was killed during a delivery kept the next emitter on the same machine from committing to that repository for up to 15 minutes. Leases now last 120 seconds. While a delivery runs, `HeldLeases` renews the repository lease and the claimed job's lease every 30 seconds on a background thread, and stops before the job's attempt is recorded and the repository is released, so both use the renewed lease. A long upload keeps its leases; a killed emitter's lease expires within 120 seconds. Co-Authored-By: Claude Opus 5.5 --- .../huggingface/service/delivery.py | 111 ++++++++++++-- .../huggingface/service/queue.py | 40 +++++ tests/huggingface/test_leases.py | 138 ++++++++++++++++++ 3 files changed, 278 insertions(+), 11 deletions(-) create mode 100644 tests/huggingface/test_leases.py diff --git a/src/microcosm_emitter/huggingface/service/delivery.py b/src/microcosm_emitter/huggingface/service/delivery.py index 57b5f27..af31d07 100644 --- a/src/microcosm_emitter/huggingface/service/delivery.py +++ b/src/microcosm_emitter/huggingface/service/delivery.py @@ -12,6 +12,7 @@ import threading import time from collections.abc import Callable +from dataclasses import replace from datetime import UTC, datetime from typing import Any @@ -20,9 +21,12 @@ from microcosm_emitter.huggingface.service.queue import ClaimedJob, HandoverQueue from microcosm_emitter.huggingface.service.run_files import commit_run_files -#: Long enough for one commit of a large upload; a crashed writer's lease -#: expires and another host on this machine resumes the repository. -REPOSITORY_LEASE_SECONDS = 900.0 +#: How long a repository or job lease lasts without renewal. While a delivery +#: runs, its leases are renewed every ``LEASE_RENEWAL_SECONDS``, so a long +#: upload keeps them; an emitter that was killed stops renewing, and another +#: emitter on this machine takes the repository over within one lease. +REPOSITORY_LEASE_SECONDS = 120.0 +LEASE_RENEWAL_SECONDS = 30.0 #: The job kinds this emitter delivers. DATASET_BUNDLE = "dataset_bundle" SUPPORTED_JOB_KINDS = frozenset({DATASET_BUNDLE}) @@ -44,6 +48,72 @@ def iso_timestamp(seconds: float) -> str: return datetime.fromtimestamp(seconds, UTC).replace(microsecond=0).isoformat() +class HeldLeases: + """A repository lease, and the job lease once one is claimed, kept renewed. + + A thread renews both every ``renewal_seconds`` until ``stop``. Read + ``repository_lease`` and ``job`` only after ``stop``: they then name the + current leases, which releasing the repository and finishing the job + must match. + """ + + def __init__( + self, + queue: HandoverQueue, + repository: str, + lease: float, + *, + clock: Callable[[], float], + lease_seconds: float, + renewal_seconds: float, + ) -> None: + self.queue, self.repository = queue, repository + self.repository_lease = lease + self.job: ClaimedJob | None = None + self.clock = clock + self.lease_seconds, self.renewal_seconds = lease_seconds, renewal_seconds + self._lock = threading.Lock() + self._stopped = threading.Event() + self._thread = threading.Thread(target=self._renew_until_stopped, daemon=True) + + def start(self) -> HeldLeases: + self._thread.start() + return self + + def hold_job(self, job: ClaimedJob) -> None: + with self._lock: + self.job = job + + def renew(self) -> None: + with self._lock: + renewed = self.clock() + self.lease_seconds + if renewed > self.repository_lease and self.queue.renew_repository( + self.repository, self.repository_lease, renewed + ): + self.repository_lease = renewed + job = self.job + if ( + job is not None + and renewed > job.lease_until + and self.queue.renew_job(job, renewed) + ): + self.job = replace(job, lease_until=renewed) + + def _renew_until_stopped(self) -> None: + while not self._stopped.wait(self.renewal_seconds): + try: + self.renew() + except Exception: + # A failed renewal leaves the lease to expire; the delivery's + # own result is recorded only if it still holds the lease. + pass + + def stop(self) -> None: + self._stopped.set() + if self._thread.is_alive(): + self._thread.join() + + class HuggingFaceDelivery: """Commit queued hand-overs with the build's credential.""" @@ -53,10 +123,12 @@ def __init__( *, api_factory: Callable[[], Any | None], clock: Callable[[], float] = time.time, + lease_renewal_seconds: float = LEASE_RENEWAL_SECONDS, ) -> None: self.queue = queue self.api_factory = api_factory self.clock = clock + self.lease_renewal_seconds = lease_renewal_seconds self._warned: set[str] = set() self._warning_lock = threading.Lock() @@ -99,20 +171,30 @@ def _deliver_repository( ) if lease is None: return 0 + held = HeldLeases( + self.queue, + repository, + lease, + clock=self.clock, + lease_seconds=REPOSITORY_LEASE_SECONDS, + renewal_seconds=self.lease_renewal_seconds, + ).start() try: if api is None: - self._credential_missing(repository, run_files_only=run_files_only) + self._credential_missing(held, run_files_only=run_files_only) return 0 hub = DatasetRepository(repository, api) if not run_files_only: job = self._claim_job(repository) if job is not None: - return int(self._deliver_job(hub, job)) + held.hold_job(job) + return int(self._deliver_job(hub, held)) if not self.queue.run_files_due(repository, self.clock()): return 0 return int(self._deliver_run_files(hub)) finally: - self.queue.release_repository(repository, lease) + held.stop() + self.queue.release_repository(repository, held.repository_lease) def _claim_job(self, repository: str) -> ClaimedJob | None: return self.queue.claim_job( @@ -149,15 +231,18 @@ def _finish_job( committed=committed, ) - def _deliver_job(self, hub: DatasetRepository, job: ClaimedJob) -> bool: - outcome = stage_bundle(hub, job.request, job.files) + def _deliver_job(self, hub: DatasetRepository, held: HeldLeases) -> bool: + assert held.job is not None, "The job is claimed before delivery." + outcome = stage_bundle(hub, held.job.request, held.job.files) if outcome.retry and outcome.error_code is not None: self.warn_once( outcome.error_code, UPLOAD_FAILED_MESSAGE.format(error_type=outcome.error_code), ) + held.stop() + assert held.job is not None self._finish_job( - job, + held.job, status=outcome.status, error_code=outcome.error_code, revision=outcome.revision, @@ -166,12 +251,16 @@ def _deliver_job(self, hub: DatasetRepository, job: ClaimedJob) -> bool: ) return outcome.committed - def _credential_missing(self, repository: str, *, run_files_only: bool) -> None: + def _credential_missing(self, held: HeldLeases, *, run_files_only: bool) -> None: + repository = held.repository self.warn_once(CREDENTIAL_MISSING, CREDENTIAL_MISSING_MESSAGE) job = None if run_files_only else self._claim_job(repository) if job is not None: + held.hold_job(job) + held.stop() + assert held.job is not None self._finish_job( - job, + held.job, status="pending", error_code=CREDENTIAL_MISSING, revision=None, diff --git a/src/microcosm_emitter/huggingface/service/queue.py b/src/microcosm_emitter/huggingface/service/queue.py index 0d7bf3a..dec5f73 100644 --- a/src/microcosm_emitter/huggingface/service/queue.py +++ b/src/microcosm_emitter/huggingface/service/queue.py @@ -140,6 +140,28 @@ def claim_repository( ) return lease if claimed is not None else None + def renew_repository(self, repository: str, lease: float, renewed: float) -> bool: + """Extend this writer's lease to ``renewed``; False if it was lost.""" + with self._sessions.begin() as session: + return ( + session.scalar( + update(RepositoryCommitSchedule) + .where( + RepositoryCommitSchedule.repository == repository, + RepositoryCommitSchedule.lease_until == lease, + ) + .values(lease_until=renewed) + .returning(RepositoryCommitSchedule.repository) + ) + is not None + ) + + def repository_lease(self, repository: str) -> float: + """When the repository's current lease expires; 0.0 when none is held.""" + with self._sessions() as session: + schedule = session.get(RepositoryCommitSchedule, repository) + return 0.0 if schedule is None else schedule.lease_until + def release_repository(self, repository: str, lease: float) -> None: with self._sessions.begin() as session: session.execute( @@ -611,6 +633,24 @@ def claim_job( files={row.name: row.content for row in copied}, ) + def renew_job(self, job: ClaimedJob, renewed: float) -> bool: + """Extend ``job``'s lease to ``renewed``; False if it was lost.""" + with self._sessions.begin() as session: + return ( + session.scalar( + update(HandoverJob) + .where( + HandoverJob.handover_id == job.handover_id, + HandoverJob.destination == job.destination, + HandoverJob.status == JOB_PENDING, + HandoverJob.lease_until == job.lease_until, + ) + .values(lease_until=renewed) + .returning(HandoverJob.handover_id) + ) + is not None + ) + def finish_job( self, job: ClaimedJob, diff --git a/tests/huggingface/test_leases.py b/tests/huggingface/test_leases.py new file mode 100644 index 0000000..05d43ab --- /dev/null +++ b/tests/huggingface/test_leases.py @@ -0,0 +1,138 @@ +"""Repository and job leases: renewed during long uploads, short after a crash.""" + +import hashlib +import time + +from fake_hugging_face import PRIVATE_REPOSITORY, STAGING_REPOSITORY + +from microcosm_emitter.huggingface.service import delivery as delivery_module +from microcosm_emitter_interface.bundles import BundleFile, StageDatasetBundle +from microcosm_emitter_interface.staging import RunFile, StageRunFiles + + +def hand_over_run_file(controls, run_id="run-1"): + local = controls.tmp_path / run_id / "progress.json" + local.parent.mkdir(parents=True, exist_ok=True) + local.write_bytes(b'{"stage":"calibrating"}') + return controls.module.query( + StageRunFiles( + repository=STAGING_REPOSITORY, + run_id=run_id, + files=(RunFile("progress.json", local),), + flush=True, + ).message() + )["receipt"] + + +def hand_over_bundle(controls, run_id="uk-run-1"): + directory = controls.tmp_path / run_id + directory.mkdir(parents=True, exist_ok=True) + files = [] + for name, role, content in ( + ("dataset.h5", "dataset", b"rows"), + ("staged_manifest.json", "staged_manifest", b'{"files": {}}'), + ): + (directory / name).write_bytes(content) + files.append( + BundleFile( + name=name, + role=role, + local_path=directory / name, + sha256=hashlib.sha256(content).hexdigest(), + bytes=len(content), + ) + ) + request = StageDatasetBundle( + repository=PRIVATE_REPOSITORY, run_id=run_id, files=tuple(files) + ) + controls.module.query(request.message()) + return request + + +def test_a_killed_emitter_holds_a_repository_for_at_most_one_short_lease( + hugging_face_module, +): + """The lease of an emitter that never released it expires within minutes.""" + controls = hugging_face_module + queue = controls.module.queue + lease_seconds = delivery_module.REPOSITORY_LEASE_SECONDS + assert lease_seconds <= 120 + # The killed emitter's lease, never released or renewed. + assert queue.claim_repository( + STAGING_REPOSITORY, now=controls.clock(), lease_seconds=lease_seconds + ) + hand_over_run_file(controls) + controls.module.delivery.run_cycle() + assert controls.hub.commits(STAGING_REPOSITORY) == [] + controls.clock.advance(121) + controls.module.delivery.run_cycle() + assert len(controls.hub.commits(STAGING_REPOSITORY)) == 1 + + +def wait_until(condition, timeout=5.0): + deadline = time.monotonic() + timeout + while not condition() and time.monotonic() < deadline: + time.sleep(0.01) + return condition() + + +def test_a_long_run_file_upload_keeps_renewing_its_repository_lease( + hugging_face_module, +): + controls = hugging_face_module + queue = controls.module.queue + controls.module.delivery.lease_renewal_seconds = 0.01 + claimed_during_upload = [] + + def slow_upload(): + # The upload outlasts several leases; renewal keeps the repository. + for _ in range(3): + controls.clock.advance(delivery_module.REPOSITORY_LEASE_SECONDS) + renewed_until = controls.clock() + delivery_module.REPOSITORY_LEASE_SECONDS + assert wait_until( + lambda until=renewed_until: ( + queue.repository_lease(STAGING_REPOSITORY) >= until + ) + ), "the lease was not renewed during the upload" + claimed_during_upload.append( + queue.claim_repository( + STAGING_REPOSITORY, now=controls.clock(), lease_seconds=1 + ) + ) + + controls.hub.run_before_next_commit(slow_upload) + hand_over_run_file(controls) + controls.module.delivery.run_cycle() + assert len(controls.hub.commits(STAGING_REPOSITORY)) == 1 + assert claimed_during_upload == [None, None, None] + assert queue.repository_lease(STAGING_REPOSITORY) == 0.0 + + +def test_a_long_bundle_upload_keeps_renewing_its_job_lease(hugging_face_module): + controls = hugging_face_module + queue = controls.module.queue + controls.module.delivery.lease_renewal_seconds = 0.01 + claimed_during_upload = [] + + def slow_upload(): + controls.clock.advance(delivery_module.REPOSITORY_LEASE_SECONDS * 2) + renewed_until = controls.clock() + delivery_module.REPOSITORY_LEASE_SECONDS + assert wait_until( + lambda: queue.repository_lease(PRIVATE_REPOSITORY) >= renewed_until + ), "the lease was not renewed during the upload" + claimed_during_upload.append( + queue.claim_job( + PRIVATE_REPOSITORY, + now=controls.clock(), + lease_seconds=1, + kinds=delivery_module.SUPPORTED_JOB_KINDS, + ) + ) + + controls.hub.run_before_next_commit(slow_upload) + request = hand_over_bundle(controls) + controls.module.delivery.run_cycle() + assert claimed_during_upload == [None] + receipt = queue.job_receipt(request.handover_id()) + # The attempt was recorded under the renewed lease. + assert (receipt["status"], receipt["attempts"]) == ("uploaded", 1) From 778f678b44a98e0718ea279d358196fae34b914b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:13:37 +0400 Subject: [PATCH 09/23] Describe the runs.json merge, commit spacing and row retention exactly Documentation only: - README: `runs.json` starts empty only when the Hub reports it missing; other download errors retry the commit (Microcosm started empty on any error). The 10-second spacing applies to every commit from this machine to a repository, of run files or bundles, forced or not, except the last commit at shutdown; it does not count other machines' commits. - `HandoverQueue`: rows are kept after delivery and nothing deletes them, except the copied hand-over bytes once the hand-over settles. The earlier docstring said records were kept only until delivered. - `StageRunFiles` and `MIN_COMMIT_SPACING_SECONDS` say the same about the interval and spacing. Co-Authored-By: Claude Opus 5.5 --- README.md | 4 ++-- .../src/microcosm_emitter_interface/staging.py | 8 +++++--- src/microcosm_emitter/huggingface/service/queue.py | 13 +++++++++++-- 3 files changed, 18 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 91dc56c..ff710a4 100644 --- a/README.md +++ b/README.md @@ -60,9 +60,9 @@ The `huggingface` module takes over the uploads Microcosm made itself. Microcosm **Staging run files** (Microcosm's `staging.py` and `staging_v2.py` uploads). A build hands over its current run files; the module copies their bytes into the spool database before answering, so the build may rewrite them at once. -- Files land at `//` (by default `runs//...`), exactly as the build's writer lays them out. Staging contract version 1 also writes `latest_staging.json` at the repository root, and its `runs.json` index: the module merges the run's summary into the repository's current copy at commit time, as Microcosm did, keeping other runs and the same JSON formatting. +- Files land at `//` (by default `runs//...`), exactly as the build's writer lays them out. Staging contract version 1 also writes `latest_staging.json` at the repository root, and its `runs.json` index: at commit time the module merges the run's summary into the repository's copy at the head it commits on, keeping other runs, with Microcosm's ordering and JSON formatting. It starts a new index only when the Hub answers that the repository has no `runs.json`; any other download error (a connection error, a timeout, a 429 or 5xx response) fails the commit and retries it after the backoff. Microcosm started from an empty index on any download error, which could drop every other run. - Each path keeps only the newest bytes, and bytes equal to those last committed are not written again. -- **One commit per repository per cycle.** Each module tick commits, for each repository with work due, every pending path in one commit. An unforced hand-over is due one interval after the repository's last commit (Microcosm's upload interval); a forced one (final states, artifacts) at the next cycle, but never sooner than 10 seconds after the previous commit. Microcosm's writers made one commit per file, up to eight per upload cycle, and each commit triggered the dataset repository's webhook. +- **One commit per repository per cycle.** Each module tick commits, for each repository with work due, every pending path in one commit. An unforced hand-over is due one interval (Microcosm's upload interval) after this machine's last commit to the repository; a forced one (final states, artifacts) at the next cycle. Either way, two commits from this machine to one repository, of run files or bundles, are at least 10 seconds apart; only the last commit at shutdown ignores that spacing, and commits from other machines are not counted. Microcosm's writers made one commit per file, up to eight per upload cycle, and each commit triggered the dataset repository's webhook. - Uploads stay best-effort. Requests answer within seconds or report that they were not accepted; a failed commit keeps the files queued and retries with backoff (up to five minutes); without a credential the files stay queued and one warning is printed. On shutdown the module commits pending files once more. Nothing raises into the build. - A commit that merges `runs.json` is pinned to the head it read (`parent_commit`); if another build commits first, the module reads the index again and merges again. - `read_back_run_files` commits pending files, then compares the Hub's copy of each committed file with the committed bytes. A run with no file this emitter committed (its hand-overs were refused, or another machine's emitter staged it) reads back as `failed` with `NO_COMMITTED_FILES`, never `passed`. diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/staging.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/staging.py index 121a502..75586d9 100644 --- a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/staging.py +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/staging.py @@ -60,9 +60,11 @@ class StageRunFiles: - ``run_index_entry`` is version 1's summary of this run; the emitter merges it into the repository's ``runs.json`` when it commits. - An unforced hand-over is committed one ``commit_interval_seconds`` after - the repository's last commit; ``flush`` asks for the next cycle (final - states and artifacts). Either way a cycle makes one commit per - repository, carrying the newest bytes of every pending path. + this machine's last commit to the repository; ``flush`` asks for the + next cycle (final states and artifacts). Either way a cycle makes one + commit per repository, carrying the newest bytes of every pending path, + and two commits from this machine to one repository are at least 10 + seconds apart. """ repository: str diff --git a/src/microcosm_emitter/huggingface/service/queue.py b/src/microcosm_emitter/huggingface/service/queue.py index dec5f73..c33a8ab 100644 --- a/src/microcosm_emitter/huggingface/service/queue.py +++ b/src/microcosm_emitter/huggingface/service/queue.py @@ -42,7 +42,9 @@ ) #: The shortest time between two commits to one repository from this machine, -#: even for hand-overs that ask to be committed at the next cycle. +#: of run files or jobs, even for hand-overs that ask to be committed at the +#: next cycle. Only the drain at shutdown ignores it; other machines' commits +#: are not counted. MIN_COMMIT_SPACING_SECONDS = 10.0 MAX_RETRY_SECONDS = 300.0 DESTINATION_HUGGING_FACE = "huggingface" @@ -92,7 +94,14 @@ class ClaimedJob: class HandoverQueue: - """Hand-over records, retained until delivered or settled.""" + """Hand-over records in the spool database. + + Rows are kept after delivery; nothing deletes them yet. Settled jobs keep + their receipts, committed run files keep their path and committed digest + (their bytes are cleared), and run deliveries and repository schedules + keep their counts and times. Only the bytes copied for a hand-over + (``handover_files``) are deleted, once every destination has settled. + """ def __init__(self, path: Path | str) -> None: self.path = Path(path) From ea3b5a30863f05bd96efb1eee29d8aa2457f0a14 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:09:01 +0400 Subject: [PATCH 10/23] Test that hand-overs leave the shared spool at a released revision Released emitters (0.2.0 to 0.4.0) refuse a spool database at a revision they don't know. Opening this branch's spool, or starting the Hugging Face module on it, moved the shared database to 20261020_03, so a build pinned to 0.2.1 on the same machine would start with no telemetry and no Orrery queue. The tests ask for the hand-over tables in a separate database next to the spool, with its own migration history, and for HuggingFaceSettings to name that file explicitly and refuse the spool file itself. Red: 6 failed, 25 passed. The compatibility test fails with "CommandError: Can't locate revision identified by '20261020_03'". Co-Authored-By: Claude Opus 5.5 --- tests/huggingface/test_queue.py | 53 +++++++++++++++++++ tests/interface/test_staging_client.py | 32 ++++++++++- tests/test_queue_compatibility.py | 73 +++++++++++++++++++++++++- 3 files changed, 156 insertions(+), 2 deletions(-) diff --git a/tests/huggingface/test_queue.py b/tests/huggingface/test_queue.py index bcb6801..7b6f2eb 100644 --- a/tests/huggingface/test_queue.py +++ b/tests/huggingface/test_queue.py @@ -69,6 +69,59 @@ def test_migration_adds_hand_over_tables_without_changing_telemetry(tmp_path): engine.dispose() +HAND_OVER_TABLES = { + "staged_run_files", + "staged_run_deliveries", + "repository_commit_schedules", + "handover_jobs", + "handover_files", +} + + +def test_a_new_hand_over_database_starts_empty_and_migrates_on_first_open(tmp_path): + from alembic import command + from sqlalchemy import inspect + + from microcosm_emitter.core.migrations import ( + HANDOVER_HISTORY, + alembic_config, + migration_head_revision, + ) + + path = tmp_path / "events.handovers.sqlite3" + assert not path.exists() + queue = HandoverQueue(path) + try: + assert current_database_revision(path) == migration_head_revision( + HANDOVER_HISTORY + ) + assert set(inspect(queue._engine).get_table_names()) == HAND_OVER_TABLES | { + "alembic_version" + } + assert queue.repositories_with_pending_run_files() == [] + with ( + queue._engine.begin() as connection, + alembic_config(connection=connection, history=HANDOVER_HISTORY) as config, + ): + command.check(config) + finally: + queue.close() + reopened = HandoverQueue(path) + reopened.close() + assert current_database_revision(path) == migration_head_revision(HANDOVER_HISTORY) + + +def test_a_spool_database_is_refused_as_a_hand_over_database(tmp_path): + from microcosm_emitter.telemetry.service.spool import EventSpool + + path = tmp_path / "events.sqlite3" + EventSpool(path).close() + revision = current_database_revision(path) + with pytest.raises(Exception, match="revision"): + HandoverQueue(path) + assert current_database_revision(path) == revision + + def test_latest_bytes_for_a_path_replace_earlier_pending_bytes(queue): stage(queue, "runs/run-1/progress.json", b'{"stage":"created"}', now=100) stage(queue, "runs/run-1/progress.json", b'{"stage":"calibrating"}', now=101) diff --git a/tests/interface/test_staging_client.py b/tests/interface/test_staging_client.py index 6a985ab..ef205eb 100644 --- a/tests/interface/test_staging_client.py +++ b/tests/interface/test_staging_client.py @@ -6,7 +6,11 @@ from fake_hugging_face import STAGING_REPOSITORY from microcosm_emitter_interface.client import EmitterClient -from microcosm_emitter_interface.staging import RunFile, StageRunFiles +from microcosm_emitter_interface.staging import ( + HuggingFaceSettings, + RunFile, + StageRunFiles, +) from microcosm_emitter_interface.transport import SocketTransport @@ -54,6 +58,32 @@ def test_requests_that_could_leave_the_run_folder_cannot_be_built( build_request(tmp_path / "progress.json") +def test_the_hand_over_database_is_named_next_to_the_spool_database(tmp_path): + spool_path = tmp_path / "events.sqlite3" + assert HuggingFaceSettings().configuration(spool_path) == { + "handover_database_path": str(tmp_path / "events.handovers.sqlite3") + } + + +def test_settings_can_name_their_own_hand_over_database(tmp_path): + chosen = tmp_path / "elsewhere" / "handovers.sqlite3" + settings = HuggingFaceSettings(handover_database_path=chosen) + assert settings.configuration(tmp_path / "events.sqlite3") == { + "handover_database_path": str(chosen) + } + + +def test_the_spool_database_cannot_be_the_hand_over_database(tmp_path): + spool_path = tmp_path / "events.sqlite3" + settings = HuggingFaceSettings(handover_database_path=spool_path) + with pytest.raises(ValueError, match="spool"): + settings.configuration(spool_path) + client = EmitterClient.start(hugging_face=settings, spool_path=spool_path) + assert not client.available + assert client.startup_error.reason == "invalid_configuration" + assert not spool_path.exists() + + def test_message_names_each_file_by_its_absolute_local_path(tmp_path, monkeypatch): monkeypatch.chdir(tmp_path) request = StageRunFiles( diff --git a/tests/test_queue_compatibility.py b/tests/test_queue_compatibility.py index 4bff66a..fdb6adc 100644 --- a/tests/test_queue_compatibility.py +++ b/tests/test_queue_compatibility.py @@ -1,14 +1,23 @@ -"""Open a synthetic #1099 queue without reassigning IDs or upload eligibility.""" +"""Open older spool databases, and leave them readable by released emitters. + +A synthetic #1099 queue opens without reassigning IDs or upload eligibility, +and hand-overs never move the shared spool database past a released revision. +""" import hashlib import json +import os import shutil from pathlib import Path +from alembic import command +from alembic.script import ScriptDirectory from sqlalchemy import select from sqlalchemy.orm import Session +from microcosm_emitter.core import migrations as migrations_module from microcosm_emitter.core.database import create_spool_engine +from microcosm_emitter.core.migrations import alembic_config, current_database_revision from microcosm_emitter.core.models import TelemetryEventRecord from microcosm_emitter.telemetry.service import spool as spool_module from microcosm_emitter.telemetry.service.spool import EventSpool @@ -51,3 +60,65 @@ def test_current_microcosm_queue_is_preserved(tmp_path, monkeypatch): assert not spool.has_deliverable() and spool.has_pending() finally: spool.close() + + +#: The spool migrations that released emitters (0.2.0 to 0.4.0) ship. +RELEASED_SPOOL_MIGRATIONS = ( + "v20261007_01_initial_spool.py", + "v20261008_02_graph_publication.py", +) +RELEASED_SPOOL_REVISION = "20261008_02" + + +def released_spool_history(directory: Path) -> ScriptDirectory: + """A migration history holding only the released spool migrations.""" + source = Path(migrations_module.__file__).parent / "alembic" + (directory / "versions").mkdir(parents=True) + shutil.copyfile(source / "script.py.mako", directory / "script.py.mako") + for name in RELEASED_SPOOL_MIGRATIONS: + shutil.copyfile(source / "versions" / name, directory / "versions" / name) + return ScriptDirectory(str(directory)) + + +def released_spool_database(path: Path) -> None: + """A spool database as a released emitter leaves it.""" + engine = create_spool_engine(path) + try: + with ( + engine.begin() as connection, + alembic_config(connection=connection) as config, + ): + command.upgrade(config, RELEASED_SPOOL_REVISION) + finally: + engine.dispose() + + +def test_hand_overs_leave_the_shared_spool_database_at_a_released_revision( + tmp_path, +): + """Released emitters refuse a spool database at a revision they don't know. + + Microcosm builds pinned to a released emitter share this database with + every newer emitter on the machine, so hand-overs must not migrate it. + """ + from microcosm_emitter.host.contracts import ModuleContext + from microcosm_emitter.huggingface.service.module import HuggingFaceModule + from microcosm_emitter_interface.staging import HuggingFaceSettings + + spool_path = tmp_path / "events.sqlite3" + released_spool_database(spool_path) + spool = EventSpool(spool_path) + module = HuggingFaceModule( + HuggingFaceSettings().configuration(spool_path), + ModuleContext(os.getpid()), + api_factory=lambda: None, + ) + module.initialize() + module.close(0) + spool.close() + + revision = current_database_revision(spool_path) + # A released emitter's check before it writes anything. + released_spool_history(tmp_path / "released-alembic").get_revision(revision) + assert revision == RELEASED_SPOOL_REVISION + assert (tmp_path / "events.handovers.sqlite3").exists() From c0b71c775cdc0c1b44ec4eef10a0f6ae44ec7d1d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:12:00 +0400 Subject: [PATCH 11/23] Keep hand-overs in their own database so the shared spool stays released The hand-over tables (staged run files, run deliveries, repository schedules, jobs and copied files) move from the shared spool database to a separate SQLite file with its own Alembic history: - core/handover_models.py holds them on their own declarative base, and core/handover_alembic holds their one migration, renamed handovers_20261020_01 so neither history accepts the other's file; - core/migrations.py applies either history through one function; - the spool database's history ends at 20261008_02 again, the revision every released emitter (0.2.0 to 0.4.0) knows. HuggingFaceSettings now sends handover_database_path: by default events.handovers.sqlite3 next to the spool, or a file the build names, never the spool file itself (invalid_configuration). The module opens only that file. The unreleased configuration fixture and its digest follow, the installed-interface smoke check confirms the wheel creates the file, and the README and changelog fragment describe it. Green: 31 passed (the three red files); 188 passed across #23's changed test files. Co-Authored-By: Claude Opus 5.5 --- README.md | 8 +- .../huggingface-handover-queue.added.md | 2 +- .../microcosm_emitter_interface/__init__.py | 6 +- .../huggingface-configuration.json | 2 +- .../src/microcosm_emitter_interface/spool.py | 13 ++- .../microcosm_emitter_interface/staging.py | 24 +++- .../core/handover_alembic/env.py | 28 +++++ .../core/handover_alembic/script.py.mako | 31 +++++ .../versions/v20261020_01_handover_queue.py} | 10 +- src/microcosm_emitter/core/handover_models.py | 106 ++++++++++++++++++ src/microcosm_emitter/core/migrations.py | 71 +++++++++--- src/microcosm_emitter/core/models.py | 88 --------------- .../huggingface/service/module.py | 4 +- .../huggingface/service/queue.py | 18 +-- tests/huggingface/conftest.py | 2 +- tests/huggingface/test_bundles.py | 2 +- tests/huggingface/test_queue.py | 19 ++-- tests/interface/conftest.py | 2 +- tests/interface/test_protocol_fixtures.py | 2 +- tests/telemetry/test_persistence.py | 7 +- tools/installed-interface-smoke.py | 3 + tools/installed-smoke.py | 2 +- 22 files changed, 307 insertions(+), 143 deletions(-) create mode 100644 src/microcosm_emitter/core/handover_alembic/env.py create mode 100644 src/microcosm_emitter/core/handover_alembic/script.py.mako rename src/microcosm_emitter/core/{alembic/versions/v20261020_03_handover_queue.py => handover_alembic/versions/v20261020_01_handover_queue.py} (92%) create mode 100644 src/microcosm_emitter/core/handover_models.py diff --git a/README.md b/README.md index ff710a4..2f8825e 100644 --- a/README.md +++ b/README.md @@ -42,7 +42,7 @@ The adapter bounds socket calls and reports telemetry failures without failing t The service uses the existing Hugging Face credential from the environment or local login cache. The collector validates PolicyEngine membership before issuing a short-lived write credential. Missing or rejected credentials permanently mark that producer's events as local-only. Permanent HTTP client errors when registering a run or uploading events also keep only the affected producer's records local; HTTP 408, 429, server errors and network failures remain retryable. HTTP 401 invalidates the short-lived credential for renewal. The production destination belongs to the package; development overrides accept loopback addresses only. -The queue upgrades the existing `20261007_01` Alembic revision to `20261008_02`, adding graph-job records, and then to `20261020_03`, adding Hugging Face hand-over records, without changing producer sequences, event IDs or upload eligibility. Unversioned databases are rejected without schema modification. There is no raw SQL, automatic schema adoption, or parallel legacy implementation. The hosted telemetry collector and calibration dashboard do not change. +The queue upgrades the existing `20261007_01` Alembic revision to `20261008_02`, adding graph-job records without changing producer sequences, event IDs or upload eligibility. Unversioned databases are rejected without schema modification. There is no raw SQL, automatic schema adoption, or parallel legacy implementation. The hosted telemetry collector and calibration dashboard do not change. ## Graph publication @@ -58,7 +58,9 @@ Preserve the selected source files until publication succeeds. `microcosm-emitte The `huggingface` module takes over the uploads Microcosm made itself. Microcosm still decides what to hand over and validates its own contracts; the module only writes the bytes it was given, with the build's own Hugging Face credential. -**Staging run files** (Microcosm's `staging.py` and `staging_v2.py` uploads). A build hands over its current run files; the module copies their bytes into the spool database before answering, so the build may rewrite them at once. +**Where hand-overs are stored.** The module keeps its queue in its own SQLite file, the hand-over database, with its own Alembic history (`core/handover_alembic`). `HuggingFaceSettings` names it: by default it sits next to the spool database (`events.sqlite3` → `events.handovers.sqlite3`), and `handover_database_path` chooses another file, never the spool database itself. The spool database stays at revision `20261008_02`, because every emitter on a machine shares it, including the released emitters (0.2.0 to 0.4.0) that Microcosm builds pin, and those refuse a revision they don't know. Each history also refuses the other's file. + +**Staging run files** (Microcosm's `staging.py` and `staging_v2.py` uploads). A build hands over its current run files; the module copies their bytes into the hand-over database before answering, so the build may rewrite them at once. - Files land at `//` (by default `runs//...`), exactly as the build's writer lays them out. Staging contract version 1 also writes `latest_staging.json` at the repository root, and its `runs.json` index: at commit time the module merges the run's summary into the repository's copy at the head it commits on, keeping other runs, with Microcosm's ordering and JSON formatting. It starts a new index only when the Hub answers that the repository has no `runs.json`; any other download error (a connection error, a timeout, a 429 or 5xx response) fails the commit and retries it after the backoff. Microcosm started from an empty index on any download error, which could drop every other run. - Each path keeps only the newest bytes, and bytes equal to those last committed are not written again. @@ -69,7 +71,7 @@ The `huggingface` module takes over the uploads Microcosm made itself. Microcosm **Staged dataset bundles** (Microcosm's `staging_dataset.stage_bundle`, used by `uk_runtime/rowwise_staging.py` and the Orrery evidence bundle). A build verifies its bundle and writes the sidecars as before, then hands the files over by path, each with its role (`dataset`, `calibration_diagnostics`, `build_manifest`, `staged_manifest`, `checksums`, ...), SHA-256 and size. The hand-over is a durable job: -- Before answering, the module copies the bundle's small files into the spool database, smallest first, each of at most 32 MiB and together at most 64 MiB, after checking each one's size and SHA-256 against the hand-over (`FILE_SIZE_DIFFERS`, `FILE_DIGEST_DIFFERS`). The receipt's `copied_files` names them, and the build may rewrite them at once: Microcosm rewrites its build manifest, its `sha256sums.txt` and, in national builds, its build record right after staging, and the commit carries the bytes as they were handed over. The copies are deleted once the hand-over settles. +- Before answering, the module copies the bundle's small files into the hand-over database, smallest first, each of at most 32 MiB and together at most 64 MiB, after checking each one's size and SHA-256 against the hand-over (`FILE_SIZE_DIFFERS`, `FILE_DIGEST_DIFFERS`). The receipt's `copied_files` names them, and the build may rewrite them at once: Microcosm rewrites its build manifest, its `sha256sums.txt` and, in national builds, its build record right after staging, and the commit carries the bytes as they were handed over. The copies are deleted once the hand-over settles. - Larger files (the datasets) are uploaded from their paths. The worker re-hashes each one first and never uploads changed bytes (`PRESERVED_BYTES_CHANGED`); a missing file waits for an operator to restore it (`PRESERVED_FILES_MISSING`). - As in Microcosm, it writes every file under `//` (by default `staged//`) in one commit pinned to the head it checked, with the message `Stage under /`. A repository copy of the bundle's `staged_manifest` whose `files` equal the outputs' digests is `already_staged` with that copy's commit; a different one is `REMOTE_DIFFERS` and is never overwritten. - Unlike Microcosm's single attempt, `REPOSITORY_UNAVAILABLE` and `UPLOAD_FAILED` are retried with backoff, and a commit that loses a race to another writer is checked and tried again. The receipt keeps the last error code while the job is pending. diff --git a/changelog.d/huggingface-handover-queue.added.md b/changelog.d/huggingface-handover-queue.added.md index 471245f..c6fe36e 100644 --- a/changelog.d/huggingface-handover-queue.added.md +++ b/changelog.d/huggingface-handover-queue.added.md @@ -1 +1 @@ -Add a durable Hugging Face hand-over queue to the emitter's spool database (Alembic revision `20261020_03`). Staging run files keep only the newest handed-over bytes per repository path; dataset bundles and releases are jobs with leases, retries and receipts; and a per-repository schedule spaces commits and lets one writer on a machine hold a repository at a time. Hugging Face calls use the build's own credential, found in the same order as telemetry's. +Add a durable Hugging Face hand-over queue in its own SQLite database next to the spool database (`events.handovers.sqlite3`, with its own Alembic history), so the shared spool database stays at a revision that released emitters accept; `HuggingFaceSettings.handover_database_path` can name another file. Staging run files keep only the newest handed-over bytes per repository path; dataset bundles and releases are jobs with leases, retries and receipts; and a per-repository schedule spaces commits and lets one writer on a machine hold a repository at a time. Hugging Face calls use the build's own credential, found in the same order as telemetry's. diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/__init__.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/__init__.py index f9e7526..4583df3 100644 --- a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/__init__.py +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/__init__.py @@ -12,7 +12,10 @@ ) from microcosm_emitter_interface.orrery import OrrerySettings from microcosm_emitter_interface.protocol import PROTOCOL_VERSION -from microcosm_emitter_interface.spool import default_spool_path +from microcosm_emitter_interface.spool import ( + default_handover_database_path, + default_spool_path, +) from microcosm_emitter_interface.staging import ( HuggingFaceSettings, RunFile, @@ -35,5 +38,6 @@ "StageRunFiles", "TelemetryRun", "TelemetrySettings", + "default_handover_database_path", "default_spool_path", ] diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-configuration.json b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-configuration.json index 0288ad2..4a9948e 100644 --- a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-configuration.json +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-configuration.json @@ -3,6 +3,6 @@ "module": "huggingface", "kind": "configuration", "body": { - "spool_path": "{spool_path}" + "handover_database_path": "{directory}/events.handovers.sqlite3" } } diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/spool.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/spool.py index 9d03d43..0e3bbea 100644 --- a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/spool.py +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/spool.py @@ -1,4 +1,4 @@ -"""Where the emitter keeps its durable queue.""" +"""Where the emitter keeps its durable queues.""" import os from pathlib import Path @@ -9,3 +9,14 @@ def default_spool_path() -> Path: configured = os.environ.get("XDG_CACHE_HOME", "").strip() root = Path(configured).expanduser() if configured else Path.home() / ".cache" return root / "microcosm" / "telemetry" / "events.sqlite3" + + +def default_handover_database_path(spool_path: Path) -> Path: + """The Hugging Face module's hand-over database, next to the spool database. + + ``events.sqlite3`` gives ``events.handovers.sqlite3``. It is a separate + file because released emitters, which Microcosm builds pin and which + share the spool database, refuse a spool database whose migration + revision they don't know. + """ + return spool_path.with_name(f"{spool_path.stem}.handovers.sqlite3") diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/staging.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/staging.py index 75586d9..0ee0e34 100644 --- a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/staging.py +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/staging.py @@ -25,6 +25,7 @@ ) from microcosm_emitter_interface.protocol import PROTOCOL_VERSION from microcosm_emitter_interface.serialization import wire_value +from microcosm_emitter_interface.spool import default_handover_database_path DEFAULT_RUN_PREFIX = "runs" DEFAULT_COMMIT_INTERVAL_SECONDS = 30.0 @@ -32,10 +33,29 @@ @dataclass(frozen=True) class HuggingFaceSettings: - """Configuration for the emitter's Hugging Face module.""" + """Configuration for the emitter's Hugging Face module. + + The module keeps its queue in its own SQLite file, the hand-over database. + ``handover_database_path`` names it; by default it sits next to the spool + database (``events.sqlite3`` → ``events.handovers.sqlite3``). It can never + be the spool database itself, which released emitters share and which + must stay at a migration revision they accept. + """ + + handover_database_path: Path | None = None def configuration(self, spool_path: Path) -> dict[str, Any]: - return {"spool_path": str(spool_path)} + handover_database = ( + Path(self.handover_database_path) + if self.handover_database_path is not None + else default_handover_database_path(spool_path) + ) + if handover_database.resolve() == Path(spool_path).resolve(): + raise ValueError( + "The hand-over database must be a different file from the " + "spool database." + ) + return {"handover_database_path": str(handover_database)} @dataclass(frozen=True) diff --git a/src/microcosm_emitter/core/handover_alembic/env.py b/src/microcosm_emitter/core/handover_alembic/env.py new file mode 100644 index 0000000..351a956 --- /dev/null +++ b/src/microcosm_emitter/core/handover_alembic/env.py @@ -0,0 +1,28 @@ +"""Alembic environment for the Hugging Face module's hand-over database.""" + +from alembic import context +from sqlalchemy.engine import Connection + +from microcosm_emitter.core.handover_models import HandoverModel + +_MISSING_CONNECTION_ERROR = ( + "hand-over database migrations require a programmatically supplied connection" +) + + +def run_migrations() -> None: + """Run migrations on the connection supplied by the emitter service.""" + + connection = context.config.attributes.get("connection") + if not isinstance(connection, Connection): + raise RuntimeError(_MISSING_CONNECTION_ERROR) + context.configure( + connection=connection, + target_metadata=HandoverModel.metadata, + render_as_batch=True, + ) + with context.begin_transaction(): + context.run_migrations() + + +run_migrations() diff --git a/src/microcosm_emitter/core/handover_alembic/script.py.mako b/src/microcosm_emitter/core/handover_alembic/script.py.mako new file mode 100644 index 0000000..f059f76 --- /dev/null +++ b/src/microcosm_emitter/core/handover_alembic/script.py.mako @@ -0,0 +1,31 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} +""" + +from __future__ import annotations + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op +${imports if imports else ""} + +revision: str = ${repr(up_revision)} +down_revision: str | None = ${repr(down_revision)} +branch_labels: str | Sequence[str] | None = ${repr(branch_labels)} +depends_on: str | Sequence[str] | None = ${repr(depends_on)} + + +def upgrade() -> None: + """Apply this schema revision.""" + + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + """Reverse this schema revision.""" + + ${downgrades if downgrades else "pass"} diff --git a/src/microcosm_emitter/core/alembic/versions/v20261020_03_handover_queue.py b/src/microcosm_emitter/core/handover_alembic/versions/v20261020_01_handover_queue.py similarity index 92% rename from src/microcosm_emitter/core/alembic/versions/v20261020_03_handover_queue.py rename to src/microcosm_emitter/core/handover_alembic/versions/v20261020_01_handover_queue.py index b6bd30f..d69da1b 100644 --- a/src/microcosm_emitter/core/alembic/versions/v20261020_03_handover_queue.py +++ b/src/microcosm_emitter/core/handover_alembic/versions/v20261020_01_handover_queue.py @@ -1,10 +1,14 @@ -"""Add durable Hugging Face hand-overs without changing telemetry or graph jobs.""" +"""Create the hand-over database's tables. + +This history is separate from the spool database's, and its revision IDs +cannot be mistaken for the spool's: each history refuses the other's file. +""" import sqlalchemy as sa from alembic import op -revision = "20261020_03" -down_revision = "20261008_02" +revision = "handovers_20261020_01" +down_revision = None branch_labels = None depends_on = None diff --git a/src/microcosm_emitter/core/handover_models.py b/src/microcosm_emitter/core/handover_models.py new file mode 100644 index 0000000..d0d78a4 --- /dev/null +++ b/src/microcosm_emitter/core/handover_models.py @@ -0,0 +1,106 @@ +"""SQLAlchemy ORM models for the Hugging Face module's hand-over database. + +These tables live in their own SQLite file with their own migration history +(``core/handover_alembic``), never in the shared spool database: released +emitters refuse a spool database at a revision they don't know, and Microcosm +builds pinned to a released emitter share that database with newer ones. +""" + +from __future__ import annotations + +from typing import Any + +from sqlalchemy import Boolean, Float, Integer, LargeBinary, Text +from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column + +from microcosm_emitter.core.models import JsonObjectText + + +class HandoverModel(DeclarativeBase): + """Declarative base for the hand-over database.""" + + +class StagedRunFile(HandoverModel): + """The latest handed-over bytes for one repository path of a staging run. + + A later hand-over of the same path replaces pending bytes, so each commit + carries only the newest state. Committed rows keep the committed digest. + """ + + __tablename__ = "staged_run_files" + repository: Mapped[str] = mapped_column(Text, primary_key=True) + path_in_repo: Mapped[str] = mapped_column(Text, primary_key=True) + run_id: Mapped[str] = mapped_column(Text, nullable=False) + content: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + content_sha256: Mapped[str] = mapped_column(Text, nullable=False) + committed_sha256: Mapped[str | None] = mapped_column(Text) + pending: Mapped[bool] = mapped_column(Boolean, nullable=False) + due_at: Mapped[float] = mapped_column(Float, nullable=False) + updated_at: Mapped[float] = mapped_column(Float, nullable=False) + + +class StagedRunDelivery(HandoverModel): + """Delivery counts for one staging run, and its pending runs-index entry.""" + + __tablename__ = "staged_run_deliveries" + repository: Mapped[str] = mapped_column(Text, primary_key=True) + run_id: Mapped[str] = mapped_column(Text, primary_key=True) + files_attempted: Mapped[int] = mapped_column(Integer, nullable=False) + files_committed: Mapped[int] = mapped_column(Integer, nullable=False) + commit_attempts: Mapped[int] = mapped_column(Integer, nullable=False) + commits: Mapped[int] = mapped_column(Integer, nullable=False) + last_error_code: Mapped[str | None] = mapped_column(Text) + last_commit: Mapped[str | None] = mapped_column(Text) + # JSON text in the build's key order: the merged runs.json keeps it. + run_index_entry: Mapped[str | None] = mapped_column(Text) + run_index_pending: Mapped[bool] = mapped_column(Boolean, nullable=False) + run_index_due_at: Mapped[float] = mapped_column(Float, nullable=False) + updated_at: Mapped[float] = mapped_column(Float, nullable=False) + + +class RepositoryCommitSchedule(HandoverModel): + """When a repository may next receive a commit from this machine.""" + + __tablename__ = "repository_commit_schedules" + repository: Mapped[str] = mapped_column(Text, primary_key=True) + last_commit_at: Mapped[float] = mapped_column(Float, nullable=False) + retry_after: Mapped[float] = mapped_column(Float, nullable=False) + consecutive_failures: Mapped[int] = mapped_column(Integer, nullable=False) + lease_until: Mapped[float] = mapped_column(Float, nullable=False) + + +class HandoverJob(HandoverModel): + """One durable hand-over for one destination, retried until it settles. + + The same hand-over can have a row per destination (today only + ``huggingface``), so a later destination needs no new table. + """ + + __tablename__ = "handover_jobs" + handover_id: Mapped[str] = mapped_column(Text, primary_key=True) + destination: Mapped[str] = mapped_column(Text, primary_key=True) + kind: Mapped[str] = mapped_column(Text, nullable=False) + target: Mapped[str] = mapped_column(Text, nullable=False) + request: Mapped[dict[str, Any]] = mapped_column(JsonObjectText, nullable=False) + status: Mapped[str] = mapped_column(Text, nullable=False) + attempts: Mapped[int] = mapped_column(Integer, nullable=False) + next_attempt_at: Mapped[float] = mapped_column(Float, nullable=False) + lease_until: Mapped[float] = mapped_column(Float, nullable=False) + receipt: Mapped[dict[str, Any]] = mapped_column(JsonObjectText, nullable=False) + created_at: Mapped[float] = mapped_column(Float, nullable=False) + + +class HandoverFile(HandoverModel): + """The bytes of one small hand-over file, copied when it was handed over. + + The build may rewrite such a file at once (Microcosm rewrites its build + manifest, checksum listing and build record right after staging), so the + delivery uploads these bytes instead of reading the file again. The rows + are deleted once every destination of the hand-over has settled. + """ + + __tablename__ = "handover_files" + handover_id: Mapped[str] = mapped_column(Text, primary_key=True) + name: Mapped[str] = mapped_column(Text, primary_key=True) + content: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) + content_sha256: Mapped[str] = mapped_column(Text, nullable=False) diff --git a/src/microcosm_emitter/core/migrations.py b/src/microcosm_emitter/core/migrations.py index 9815d9d..287f9d4 100644 --- a/src/microcosm_emitter/core/migrations.py +++ b/src/microcosm_emitter/core/migrations.py @@ -1,10 +1,22 @@ -"""Programmatic Alembic migration entry points for the telemetry spool.""" +"""Programmatic Alembic migration entry points for the emitter's databases. + +The emitter keeps two SQLite databases, each with its own migration history: + +- the **spool database** (telemetry events and graph publication jobs). Every + emitter on a machine shares it, including the released emitters that + Microcosm builds pin, and those refuse a revision they don't know. Its + history must stay at revisions that released emitters ship until they are + no longer in use; +- the **hand-over database** of the Hugging Face module, next to the spool + database. Only emitters with that module open it. +""" from __future__ import annotations import fcntl from collections.abc import Iterator from contextlib import contextmanager +from dataclasses import dataclass from importlib.resources import as_file, files from pathlib import Path @@ -20,14 +32,28 @@ _MIGRATION_TARGET = "head" +@dataclass(frozen=True) +class MigrationHistory: + """One database's migration scripts, in a directory of this package.""" + + directory: str + #: Names the database in the refusal of an unversioned file. + label: str + + +SPOOL_HISTORY = MigrationHistory(directory="alembic", label="telemetry") +HANDOVER_HISTORY = MigrationHistory(directory="handover_alembic", label="hand-over") + + @contextmanager def alembic_config( *, connection: Connection | None = None, + history: MigrationHistory = SPOOL_HISTORY, ) -> Iterator[Config]: - """Yield Alembic configuration with migrations on the filesystem.""" + """Yield Alembic configuration with ``history``'s migrations on the filesystem.""" - migration_resources = files(__package__).joinpath("alembic") + migration_resources = files(__package__).joinpath(history.directory) with as_file(migration_resources) as migration_directory: config = Config() config.set_main_option("script_location", str(migration_directory)) @@ -36,10 +62,10 @@ def alembic_config( yield config -def upgrade_spool_database(engine: Engine) -> None: - """Apply every committed spool migration to an engine.""" +def upgrade_database(engine: Engine, history: MigrationHistory) -> None: + """Apply every committed migration of ``history`` to an engine.""" - # All build processes share this database. Serialize migration checks and + # All build processes share these databases. Serialize migration checks and # upgrades; SQLite transaction locks alone do not serialize schema inspection. path = Path(engine.url.database) with path.with_suffix(".migration.lock").open("a") as lock: @@ -49,19 +75,32 @@ def upgrade_spool_database(engine: Engine) -> None: revision = MigrationContext.configure(connection).get_current_revision() if names and revision is None: raise RuntimeError( - "Refusing an unversioned telemetry database; no changes made." + f"Refusing an unversioned {history.label} database; " + "no changes made." ) - with alembic_config(connection=connection) as config: - history = ScriptDirectory.from_config(config) + with alembic_config(connection=connection, history=history) as config: + scripts = ScriptDirectory.from_config(config) if revision is not None: - history.get_revision( - revision - ) # Reject unknown revisions before writes. + # Reject unknown revisions, including the other history's, + # before writes. + scripts.get_revision(revision) command.upgrade(config, _MIGRATION_TARGET) +def upgrade_spool_database(engine: Engine) -> None: + """Apply every committed spool migration to an engine.""" + + upgrade_database(engine, SPOOL_HISTORY) + + +def upgrade_handover_database(engine: Engine) -> None: + """Apply every committed hand-over migration to an engine.""" + + upgrade_database(engine, HANDOVER_HISTORY) + + def current_database_revision(path: Path | str) -> str | None: - """Return the Alembic revision recorded by one spool database.""" + """Return the Alembic revision recorded by one emitter database.""" engine = create_spool_engine(path) try: @@ -71,8 +110,8 @@ def current_database_revision(path: Path | str) -> str | None: engine.dispose() -def migration_head_revision() -> str | None: - """Return the single current head from the packaged migration history.""" +def migration_head_revision(history: MigrationHistory = SPOOL_HISTORY) -> str | None: + """Return the single current head from one packaged migration history.""" - with alembic_config() as config: + with alembic_config(history=history) as config: return ScriptDirectory.from_config(config).get_current_head() diff --git a/src/microcosm_emitter/core/models.py b/src/microcosm_emitter/core/models.py index f4380c5..0bca2a9 100644 --- a/src/microcosm_emitter/core/models.py +++ b/src/microcosm_emitter/core/models.py @@ -6,12 +6,10 @@ from typing import Any from sqlalchemy import ( - Boolean, Float, ForeignKeyConstraint, Index, Integer, - LargeBinary, Text, UniqueConstraint, ) @@ -145,89 +143,3 @@ class GraphPublicationJob(SpoolModel): next_attempt_at: Mapped[float] = mapped_column(Float, nullable=False) lease_until: Mapped[float] = mapped_column(Float, nullable=False) receipt: Mapped[dict[str, Any]] = mapped_column(JsonObjectText, nullable=False) - - -class StagedRunFile(SpoolModel): - """The latest handed-over bytes for one repository path of a staging run. - - A later hand-over of the same path replaces pending bytes, so each commit - carries only the newest state. Committed rows keep the committed digest. - """ - - __tablename__ = "staged_run_files" - repository: Mapped[str] = mapped_column(Text, primary_key=True) - path_in_repo: Mapped[str] = mapped_column(Text, primary_key=True) - run_id: Mapped[str] = mapped_column(Text, nullable=False) - content: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) - content_sha256: Mapped[str] = mapped_column(Text, nullable=False) - committed_sha256: Mapped[str | None] = mapped_column(Text) - pending: Mapped[bool] = mapped_column(Boolean, nullable=False) - due_at: Mapped[float] = mapped_column(Float, nullable=False) - updated_at: Mapped[float] = mapped_column(Float, nullable=False) - - -class StagedRunDelivery(SpoolModel): - """Delivery counts for one staging run, and its pending runs-index entry.""" - - __tablename__ = "staged_run_deliveries" - repository: Mapped[str] = mapped_column(Text, primary_key=True) - run_id: Mapped[str] = mapped_column(Text, primary_key=True) - files_attempted: Mapped[int] = mapped_column(Integer, nullable=False) - files_committed: Mapped[int] = mapped_column(Integer, nullable=False) - commit_attempts: Mapped[int] = mapped_column(Integer, nullable=False) - commits: Mapped[int] = mapped_column(Integer, nullable=False) - last_error_code: Mapped[str | None] = mapped_column(Text) - last_commit: Mapped[str | None] = mapped_column(Text) - # JSON text in the build's key order: the merged runs.json keeps it. - run_index_entry: Mapped[str | None] = mapped_column(Text) - run_index_pending: Mapped[bool] = mapped_column(Boolean, nullable=False) - run_index_due_at: Mapped[float] = mapped_column(Float, nullable=False) - updated_at: Mapped[float] = mapped_column(Float, nullable=False) - - -class RepositoryCommitSchedule(SpoolModel): - """When a repository may next receive a commit from this machine.""" - - __tablename__ = "repository_commit_schedules" - repository: Mapped[str] = mapped_column(Text, primary_key=True) - last_commit_at: Mapped[float] = mapped_column(Float, nullable=False) - retry_after: Mapped[float] = mapped_column(Float, nullable=False) - consecutive_failures: Mapped[int] = mapped_column(Integer, nullable=False) - lease_until: Mapped[float] = mapped_column(Float, nullable=False) - - -class HandoverJob(SpoolModel): - """One durable hand-over for one destination, retried until it settles. - - The same hand-over can have a row per destination (today only - ``huggingface``), so a later destination needs no new table. - """ - - __tablename__ = "handover_jobs" - handover_id: Mapped[str] = mapped_column(Text, primary_key=True) - destination: Mapped[str] = mapped_column(Text, primary_key=True) - kind: Mapped[str] = mapped_column(Text, nullable=False) - target: Mapped[str] = mapped_column(Text, nullable=False) - request: Mapped[dict[str, Any]] = mapped_column(JsonObjectText, nullable=False) - status: Mapped[str] = mapped_column(Text, nullable=False) - attempts: Mapped[int] = mapped_column(Integer, nullable=False) - next_attempt_at: Mapped[float] = mapped_column(Float, nullable=False) - lease_until: Mapped[float] = mapped_column(Float, nullable=False) - receipt: Mapped[dict[str, Any]] = mapped_column(JsonObjectText, nullable=False) - created_at: Mapped[float] = mapped_column(Float, nullable=False) - - -class HandoverFile(SpoolModel): - """The bytes of one small hand-over file, copied when it was handed over. - - The build may rewrite such a file at once (Microcosm rewrites its build - manifest, checksum listing and build record right after staging), so the - delivery uploads these bytes instead of reading the file again. The rows - are deleted once every destination of the hand-over has settled. - """ - - __tablename__ = "handover_files" - handover_id: Mapped[str] = mapped_column(Text, primary_key=True) - name: Mapped[str] = mapped_column(Text, primary_key=True) - content: Mapped[bytes] = mapped_column(LargeBinary, nullable=False) - content_sha256: Mapped[str] = mapped_column(Text, nullable=False) diff --git a/src/microcosm_emitter/huggingface/service/module.py b/src/microcosm_emitter/huggingface/service/module.py index 73b3197..af38ef3 100644 --- a/src/microcosm_emitter/huggingface/service/module.py +++ b/src/microcosm_emitter/huggingface/service/module.py @@ -75,7 +75,9 @@ def __init__( self._closed = False def initialize(self) -> None: - self.queue = HandoverQueue(Path(self.configuration["spool_path"])) + # The hand-over database, never the shared spool database (see + # ``HandoverQueue``); ``HuggingFaceSettings`` names it. + self.queue = HandoverQueue(Path(self.configuration["handover_database_path"])) self.delivery = HuggingFaceDelivery( self.queue, api_factory=self.api_factory, clock=self.clock ) diff --git a/src/microcosm_emitter/huggingface/service/queue.py b/src/microcosm_emitter/huggingface/service/queue.py index c33a8ab..86e518e 100644 --- a/src/microcosm_emitter/huggingface/service/queue.py +++ b/src/microcosm_emitter/huggingface/service/queue.py @@ -1,7 +1,11 @@ -"""Durable hand-overs, kept in the emitter's spool database. +"""Durable hand-overs, kept in the Hugging Face module's hand-over database. -Two record shapes share one SQLite database (and its Alembic history) with -telemetry and graph jobs: +The hand-over database is its own SQLite file, next to the spool database +(``events.sqlite3`` → ``events.handovers.sqlite3``), with its own Alembic +history. The shared spool database stays at a revision that released +emitters accept: Microcosm builds pinned to one share it on the same machine. + +Two record shapes share the hand-over database: - **Staging run files** coalesce: each repository path keeps only the newest handed-over bytes, so a commit carries the latest state of every pending @@ -32,14 +36,14 @@ create_spool_engine, create_spool_session_factory, ) -from microcosm_emitter.core.migrations import upgrade_spool_database -from microcosm_emitter.core.models import ( +from microcosm_emitter.core.handover_models import ( HandoverFile, HandoverJob, RepositoryCommitSchedule, StagedRunDelivery, StagedRunFile, ) +from microcosm_emitter.core.migrations import upgrade_handover_database #: The shortest time between two commits to one repository from this machine, #: of run files or jobs, even for hand-overs that ask to be committed at the @@ -94,7 +98,7 @@ class ClaimedJob: class HandoverQueue: - """Hand-over records in the spool database. + """Hand-over records in the hand-over database at ``path``. Rows are kept after delivery; nothing deletes them yet. Settled jobs keep their receipts, committed run files keep their path and committed digest @@ -107,7 +111,7 @@ def __init__(self, path: Path | str) -> 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) + upgrade_handover_database(self._engine) self._sessions = create_spool_session_factory(self._engine) def close(self) -> None: diff --git a/tests/huggingface/conftest.py b/tests/huggingface/conftest.py index 63542de..fc8e1a0 100644 --- a/tests/huggingface/conftest.py +++ b/tests/huggingface/conftest.py @@ -41,7 +41,7 @@ def hugging_face_module(tmp_path, hub): clock = ManualClock() available_api = {"value": hub} module = HuggingFaceModule( - {"spool_path": str(tmp_path / "events.sqlite3")}, + {"handover_database_path": str(tmp_path / "events.handovers.sqlite3")}, ModuleContext(os.getpid()), api_factory=lambda: available_api["value"], clock=clock, diff --git a/tests/huggingface/test_bundles.py b/tests/huggingface/test_bundles.py index 5822942..ff04080 100644 --- a/tests/huggingface/test_bundles.py +++ b/tests/huggingface/test_bundles.py @@ -480,7 +480,7 @@ def test_write_access_check_reports_reach_and_scope( if fail_head: hub.fail_next("repo_info", FakeHubHTTPError(503)) module = HuggingFaceModule( - {"spool_path": str(tmp_path / "events.sqlite3")}, + {"handover_database_path": str(tmp_path / "events.handovers.sqlite3")}, ModuleContext(os.getpid()), api_factory=lambda: hub, ) diff --git a/tests/huggingface/test_queue.py b/tests/huggingface/test_queue.py index 7b6f2eb..942301a 100644 --- a/tests/huggingface/test_queue.py +++ b/tests/huggingface/test_queue.py @@ -1,4 +1,4 @@ -"""The durable hand-over queue shares the spool database and its migrations.""" +"""The durable hand-over queue, in its own database next to the spool database.""" import shutil from pathlib import Path @@ -22,7 +22,7 @@ @pytest.fixture def queue(tmp_path): - handle = HandoverQueue(tmp_path / "events.sqlite3") + handle = HandoverQueue(tmp_path / "events.handovers.sqlite3") yield handle handle.close() @@ -49,17 +49,20 @@ def commit_everything(queue, *, repository=REPOSITORY, now, revision="commit-1") ) -def test_migration_adds_hand_over_tables_without_changing_telemetry(tmp_path): +def test_the_hand_over_database_leaves_the_spool_database_unchanged(tmp_path): path = tmp_path / "events.sqlite3" shutil.copyfile(FIXTURES / "microcosm-1099.sqlite3", path) + spool_before = path.read_bytes() + revision_before = current_database_revision(path) engine = create_spool_engine(path) with Session(engine) as session: events_before = session.scalar( select(func.count(TelemetryEventRecord.event_id)) ) engine.dispose() - HandoverQueue(path).close() - assert current_database_revision(path) == "20261020_03" + HandoverQueue(tmp_path / "events.handovers.sqlite3").close() + assert path.read_bytes() == spool_before + assert current_database_revision(path) == revision_before engine = create_spool_engine(path) with Session(engine) as session: assert ( @@ -196,8 +199,8 @@ def test_repositories_are_due_independently(queue): def test_one_writer_holds_a_repository_at_a_time(tmp_path): - first = HandoverQueue(tmp_path / "events.sqlite3") - second = HandoverQueue(tmp_path / "events.sqlite3") + first = HandoverQueue(tmp_path / "events.handovers.sqlite3") + second = HandoverQueue(tmp_path / "events.handovers.sqlite3") try: lease = first.claim_repository(REPOSITORY, now=1000, lease_seconds=60) assert lease is not None @@ -266,7 +269,7 @@ def test_receipt_counts_committed_files_and_commits_per_run(queue): def test_jobs_are_idempotent_by_hand_over_and_survive_a_restart(tmp_path): - path = tmp_path / "events.sqlite3" + path = tmp_path / "events.handovers.sqlite3" queue = HandoverQueue(path) request = {"repository": REPOSITORY, "run_id": "run-1"} first = queue.enqueue_job( diff --git a/tests/interface/conftest.py b/tests/interface/conftest.py index 39caf58..c067537 100644 --- a/tests/interface/conftest.py +++ b/tests/interface/conftest.py @@ -117,7 +117,7 @@ def emitter_with_hub(tmp_path): runtime_directory = Path(tempfile.mkdtemp(prefix="pe-stage-", dir="/tmp")) socket_path = runtime_directory / "service.sock" module = HuggingFaceModule( - {"spool_path": str(tmp_path / "events.sqlite3")}, + {"handover_database_path": str(tmp_path / "events.handovers.sqlite3")}, ModuleContext(os.getpid()), api_factory=lambda: hub, ) diff --git a/tests/interface/test_protocol_fixtures.py b/tests/interface/test_protocol_fixtures.py index fcc9ff6..10aff54 100644 --- a/tests/interface/test_protocol_fixtures.py +++ b/tests/interface/test_protocol_fixtures.py @@ -47,7 +47,7 @@ "6f576285b7a87607c47755ebee00ca051694b39455812a0068ebee2c27277728" ), "protocol-v1/huggingface-configuration.json": ( - "5a5ccd639953b24c774418e6e23735f515e46a673a68624e26d640839ef4b1c2" + "0ecea1a8e7ba6c4458c530b58fa90079207ee3f48da3e522479e9afde581551a" ), "protocol-v1/huggingface-read-back-run-files.json": ( "ab643a81520cbf35fada4362d6bf08f63c84d45ed41d9ddf1a0a592855c91313" diff --git a/tests/telemetry/test_persistence.py b/tests/telemetry/test_persistence.py index 5161166..c8fdbe3 100644 --- a/tests/telemetry/test_persistence.py +++ b/tests/telemetry/test_persistence.py @@ -37,7 +37,7 @@ def test_migrations_match_models_and_reopening_preserves_sequences(tmp_path): ] == [1, 2] assert spool.batch("fixture-run", "producer-a")[0] == first assert first["event_id"] != second["event_id"] - assert current_database_revision(path) == "20261020_03" + assert current_database_revision(path) == "20261008_02" with ( spool._engine.begin() as connection, alembic_config(connection=connection) as config, @@ -83,10 +83,5 @@ def test_multiple_producers_and_local_only_survive_reopening(tmp_path): "telemetry_runs", "telemetry_events", "graph_publication_jobs", - "staged_run_files", - "staged_run_deliveries", - "repository_commit_schedules", - "handover_jobs", - "handover_files", } spool.close() diff --git a/tools/installed-interface-smoke.py b/tools/installed-interface-smoke.py index edb5ffd..348cd11 100644 --- a/tools/installed-interface-smoke.py +++ b/tools/installed-interface-smoke.py @@ -63,6 +63,9 @@ def main(emitter_command: str) -> None: ) ) assert staged["status"] == "recorded" and staged["pending_files"] == 1 + # The installed module keeps hand-overs in their own database, created + # from the migrations shipped in the wheel, next to the spool database. + assert (Path.cwd() / "interface.handovers.sqlite3").is_file() finally: process = client._process client.close() diff --git a/tools/installed-smoke.py b/tools/installed-smoke.py index 1ab3b5c..a886317 100644 --- a/tools/installed-smoke.py +++ b/tools/installed-smoke.py @@ -56,7 +56,7 @@ def telemetry_smoke(): ) from microcosm_emitter.telemetry.service.spool import EventSpool - assert current_database_revision(path) == "20261020_03" + assert current_database_revision(path) == "20261008_02" spool = EventSpool(path) try: assert not spool.has_deliverable() From 94136362a2a14035ca2adce077eb2544be315966 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:13:12 +0400 Subject: [PATCH 12/23] Test that a long bundle upload holds up no other commit All repositories' commits ran one after another on the module's single worker thread, and the commit at shutdown ran on that thread after the current cycle. While a bundle upload to one repository was in progress, run files for another repository waited, the shutdown commit never ran before the host gave up, and a requested read-back answered pending. The tests run the module in a host runtime with a 2-second shutdown period and an in-memory Hub whose bundle commit blocks on an event. Red: 3 failed, 1 passed. Run files never reached the staging repository; the runtime reported "LocalServiceError: Service worker failed (TimeoutError)" at shutdown without the staging commit; the read-back answered "pending". The passing test checks that the blocked bundle still lands once released. Co-Authored-By: Claude Opus 5.5 --- tests/huggingface/test_delivery_workers.py | 237 +++++++++++++++++++++ 1 file changed, 237 insertions(+) create mode 100644 tests/huggingface/test_delivery_workers.py diff --git a/tests/huggingface/test_delivery_workers.py b/tests/huggingface/test_delivery_workers.py new file mode 100644 index 0000000..28c80fe --- /dev/null +++ b/tests/huggingface/test_delivery_workers.py @@ -0,0 +1,237 @@ +"""A long upload to one repository never holds up commits to another. + +The module runs inside a real host runtime with the in-memory Hub. A bundle +commit to the private repository blocks until the test releases it; run-file +commits to the staging repository, the commit at shutdown, and a read-back +must all complete meanwhile. +""" + +import hashlib +import os +import shutil +import tempfile +import threading +import time +from dataclasses import dataclass, field +from pathlib import Path +from types import SimpleNamespace + +import pytest +from fake_hugging_face import ( + PRIVATE_REPOSITORY, + STAGING_REPOSITORY, + FakeHuggingFaceApi, +) + +from microcosm_emitter_interface.bundles import BundleFile, StageDatasetBundle +from microcosm_emitter_interface.client import EmitterClient +from microcosm_emitter_interface.staging import RunFile, StageRunFiles +from microcosm_emitter_interface.transport import SocketTransport + +SHUTDOWN_SECONDS = 2.0 +BLOCKED_COMMIT_LIMIT_SECONDS = 30.0 + + +@dataclass +class InThreadEmitter: + """The client's view of an emitter host running in a test thread.""" + + transport: object + description: dict = field(default_factory=dict) + + def close(self): + self.transport.close() + + +def wait_until(condition, *, seconds: float) -> bool: + deadline = time.monotonic() + seconds + while time.monotonic() < deadline: + if condition(): + return True + time.sleep(0.02) + return condition() + + +@pytest.fixture +def running_module(tmp_path): + """The Hugging Face module in a host runtime; its next commit blocks.""" + from microcosm_emitter.host.contracts import ModuleContext + from microcosm_emitter.host.runtime import ServiceRuntime + from microcosm_emitter.huggingface.service.module import HuggingFaceModule + + hub = FakeHuggingFaceApi(whoami={"auth": {"accessToken": {"role": "write"}}}) + hub.create_repository(STAGING_REPOSITORY, {"README.md": b"staging fixture"}) + hub.create_repository(PRIVATE_REPOSITORY, {"README.md": b"private fixture"}) + commit_entered = threading.Event() + commit_released = threading.Event() + + def block_the_next_commit(): + commit_entered.set() + commit_released.wait(BLOCKED_COMMIT_LIMIT_SECONDS) + + hub.run_before_next_commit(block_the_next_commit) + runtime_directory = Path(tempfile.mkdtemp(prefix="pe-workers-", dir="/tmp")) + socket_path = runtime_directory / "service.sock" + module = HuggingFaceModule( + {"handover_database_path": str(tmp_path / "events.handovers.sqlite3")}, + ModuleContext(os.getpid()), + api_factory=lambda: hub, + ) + runtime = ServiceRuntime( + socket_path, + modules={"huggingface": module}, + is_parent_alive=lambda: True, + tick_seconds=0.01, + drain_seconds=SHUTDOWN_SECONDS, + ) + runtime_errors: list[BaseException] = [] + + def run_runtime(): + try: + runtime.run() + except BaseException as error: + runtime_errors.append(error) + + runtime_thread = threading.Thread(target=run_runtime, daemon=True) + runtime_thread.start() + transport = SocketTransport(socket_path) + assert wait_until(lambda: _answers_ping(transport), seconds=5) + client = EmitterClient( + InThreadEmitter( + transport, {"modules": ["huggingface"], "protocol_versions": [1]} + ) + ) + controls = SimpleNamespace( + hub=hub, + client=client, + tmp_path=tmp_path, + commit_entered=commit_entered, + commit_released=commit_released, + runtime_thread=runtime_thread, + runtime_errors=runtime_errors, + ) + try: + yield controls + finally: + commit_released.set() + client.close() + runtime_thread.join(timeout=SHUTDOWN_SECONDS + 10) + shutil.rmtree(runtime_directory, ignore_errors=True) + + +def _answers_ping(transport) -> bool: + try: + transport.ping() + except Exception: + return False + return True + + +def bundle_file(directory: Path, name: str, role: str, content: bytes) -> BundleFile: + path = directory / name + path.write_bytes(content) + return BundleFile( + name=name, + role=role, + local_path=path, + sha256=hashlib.sha256(content).hexdigest(), + bytes=len(content), + ) + + +def start_a_blocked_bundle_upload(controls) -> None: + """Hand over a bundle; return once its commit is in progress and blocked.""" + directory = controls.tmp_path / "bundle" + directory.mkdir() + request = StageDatasetBundle( + repository=PRIVATE_REPOSITORY, + run_id="uk-run-1", + files=( + bundle_file(directory, "dataset.h5", "dataset", b"rows"), + bundle_file( + directory, + "calibration_diagnostics.json", + "calibration_diagnostics", + b"{}", + ), + bundle_file( + directory, "staged_manifest.json", "staged_manifest", b'{"files": {}}' + ), + ), + ) + receipt = controls.client.stage_dataset_bundle(request) + assert receipt["status"] == "pending", receipt + assert controls.commit_entered.wait(10), "the bundle commit never started" + + +def stage_progress(controls, *, flush: bool) -> None: + progress = controls.tmp_path / "progress.json" + progress.write_bytes(b'{"stage":"completed"}') + receipt = controls.client.stage_run_files( + StageRunFiles( + repository=STAGING_REPOSITORY, + run_id="staging-run-1", + files=(RunFile("progress.json", progress),), + flush=flush, + commit_interval_seconds=3600, + ) + ) + assert receipt["status"] == "recorded", receipt + + +def staging_holds_the_progress_file(controls) -> bool: + return any( + "runs/staging-run-1/progress.json" in commit.added + for commit in controls.hub.commits(STAGING_REPOSITORY) + ) + + +def test_run_files_reach_another_repository_during_a_bundle_upload(running_module): + controls = running_module + start_a_blocked_bundle_upload(controls) + + stage_progress(controls, flush=True) + + assert wait_until(lambda: staging_holds_the_progress_file(controls), seconds=5) + assert controls.hub.commits(PRIVATE_REPOSITORY) == [] + + +def test_the_shutdown_commit_runs_during_a_bundle_upload(running_module): + controls = running_module + start_a_blocked_bundle_upload(controls) + # Not due for an hour: only the commit at shutdown can deliver it. + stage_progress(controls, flush=False) + assert not staging_holds_the_progress_file(controls) + + controls.client.close() + controls.runtime_thread.join(timeout=SHUTDOWN_SECONDS + 2) + + assert not controls.runtime_thread.is_alive() + assert controls.runtime_errors == [] + assert staging_holds_the_progress_file(controls) + assert controls.hub.commits(PRIVATE_REPOSITORY) == [] + + +def test_a_read_back_passes_during_a_bundle_upload(running_module): + controls = running_module + start_a_blocked_bundle_upload(controls) + stage_progress(controls, flush=False) + + result = controls.client.read_back_run_files( + STAGING_REPOSITORY, "staging-run-1", wait_seconds=10 + ) + + assert result["status"] == "passed", result + assert result["files"] == 1 + assert controls.hub.commits(PRIVATE_REPOSITORY) == [] + + +def test_the_bundle_commit_lands_once_its_upload_finishes(running_module): + controls = running_module + start_a_blocked_bundle_upload(controls) + + controls.commit_released.set() + + assert wait_until( + lambda: len(controls.hub.commits(PRIVATE_REPOSITORY)) == 1, seconds=10 + ) From 2f25a3a59a300f4141e378108d4277eeb053130c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:16:32 +0400 Subject: [PATCH 13/23] Upload bundles on their own worker so other commits never wait The host gives each module one worker thread, and runs the module's close (the commit at shutdown) on it only after the current tick returns. Bundle uploads ran inside that tick, so a long one held up every other repository's commits, the shutdown commit and read-backs. - workers.DeliveryWorker runs one kind of delivery cycle on a thread of its own. It can tell an upload in progress from an idle wait, so close knows when the queue is safe to release. Later kinds (release changes, collector submissions) use the same class. - HuggingFaceDelivery splits its cycle: run_file_cycle commits run files only, and job_cycle(kinds) delivers jobs of those kinds only. run_cycle still does both on the calling thread. - HuggingFaceModule starts a dataset-bundle worker. The host's tick commits run files only. close first commits every repository's pending run files that no worker is uploading to, then gives an upload in progress until the deadline, and commits that repository's run files after it if time remains. Leases in the hand-over database still keep two writers off one repository. - background_workers=False keeps every delivery on tick, for tests that drive the module by hand. The README's statement that the module commits pending files once more at shutdown now says which repository can miss that commit and why. Green: 4 passed; 192 passed across #23's changed test files. Co-Authored-By: Claude Opus 5.5 --- README.md | 6 +- .../huggingface-dataset-bundles.added.md | 2 +- .../huggingface/service/delivery.py | 102 +++++++++++++----- .../huggingface/service/module.py | 80 ++++++++++++-- .../huggingface/service/queue.py | 55 ++++++---- .../huggingface/service/workers.py | 83 ++++++++++++++ tests/huggingface/conftest.py | 2 + 7 files changed, 269 insertions(+), 61 deletions(-) create mode 100644 src/microcosm_emitter/huggingface/service/workers.py diff --git a/README.md b/README.md index 2f8825e..0c05923 100644 --- a/README.md +++ b/README.md @@ -64,8 +64,8 @@ The `huggingface` module takes over the uploads Microcosm made itself. Microcosm - Files land at `//` (by default `runs//...`), exactly as the build's writer lays them out. Staging contract version 1 also writes `latest_staging.json` at the repository root, and its `runs.json` index: at commit time the module merges the run's summary into the repository's copy at the head it commits on, keeping other runs, with Microcosm's ordering and JSON formatting. It starts a new index only when the Hub answers that the repository has no `runs.json`; any other download error (a connection error, a timeout, a 429 or 5xx response) fails the commit and retries it after the backoff. Microcosm started from an empty index on any download error, which could drop every other run. - Each path keeps only the newest bytes, and bytes equal to those last committed are not written again. -- **One commit per repository per cycle.** Each module tick commits, for each repository with work due, every pending path in one commit. An unforced hand-over is due one interval (Microcosm's upload interval) after this machine's last commit to the repository; a forced one (final states, artifacts) at the next cycle. Either way, two commits from this machine to one repository, of run files or bundles, are at least 10 seconds apart; only the last commit at shutdown ignores that spacing, and commits from other machines are not counted. Microcosm's writers made one commit per file, up to eight per upload cycle, and each commit triggered the dataset repository's webhook. -- Uploads stay best-effort. Requests answer within seconds or report that they were not accepted; a failed commit keeps the files queued and retries with backoff (up to five minutes); without a credential the files stay queued and one warning is printed. On shutdown the module commits pending files once more. Nothing raises into the build. +- **One commit per repository per cycle.** Each module tick (once a second, on the host's worker thread) commits, for each repository with run files due, every pending path in one commit. An unforced hand-over is due one interval (Microcosm's upload interval) after this machine's last commit to the repository; a forced one (final states, artifacts) at the next cycle. Either way, two commits from this machine to one repository, of run files or bundles, are at least 10 seconds apart; only the last commit at shutdown ignores that spacing, and commits from other machines are not counted. Microcosm's writers made one commit per file, up to eight per upload cycle, and each commit triggered the dataset repository's webhook. +- Uploads stay best-effort. Requests answer within seconds or report that they were not accepted; a failed commit keeps the files queued and retries with backoff (up to five minutes); without a credential the files stay queued and one warning is printed. On shutdown the module commits every repository's pending run files once more, on the host's thread and before it waits for any bundle upload. The one exception is a repository a bundle upload holds at that moment: its run files are committed after that upload, if the upload finishes within the shutdown period (15 seconds); otherwise they stay queued for the next emitter. Nothing raises into the build. - A commit that merges `runs.json` is pinned to the head it read (`parent_commit`); if another build commits first, the module reads the index again and merges again. - `read_back_run_files` commits pending files, then compares the Hub's copy of each committed file with the committed bytes. A run with no file this emitter committed (its hand-overs were refused, or another machine's emitter staged it) reads back as `failed` with `NO_COMMITTED_FILES`, never `passed`. @@ -77,7 +77,7 @@ The `huggingface` module takes over the uploads Microcosm made itself. Microcosm - Unlike Microcosm's single attempt, `REPOSITORY_UNAVAILABLE` and `UPLOAD_FAILED` are retried with backoff, and a commit that loses a race to another writer is checked and tried again. The receipt keeps the last error code while the job is pending. - `check_write_access` is Microcosm's staged-dataset preflight: whether a credential is present, the repository is reachable, and the Hub reports a write scope for it. -A repository with both a due bundle and due run files gets the bundle in one cycle and the run files in a later one, so each cycle still makes one commit per repository. +**Bundle uploads run on their own thread.** A bundle upload can take minutes, so the module uploads bundles on a worker thread of its own, while the host's worker thread commits run files and makes the commit at shutdown. A bundle upload to one repository therefore never delays run-file commits to another repository, the commit at shutdown, or a read-back. Each delivery leases its repository first, so this machine never commits to one repository twice at once: run files for the repository a bundle is uploading to wait for that upload, and then for the 10-second spacing. ## Development diff --git a/changelog.d/huggingface-dataset-bundles.added.md b/changelog.d/huggingface-dataset-bundles.added.md index fb557cf..1db66a8 100644 --- a/changelog.d/huggingface-dataset-bundles.added.md +++ b/changelog.d/huggingface-dataset-bundles.added.md @@ -1 +1 @@ -Add `stage_dataset_bundle`, `dataset_bundle_receipt` and `check_write_access`. A build hands over its staged dataset bundle with each file's role; the emitter writes it under `staged//` in one commit with Microcosm's message, paths and idempotency (`already_staged`, `REMOTE_DIFFERS`), and retries Hub failures from the durable queue. The emitter copies the bundle's small files (each up to 32 MiB, up to 64 MiB per hand-over) into its queue before answering, so the build may rewrite them at once, as Microcosm does with its build manifest, checksum listing and build record; it uploads larger files from their paths and never uploads one whose bytes changed after the hand-over. +Add `stage_dataset_bundle`, `dataset_bundle_receipt` and `check_write_access`. A build hands over its staged dataset bundle with each file's role; the emitter writes it under `staged//` in one commit with Microcosm's message, paths and idempotency (`already_staged`, `REMOTE_DIFFERS`), and retries Hub failures from the durable queue. The emitter copies the bundle's small files (each up to 32 MiB, up to 64 MiB per hand-over) into its queue before answering, so the build may rewrite them at once, as Microcosm does with its build manifest, checksum listing and build record; it uploads larger files from their paths and never uploads one whose bytes changed after the hand-over. Bundles upload on a worker thread of their own, so a long upload never delays run-file commits to another repository, the commit at shutdown, or a read-back. diff --git a/src/microcosm_emitter/huggingface/service/delivery.py b/src/microcosm_emitter/huggingface/service/delivery.py index af31d07..211a1c5 100644 --- a/src/microcosm_emitter/huggingface/service/delivery.py +++ b/src/microcosm_emitter/huggingface/service/delivery.py @@ -1,9 +1,11 @@ -"""The delivery cycle: at most one Hugging Face commit per repository per cycle. +"""Delivery cycles: at most one Hugging Face commit per repository per cycle. -Each cycle (one module tick) finds the repositories with work due, leases each -one, and makes one commit carrying every pending staging run file of that -repository. Failures keep the work queued and back off; nothing here raises -into the build. +Each cycle finds the repositories with work due, leases each one, and makes +one commit: a due job (a dataset bundle), or every pending staging run file of +that repository at once. The module runs the run-file cycle on the host's +thread and each job kind's cycle on its own worker thread (see ``workers``), +so a long upload holds only its own repository. Failures keep the work queued +and back off; nothing here raises into the build. """ from __future__ import annotations @@ -140,32 +142,64 @@ def warn_once(self, key: str, message: str) -> None: print(message, file=sys.stderr, flush=True) def run_cycle(self) -> int: - """Make at most one commit per due repository; return how many.""" + """Every kind of work, on the calling thread; return how many commits. + + A repository with a due job gets the job, and its run files wait for + a later cycle. Used when nothing runs the cycles below on their own + threads, for example a module driven by hand in tests. + """ + return self._cycle(job_kinds=SUPPORTED_JOB_KINDS, run_files=True) + + def run_file_cycle(self) -> int: + """Commit each due repository's pending run files; return how many.""" + return self._cycle(job_kinds=frozenset(), run_files=True) + + def job_cycle(self, kinds: frozenset[str]) -> int: + """Deliver one due job of ``kinds`` per repository; return how many.""" + return self._cycle(job_kinds=kinds & SUPPORTED_JOB_KINDS, run_files=False) + + def _cycle(self, *, job_kinds: frozenset[str], run_files: bool) -> int: repositories = self.queue.due_repositories( - self.clock(), kinds=SUPPORTED_JOB_KINDS + self.clock(), kinds=job_kinds, run_files=run_files ) if not repositories: return 0 api = self.api_factory() return sum( - self._deliver_repository(repository, api) for repository in repositories + self._deliver_repository( + repository, api, job_kinds=job_kinds, run_files=run_files + ) + for repository in repositories ) def drain(self, deadline: float) -> None: """At shutdown, commit every repository's pending run files once. Ignores the commit interval and spacing, but still leases each - repository. ``deadline`` is a ``time.monotonic`` value. + repository: one whose lease a worker holds, during a long upload, is + skipped. ``deadline`` is a ``time.monotonic`` value. """ api = self.api_factory() for repository in self.queue.repositories_with_pending_run_files(): if time.monotonic() >= deadline: return - self._deliver_repository(repository, api, run_files_only=True) + self._deliver_repository( + repository, api, job_kinds=frozenset(), run_files=True, at_shutdown=True + ) def _deliver_repository( - self, repository: str, api: Any | None, *, run_files_only: bool = False + self, + repository: str, + api: Any | None, + *, + job_kinds: frozenset[str], + run_files: bool, + at_shutdown: bool = False, ) -> int: + """One commit to ``repository``: a job of ``job_kinds``, else run files. + + ``at_shutdown`` commits pending run files whether or not they are due. + """ lease = self.queue.claim_repository( repository, now=self.clock(), lease_seconds=REPOSITORY_LEASE_SECONDS ) @@ -181,27 +215,38 @@ def _deliver_repository( ).start() try: if api is None: - self._credential_missing(held, run_files_only=run_files_only) + self._credential_missing( + held, + job_kinds=job_kinds, + run_files=run_files, + at_shutdown=at_shutdown, + ) return 0 hub = DatasetRepository(repository, api) - if not run_files_only: - job = self._claim_job(repository) - if job is not None: - held.hold_job(job) - return int(self._deliver_job(hub, held)) - if not self.queue.run_files_due(repository, self.clock()): - return 0 + job = self._claim_job(repository, job_kinds) if job_kinds else None + if job is not None: + held.hold_job(job) + return int(self._deliver_job(hub, held)) + if not self._run_files_to_commit(repository, run_files, at_shutdown): + return 0 return int(self._deliver_run_files(hub)) finally: held.stop() self.queue.release_repository(repository, held.repository_lease) - def _claim_job(self, repository: str) -> ClaimedJob | None: + def _run_files_to_commit( + self, repository: str, run_files: bool, at_shutdown: bool + ) -> bool: + return run_files and ( + at_shutdown or self.queue.run_files_due(repository, self.clock()) + ) + + def _claim_job(self, repository: str, kinds: frozenset[str]) -> ClaimedJob | None: return self.queue.claim_job( repository, now=self.clock(), lease_seconds=REPOSITORY_LEASE_SECONDS, - kinds=SUPPORTED_JOB_KINDS, + kinds=kinds, ) def _finish_job( @@ -251,10 +296,17 @@ def _deliver_job(self, hub: DatasetRepository, held: HeldLeases) -> bool: ) return outcome.committed - def _credential_missing(self, held: HeldLeases, *, run_files_only: bool) -> None: + def _credential_missing( + self, + held: HeldLeases, + *, + job_kinds: frozenset[str], + run_files: bool, + at_shutdown: bool, + ) -> None: repository = held.repository self.warn_once(CREDENTIAL_MISSING, CREDENTIAL_MISSING_MESSAGE) - job = None if run_files_only else self._claim_job(repository) + job = self._claim_job(repository, job_kinds) if job_kinds else None if job is not None: held.hold_job(job) held.stop() @@ -267,9 +319,7 @@ def _credential_missing(self, held: HeldLeases, *, run_files_only: bool) -> None retry=True, ) return - if not run_files_only and not self.queue.run_files_due( - repository, self.clock() - ): + if not self._run_files_to_commit(repository, run_files, at_shutdown): return entries = self.queue.pending_run_files(repository) index_entries = self.queue.pending_run_index_entries(repository) diff --git a/src/microcosm_emitter/huggingface/service/module.py b/src/microcosm_emitter/huggingface/service/module.py index af38ef3..88dbbd4 100644 --- a/src/microcosm_emitter/huggingface/service/module.py +++ b/src/microcosm_emitter/huggingface/service/module.py @@ -1,7 +1,12 @@ """The Hugging Face hand-over module of the local emitter host. Requests are queries: the module checks and durably records a hand-over, then -answers with its receipt. The module's worker makes the Hugging Face commits. +answers with its receipt. Commits happen on worker threads: + +- staging run-file commits, and the commit at shutdown, on the host's worker + thread (``tick`` and ``close``); +- dataset bundle uploads, which can take minutes, on the module's own + ``DeliveryWorker``, so they never delay those (see ``workers``). """ from __future__ import annotations @@ -11,6 +16,7 @@ from pathlib import Path from typing import Any +from microcosm_emitter.host.constants import TICK_SECONDS from microcosm_emitter.host.contracts import JsonObject, ModuleContext from microcosm_emitter.huggingface.service.delivery import ( CREDENTIAL_MISSING, @@ -29,6 +35,7 @@ require_supported_version, run_reference, ) +from microcosm_emitter.huggingface.service.workers import DeliveryWorker from microcosm_emitter_interface.identifiers import repository_identifier #: A read-back of a run with no file this emitter committed: nothing to compare. @@ -58,7 +65,12 @@ def _refused_run_files_receipt(message: Mapping[str, Any], code: str) -> dict[st class HuggingFaceModule: - """Queue hand-overs on request; commit them on the worker.""" + """Queue hand-overs on request; commit them on worker threads. + + ``background_workers`` false runs every delivery on the host's ``tick`` + instead of starting the module's own workers; tests that drive the + module by hand use it. + """ def __init__( self, @@ -67,11 +79,16 @@ def __init__( *, api_factory: Callable[[], Any | None] = hugging_face_api, clock: Callable[[], float] = time.time, + background_workers: bool = True, + worker_interval_seconds: float = TICK_SECONDS, ) -> None: self.configuration, self.context = configuration, context self.api_factory, self.clock = api_factory, clock + self.background_workers = background_workers + self.worker_interval_seconds = worker_interval_seconds self.queue: HandoverQueue | None = None self.delivery: HuggingFaceDelivery | None = None + self.workers: dict[str, DeliveryWorker] = {} self._closed = False def initialize(self) -> None: @@ -81,6 +98,27 @@ def initialize(self) -> None: self.delivery = HuggingFaceDelivery( self.queue, api_factory=self.api_factory, clock=self.clock ) + if self.background_workers: + self.workers = self._delivery_workers(self.delivery) + for worker in self.workers.values(): + worker.start() + + def _delivery_workers( + self, delivery: HuggingFaceDelivery + ) -> dict[str, DeliveryWorker]: + """One worker per kind of upload that can take minutes.""" + return { + DATASET_BUNDLE: DeliveryWorker( + "dataset-bundles", + lambda: delivery.job_cycle(frozenset({DATASET_BUNDLE})), + interval_seconds=self.worker_interval_seconds, + on_error=self._worker_error, + ) + } + + def _worker_error(self) -> None: + if self.delivery is not None: + self.delivery.warn_once("worker", WORKER_ERROR_MESSAGE) def handle_message(self, message: JsonObject) -> None: raise ValueError("Hugging Face hand-overs are queries with receipts.") @@ -246,25 +284,49 @@ def _read_back_run_files(self, message: JsonObject) -> dict[str, Any]: def tick(self, now: float) -> None: assert self.delivery is not None, "The host initializes modules first." try: - self.delivery.run_cycle() + if self.workers: + # Uploads that can take minutes run on the module's workers. + self.delivery.run_file_cycle() + else: + self.delivery.run_cycle() except Exception: # The host stops a module whose tick raises; a local queue error # leaves the work queued and must not end every later delivery. - self.delivery.warn_once("worker", WORKER_ERROR_MESSAGE) + self._worker_error() def parent_exited(self) -> None: """Queued hand-overs stay eligible for the next emitter host.""" def close(self, deadline: float) -> None: + """Commit pending run files once more, then stop by ``deadline``. + + The commit at shutdown runs first, on the calling thread, for every + repository a worker is not uploading to at that moment. A worker's + upload in progress then has until ``deadline`` to finish, and its + repository's run files are committed after it if time remains. An + upload still running at ``deadline`` is abandoned when the process + exits; its job lease expires and the next emitter retries the job. + """ if self._closed: return self._closed = True + for worker in self.workers.values(): + worker.stop() try: - if self.delivery is not None and time.monotonic() < deadline: - self.delivery.drain(deadline) + self._drain(deadline) + workers_stopped = all( + [worker.wait(deadline) for worker in self.workers.values()] + ) + if self.workers and workers_stopped: + self._drain(deadline) except Exception: - if self.delivery is not None: - self.delivery.warn_once("worker", WORKER_ERROR_MESSAGE) + self._worker_error() + workers_stopped = False finally: - if self.queue is not None: + # A worker still uploading needs the queue to record its result. + if self.queue is not None and (not self.workers or workers_stopped): self.queue.close() + + def _drain(self, deadline: float) -> None: + if self.delivery is not None and time.monotonic() < deadline: + self.delivery.drain(deadline) diff --git a/src/microcosm_emitter/huggingface/service/queue.py b/src/microcosm_emitter/huggingface/service/queue.py index 86e518e..f0c4b5b 100644 --- a/src/microcosm_emitter/huggingface/service/queue.py +++ b/src/microcosm_emitter/huggingface/service/queue.py @@ -191,37 +191,48 @@ def record_commit(self, repository: str, *, now: float) -> None: self._schedule(session, repository).last_commit_at = now def due_repositories( - self, now: float, *, kinds: Iterable[str] | None = None + self, + now: float, + *, + kinds: Iterable[str] | None = None, + run_files: bool = True, ) -> list[str]: """Repositories with work due now that may receive a commit now. ``kinds`` limits jobs to those this emitter can deliver, so a job - queued by a newer emitter sharing the database waits for it. + queued by a newer emitter sharing the database waits for it; an empty + ``kinds`` leaves jobs out. ``run_files`` false leaves out run files. """ + job_kinds = None if kinds is None else list(kinds) with self._sessions() as session: - run_file_repositories = set( - session.scalars( - select(StagedRunFile.repository).where( - StagedRunFile.pending.is_(True), StagedRunFile.due_at <= now + run_file_repositories: set[str] = set() + if run_files: + run_file_repositories = set( + session.scalars( + select(StagedRunFile.repository).where( + StagedRunFile.pending.is_(True), + StagedRunFile.due_at <= now, + ) ) - ) - ) | set( - session.scalars( - select(StagedRunDelivery.repository).where( - StagedRunDelivery.run_index_pending.is_(True), - StagedRunDelivery.run_index_due_at <= now, + ) | set( + session.scalars( + select(StagedRunDelivery.repository).where( + StagedRunDelivery.run_index_pending.is_(True), + StagedRunDelivery.run_index_due_at <= now, + ) ) ) - ) - job_query = select(HandoverJob.target).where( - HandoverJob.destination == DESTINATION_HUGGING_FACE, - HandoverJob.status == JOB_PENDING, - HandoverJob.next_attempt_at <= now, - HandoverJob.lease_until <= now, - ) - if kinds is not None: - job_query = job_query.where(HandoverJob.kind.in_(list(kinds))) - job_repositories = set(session.scalars(job_query)) + job_repositories: set[str] = set() + if job_kinds != []: + job_query = select(HandoverJob.target).where( + HandoverJob.destination == DESTINATION_HUGGING_FACE, + HandoverJob.status == JOB_PENDING, + HandoverJob.next_attempt_at <= now, + HandoverJob.lease_until <= now, + ) + if job_kinds is not None: + job_query = job_query.where(HandoverJob.kind.in_(job_kinds)) + job_repositories = set(session.scalars(job_query)) due = [] for repository in sorted(run_file_repositories | job_repositories): schedule = session.get(RepositoryCommitSchedule, repository) diff --git a/src/microcosm_emitter/huggingface/service/workers.py b/src/microcosm_emitter/huggingface/service/workers.py new file mode 100644 index 0000000..9c66d0f --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/workers.py @@ -0,0 +1,83 @@ +"""Delivery workers: each kind of long upload on its own thread. + +The host gives a module one worker thread. It calls ``tick`` once per tick and, +when the host stops, ``close``, which makes the commit at shutdown; ``close`` +runs only after the current ``tick`` returns. So nothing that can take minutes +may run inside ``tick``. The Hugging Face module therefore commits staging run +files (small, at most one commit per repository) on the host's thread, and +gives each kind of upload that can take minutes a ``DeliveryWorker``: its own +thread, running that kind's delivery cycle once per interval. + +A worker's upload then never delays run-file commits to another repository, +the commit at shutdown, or a read-back. Two writers never commit to one +repository at once: each delivery first leases its repository in the +hand-over database, and one that finds it leased tries again next cycle. +""" + +from __future__ import annotations + +import threading +import time +from collections.abc import Callable + + +class DeliveryWorker: + """Runs ``cycle`` on its own thread every ``interval_seconds`` until stopped. + + An exception from ``cycle`` is reported through ``on_error``; the work it + was doing stays queued, and the next cycle tries again. + """ + + def __init__( + self, + name: str, + cycle: Callable[[], object], + *, + interval_seconds: float, + on_error: Callable[[], None], + ) -> None: + self.name = name + self.cycle = cycle + self.interval_seconds = interval_seconds + self.on_error = on_error + self._stopping = threading.Event() + # Held for the length of each cycle, so ``wait`` can tell a worker + # that is uploading from one that is only waiting for its next cycle. + self._in_cycle = threading.Lock() + self._thread = threading.Thread( + target=self._run, name=f"microcosm-emitter-{name}", daemon=True + ) + + def start(self) -> None: + self._thread.start() + + def stop(self) -> None: + """Start no further cycle; a cycle in progress runs on.""" + self._stopping.set() + + def wait(self, deadline: float) -> bool: + """After ``stop``, wait for a cycle in progress until ``deadline``. + + ``deadline`` is a ``time.monotonic`` value. True means no cycle is + running and none will start; False means one is still running, for + example an upload that outlasts the shutdown period. + """ + remaining = deadline - time.monotonic() + acquired = ( + self._in_cycle.acquire(timeout=remaining) + if remaining > 0 + else self._in_cycle.acquire(blocking=False) + ) + if acquired: + self._in_cycle.release() + return acquired + + def _run(self) -> None: + while not self._stopping.wait(self.interval_seconds): + with self._in_cycle: + if self._stopping.is_set(): + return + try: + self.cycle() + except Exception: + self.on_error() diff --git a/tests/huggingface/conftest.py b/tests/huggingface/conftest.py index fc8e1a0..00c0345 100644 --- a/tests/huggingface/conftest.py +++ b/tests/huggingface/conftest.py @@ -45,6 +45,8 @@ def hugging_face_module(tmp_path, hub): ModuleContext(os.getpid()), api_factory=lambda: available_api["value"], clock=clock, + # Each test runs the delivery cycles itself, on its own thread. + background_workers=False, ) module.initialize() controls = SimpleNamespace( From 97a59d9f13787177c9bde1ee7a1b2980ffd94863 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:19:55 +0400 Subject: [PATCH 14/23] Test the retry cap, a stale-bytes race and a receipt without an ID From the #23 review's code-quality notes: - retry_delay capped its exponent at 8, so retries never waited the 300 seconds MAX_RETRY_SECONDS and the README promise; - bytes equal to the last committed ones, handed over again while other bytes for the path were being committed, were marked not pending and their copy cleared, so the repository kept the older commit's bytes; - the interface's wait for a bundle to settle indexed the receipt's handover_id outside its error handling. Red: 3 failed, 34 passed. "assert [] == [b'X']"; "At index 3 diff: 256.0 != 300.0"; "KeyError: 'handover_id'". Co-Authored-By: Claude Opus 5.5 --- tests/huggingface/test_queue.py | 32 ++++++++++++++++++++++ tests/interface/test_client_errors.py | 39 +++++++++++++++++++++++++++ 2 files changed, 71 insertions(+) diff --git a/tests/huggingface/test_queue.py b/tests/huggingface/test_queue.py index 942301a..2568b49 100644 --- a/tests/huggingface/test_queue.py +++ b/tests/huggingface/test_queue.py @@ -243,6 +243,38 @@ def test_a_newer_hand_over_during_a_commit_stays_pending(queue): } +def test_earlier_bytes_handed_over_again_during_a_commit_are_committed_after_it( + queue, +): + """The repository holds X; Y is being committed when X is handed over again.""" + path_in_repo = "runs/run-1/progress.json" + stage(queue, path_in_repo, b"X", now=1000) + commit_everything(queue, now=1001, revision="commit-x") + stage(queue, path_in_repo, b"Y", now=1100) + claimed = queue.pending_run_files(REPOSITORY) + stage(queue, path_in_repo, b"X", now=1101) + queue.record_run_files_committed( + REPOSITORY, claimed, [], revision="commit-y", now=1102 + ) + assert [entry.content for entry in queue.pending_run_files(REPOSITORY)] == [b"X"] + + +def test_retries_back_off_up_to_five_minutes(): + from microcosm_emitter.huggingface.service.queue import ( + MAX_RETRY_SECONDS, + retry_delay, + ) + + assert MAX_RETRY_SECONDS == 300.0 + assert [retry_delay(failures) for failures in (1, 2, 8, 9, 64)] == [ + 2.0, + 4.0, + 256.0, + 300.0, + 300.0, + ] + + def test_receipt_counts_committed_files_and_commits_per_run(queue): queue.stage_run_files( repository=REPOSITORY, diff --git a/tests/interface/test_client_errors.py b/tests/interface/test_client_errors.py index d297715..e8de8cb 100644 --- a/tests/interface/test_client_errors.py +++ b/tests/interface/test_client_errors.py @@ -321,3 +321,42 @@ def bundle_file(name, role, local_path): if running: assert transport.requests == [] assert client.last_failure.reason == "invalid_arguments" + + +def test_a_pending_bundle_receipt_without_an_identifier_is_returned_not_raised( + tmp_path, +): + import hashlib + + from microcosm_emitter_interface.bundles import BundleFile, StageDatasetBundle + + client, transport = hugging_face_client() + transport.query = lambda module, message, *, timeout=None: { + "receipt": {"status": "pending"} + } + manifest = tmp_path / "staged_manifest.json" + manifest.write_bytes(b"{}") + request = StageDatasetBundle( + repository="policyengine/fixture-staging", + run_id="run-1", + files=( + BundleFile( + name="staged_manifest.json", + role="staged_manifest", + local_path=manifest, + sha256=hashlib.sha256(b"{}").hexdigest(), + bytes=2, + ), + BundleFile( + name="dataset.h5", + role="dataset", + local_path=manifest, + sha256=hashlib.sha256(b"{}").hexdigest(), + bytes=2, + ), + ), + ) + + receipt = client.stage_dataset_bundle(request, wait_seconds=1) + + assert receipt == {"status": "pending"} From 191f2ffe8820cdd829d9a9776242ce75420185e1 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 21:20:03 +0400 Subject: [PATCH 15/23] Reach the five-minute retry cap, recommit re-handed bytes, guard waits - retry_delay lets its exponent grow until MAX_RETRY_SECONDS (300 s) applies, as the README says. - A path handed over again with the bytes last committed, while other bytes for it are pending, keeps those bytes; when the in-progress commit is recorded and the row's latest bytes differ from what it committed, the row is pending again, so the next commit restores the newest bytes. - The interface waits for a hand-over to settle only when its receipt names a hand-over ID, and otherwise returns the receipt as answered. Green: 37 passed. Co-Authored-By: Claude Opus 5.5 --- .../src/microcosm_emitter_interface/client.py | 6 +++++- src/microcosm_emitter/huggingface/service/queue.py | 14 +++++++++++--- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py index ceed801..04d6d56 100644 --- a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py @@ -430,6 +430,10 @@ def _handover_receipt(self, handover_id: str) -> dict[str, Any] | None: def _wait_until_settled( self, receipt: dict[str, Any], wait_seconds: float ) -> dict[str, Any]: + handover_id = receipt.get("handover_id") + if not isinstance(handover_id, str) or not handover_id: + # Nothing to ask about: return the receipt as answered. + return receipt deadline = time.monotonic() + _wait_seconds(wait_seconds) while ( receipt.get("status") not in SETTLED_STATUSES @@ -437,7 +441,7 @@ def _wait_until_settled( and time.monotonic() < deadline ): time.sleep(0.1) - latest = self._handover_receipt(receipt["handover_id"]) + latest = self._handover_receipt(handover_id) if latest is None: break receipt = latest diff --git a/src/microcosm_emitter/huggingface/service/queue.py b/src/microcosm_emitter/huggingface/service/queue.py index f0c4b5b..f01927a 100644 --- a/src/microcosm_emitter/huggingface/service/queue.py +++ b/src/microcosm_emitter/huggingface/service/queue.py @@ -69,7 +69,8 @@ def _entry(text: str | None) -> dict[str, Any]: def retry_delay(failures: int) -> float: - return min(MAX_RETRY_SECONDS, float(2 ** min(failures, 8))) + """2, 4, 8 … seconds after each failure, at most ``MAX_RETRY_SECONDS``.""" + return min(MAX_RETRY_SECONDS, float(2 ** min(failures, 16))) @dataclass(frozen=True) @@ -312,8 +313,12 @@ def stage_run_files( ) continue if digest == row.committed_sha256: - # The repository already holds exactly these bytes. - row.pending, row.content = False, b"" + # The repository already holds exactly these bytes. If + # other bytes for this path are pending, a commit of them + # may be in progress: keep these bytes, which + # record_run_files_committed queues again after it. + row.content = content if row.pending else b"" + row.pending = False elif not (row.pending and digest == row.content_sha256): row.due_at = min(row.due_at, due_at) if row.pending else due_at row.content, row.pending = content, True @@ -446,6 +451,9 @@ def record_run_files_committed( row.committed_sha256 = entry.sha256 if row.content_sha256 == entry.sha256: row.pending, row.content = False, b"" + else: + # Other bytes were handed over during this commit. + row.pending = True counts = self._files_per_run(entries) index_runs = {run_id for run_id, _ in index_entries} for run_id in sorted(set(counts) | index_runs): From 792b9243aac27e3723dc2a72bfb11b71734992fe Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:18:29 +0400 Subject: [PATCH 16/23] Test that no bundle upload starts once shutdown begins The #23 fix review found that after `close` calls the bundle worker's `stop`, the cycle already in progress still leases, claims and uploads the remaining repositories on its list. When the upload to a second repository outlasts the shutdown period, `wait` returns false, the second shutdown commit is skipped, and the first repository's run files miss it although that repository's upload finished in time; the README says they are committed. The test queues bundles for two repositories before the worker starts, holds the first upload, begins shutdown, then releases it after half a second. It asserts that no upload to the second repository starts and that the first repository's run files land at shutdown. Red: 1 failed, 4 passed (`AssertionError: an upload to the second repository started after shutdown began`). Co-Authored-By: Claude Opus 5.5 --- tests/huggingface/test_delivery_workers.py | 120 +++++++++++++++++++++ 1 file changed, 120 insertions(+) diff --git a/tests/huggingface/test_delivery_workers.py b/tests/huggingface/test_delivery_workers.py index 28c80fe..22c4625 100644 --- a/tests/huggingface/test_delivery_workers.py +++ b/tests/huggingface/test_delivery_workers.py @@ -235,3 +235,123 @@ def test_the_bundle_commit_lands_once_its_upload_finishes(running_module): assert wait_until( lambda: len(controls.hub.commits(PRIVATE_REPOSITORY)) == 1, seconds=10 ) + + +def bundle_message(directory: Path, repository: str, run_id: str) -> dict: + directory.mkdir(parents=True) + return StageDatasetBundle( + repository=repository, + run_id=run_id, + files=( + bundle_file(directory, "dataset.h5", "dataset", b"rows"), + bundle_file( + directory, "staged_manifest.json", "staged_manifest", b'{"files": {}}' + ), + ), + ).message() + + +def queue_bundles_before_the_workers_start( + database_path: Path, hub, directory: Path, repositories: list[str] +) -> None: + """Queue one bundle per repository, so the worker's first cycle lists all.""" + from microcosm_emitter.host.contracts import ModuleContext + from microcosm_emitter.huggingface.service.module import HuggingFaceModule + + queueing_module = HuggingFaceModule( + {"handover_database_path": str(database_path)}, + ModuleContext(os.getpid()), + api_factory=lambda: hub, + background_workers=False, + ) + queueing_module.initialize() + for position, repository in enumerate(repositories): + receipt = queueing_module.query( + bundle_message( + directory / f"bundle-{position}", repository, f"run-{position}" + ) + )["receipt"] + assert receipt["status"] == "pending", receipt + queueing_module.close(time.monotonic() + SHUTDOWN_SECONDS) + + +def test_no_bundle_upload_starts_once_shutdown_begins(tmp_path): + """Shutdown during one bundle upload: the worker starts no other upload. + + The worker's cycle lists two repositories. Shutdown begins while the + upload to the first is held; that upload then finishes within the + shutdown period. No upload to the second repository may start, and the + first repository's run files, which the first shutdown commit skipped + because the upload held that repository, get the second shutdown commit. + """ + from microcosm_emitter.host.contracts import ModuleContext + from microcosm_emitter.huggingface.service.module import HuggingFaceModule + + hub = FakeHuggingFaceApi(whoami={"auth": {"accessToken": {"role": "write"}}}) + hub.create_repository(PRIVATE_REPOSITORY, {"README.md": b"private fixture"}) + hub.create_repository(STAGING_REPOSITORY, {"README.md": b"staging fixture"}) + # The cycle visits repositories in name order: the private one first. + assert PRIVATE_REPOSITORY < STAGING_REPOSITORY + database_path = tmp_path / "events.handovers.sqlite3" + queue_bundles_before_the_workers_start( + database_path, hub, tmp_path, [PRIVATE_REPOSITORY, STAGING_REPOSITORY] + ) + + first_upload_entered = threading.Event() + first_upload_released = threading.Event() + second_upload_started = threading.Event() + test_finished = threading.Event() + + def hold_the_first_upload(): + first_upload_entered.set() + first_upload_released.wait(BLOCKED_COMMIT_LIMIT_SECONDS) + + hub.run_before_next_commit(hold_the_first_upload) + create_commit_on_the_hub = hub.create_commit + + def create_commit_holding_the_second_repository(repo_id, operations, **options): + if repo_id == STAGING_REPOSITORY: + # An upload that outlasts the shutdown period. + second_upload_started.set() + test_finished.wait(BLOCKED_COMMIT_LIMIT_SECONDS) + return create_commit_on_the_hub(repo_id, operations, **options) + + hub.create_commit = create_commit_holding_the_second_repository + module = HuggingFaceModule( + {"handover_database_path": str(database_path)}, + ModuleContext(os.getpid()), + api_factory=lambda: hub, + worker_interval_seconds=0.05, + ) + module.initialize() + try: + assert first_upload_entered.wait(10), "the first bundle upload never started" + progress = tmp_path / "progress.json" + progress.write_bytes(b'{"stage":"completed"}') + # Not due for an hour: only a commit at shutdown can deliver it. + run_files_receipt = module.query( + StageRunFiles( + repository=PRIVATE_REPOSITORY, + run_id="staging-run-1", + files=(RunFile("progress.json", progress),), + flush=False, + commit_interval_seconds=3600, + ).message() + )["receipt"] + assert run_files_receipt["status"] == "recorded", run_files_receipt + releaser = threading.Timer(0.5, first_upload_released.set) + releaser.start() + + module.close(time.monotonic() + SHUTDOWN_SECONDS) + + releaser.join() + assert not second_upload_started.is_set(), ( + "an upload to the second repository started after shutdown began" + ) + assert any( + "runs/staging-run-1/progress.json" in commit.added + for commit in hub.commits(PRIVATE_REPOSITORY) + ), "the first repository's run files missed the second shutdown commit" + finally: + first_upload_released.set() + test_finished.set() From fdbbf105dbf9115165998252b8c3afeed9b3f2a2 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:19:21 +0400 Subject: [PATCH 17/23] Start no bundle upload once shutdown begins `DeliveryWorker` now calls its cycle with `stop_requested`, which turns true when `close` calls `stop`. The job cycle checks it before each repository on its list and again once that repository is leased, so after shutdown begins the bundle worker starts no new upload. Only an upload already under way can then hold a repository during the commit at shutdown, which is the one exception the README names; that repository's run files get the second shutdown commit when the upload finishes in time. Cycles that no worker runs (the host's tick, and modules driven by hand) use `never_stop`. Green: 5 passed. Co-Authored-By: Claude Opus 5.5 --- README.md | 2 +- .../huggingface/service/delivery.py | 53 +++++++++++++++---- .../huggingface/service/module.py | 15 +++--- .../huggingface/service/workers.py | 18 +++++-- 4 files changed, 68 insertions(+), 20 deletions(-) diff --git a/README.md b/README.md index 0c05923..0003d45 100644 --- a/README.md +++ b/README.md @@ -77,7 +77,7 @@ The `huggingface` module takes over the uploads Microcosm made itself. Microcosm - Unlike Microcosm's single attempt, `REPOSITORY_UNAVAILABLE` and `UPLOAD_FAILED` are retried with backoff, and a commit that loses a race to another writer is checked and tried again. The receipt keeps the last error code while the job is pending. - `check_write_access` is Microcosm's staged-dataset preflight: whether a credential is present, the repository is reachable, and the Hub reports a write scope for it. -**Bundle uploads run on their own thread.** A bundle upload can take minutes, so the module uploads bundles on a worker thread of its own, while the host's worker thread commits run files and makes the commit at shutdown. A bundle upload to one repository therefore never delays run-file commits to another repository, the commit at shutdown, or a read-back. Each delivery leases its repository first, so this machine never commits to one repository twice at once: run files for the repository a bundle is uploading to wait for that upload, and then for the 10-second spacing. +**Bundle uploads run on their own thread.** A bundle upload can take minutes, so the module uploads bundles on a worker thread of its own, while the host's worker thread commits run files and makes the commit at shutdown. A bundle upload to one repository therefore never delays run-file commits to another repository, the commit at shutdown, or a read-back. Once shutdown begins, the bundle worker starts no new upload, even to a repository left in its current pass, so only an upload already under way can hold a repository during the commit at shutdown. Each delivery leases its repository first, so this machine never commits to one repository twice at once: run files for the repository a bundle is uploading to wait for that upload, and then for the 10-second spacing. ## Development diff --git a/src/microcosm_emitter/huggingface/service/delivery.py b/src/microcosm_emitter/huggingface/service/delivery.py index 211a1c5..8e5f496 100644 --- a/src/microcosm_emitter/huggingface/service/delivery.py +++ b/src/microcosm_emitter/huggingface/service/delivery.py @@ -45,6 +45,11 @@ ) +def never_stop() -> bool: + """The stop check of a cycle that no worker runs: it is never stopped.""" + return False + + def iso_timestamp(seconds: float) -> str: """Microcosm ``staging._now()`` format, at ``seconds``.""" return datetime.fromtimestamp(seconds, UTC).replace(microsecond=0).isoformat() @@ -154,23 +159,48 @@ def run_file_cycle(self) -> int: """Commit each due repository's pending run files; return how many.""" return self._cycle(job_kinds=frozenset(), run_files=True) - def job_cycle(self, kinds: frozenset[str]) -> int: - """Deliver one due job of ``kinds`` per repository; return how many.""" - return self._cycle(job_kinds=kinds & SUPPORTED_JOB_KINDS, run_files=False) + def job_cycle( + self, + kinds: frozenset[str], + *, + stop_requested: Callable[[], bool] = never_stop, + ) -> int: + """Deliver one due job of ``kinds`` per repository; return how many. - def _cycle(self, *, job_kinds: frozenset[str], run_files: bool) -> int: + Once ``stop_requested`` turns true (shutdown has begun), the cycle + starts no new upload and returns. + """ + return self._cycle( + job_kinds=kinds & SUPPORTED_JOB_KINDS, + run_files=False, + stop_requested=stop_requested, + ) + + def _cycle( + self, + *, + job_kinds: frozenset[str], + run_files: bool, + stop_requested: Callable[[], bool] = never_stop, + ) -> int: repositories = self.queue.due_repositories( self.clock(), kinds=job_kinds, run_files=run_files ) if not repositories: return 0 api = self.api_factory() - return sum( - self._deliver_repository( - repository, api, job_kinds=job_kinds, run_files=run_files + delivered = 0 + for repository in repositories: + if stop_requested(): + break + delivered += self._deliver_repository( + repository, + api, + job_kinds=job_kinds, + run_files=run_files, + stop_requested=stop_requested, ) - for repository in repositories - ) + return delivered def drain(self, deadline: float) -> None: """At shutdown, commit every repository's pending run files once. @@ -195,10 +225,13 @@ def _deliver_repository( job_kinds: frozenset[str], run_files: bool, at_shutdown: bool = False, + stop_requested: Callable[[], bool] = never_stop, ) -> int: """One commit to ``repository``: a job of ``job_kinds``, else run files. ``at_shutdown`` commits pending run files whether or not they are due. + ``stop_requested`` is checked again once the repository is leased, so + a shutdown that began while the lease was being taken starts no upload. """ lease = self.queue.claim_repository( repository, now=self.clock(), lease_seconds=REPOSITORY_LEASE_SECONDS @@ -214,6 +247,8 @@ def _deliver_repository( renewal_seconds=self.lease_renewal_seconds, ).start() try: + if stop_requested(): + return 0 if api is None: self._credential_missing( held, diff --git a/src/microcosm_emitter/huggingface/service/module.py b/src/microcosm_emitter/huggingface/service/module.py index 88dbbd4..48c6fb8 100644 --- a/src/microcosm_emitter/huggingface/service/module.py +++ b/src/microcosm_emitter/huggingface/service/module.py @@ -110,7 +110,9 @@ def _delivery_workers( return { DATASET_BUNDLE: DeliveryWorker( "dataset-bundles", - lambda: delivery.job_cycle(frozenset({DATASET_BUNDLE})), + lambda stop_requested: delivery.job_cycle( + frozenset({DATASET_BUNDLE}), stop_requested=stop_requested + ), interval_seconds=self.worker_interval_seconds, on_error=self._worker_error, ) @@ -301,11 +303,12 @@ def close(self, deadline: float) -> None: """Commit pending run files once more, then stop by ``deadline``. The commit at shutdown runs first, on the calling thread, for every - repository a worker is not uploading to at that moment. A worker's - upload in progress then has until ``deadline`` to finish, and its - repository's run files are committed after it if time remains. An - upload still running at ``deadline`` is abandoned when the process - exits; its job lease expires and the next emitter retries the job. + repository a worker is not uploading to at that moment. Once stopped, + a worker starts no new upload. Its upload in progress has until + ``deadline`` to finish, and that repository's run files are committed + after it if time remains. An upload still running at ``deadline`` is + abandoned when the process exits; its job lease expires and the next + emitter retries the job. """ if self._closed: return diff --git a/src/microcosm_emitter/huggingface/service/workers.py b/src/microcosm_emitter/huggingface/service/workers.py index 9c66d0f..a5e42a6 100644 --- a/src/microcosm_emitter/huggingface/service/workers.py +++ b/src/microcosm_emitter/huggingface/service/workers.py @@ -12,6 +12,10 @@ the commit at shutdown, or a read-back. Two writers never commit to one repository at once: each delivery first leases its repository in the hand-over database, and one that finds it leased tries again next cycle. + +Once shutdown begins (``stop``), a worker starts no new upload: its cycle +receives ``stop_requested`` and checks it before each repository, so only +an upload already under way when ``stop`` was called can still be running. """ from __future__ import annotations @@ -24,14 +28,16 @@ class DeliveryWorker: """Runs ``cycle`` on its own thread every ``interval_seconds`` until stopped. - An exception from ``cycle`` is reported through ``on_error``; the work it + ``cycle`` is called with ``stop_requested``, which turns true once + ``stop`` is called; a cycle checks it before it starts each upload. An + exception from ``cycle`` is reported through ``on_error``; the work it was doing stays queued, and the next cycle tries again. """ def __init__( self, name: str, - cycle: Callable[[], object], + cycle: Callable[[Callable[[], bool]], object], *, interval_seconds: float, on_error: Callable[[], None], @@ -52,9 +58,13 @@ def start(self) -> None: self._thread.start() def stop(self) -> None: - """Start no further cycle; a cycle in progress runs on.""" + """Start no further cycle or upload; an upload under way runs on.""" self._stopping.set() + def stop_requested(self) -> bool: + """Whether ``stop`` was called: the cycle then starts no new upload.""" + return self._stopping.is_set() + def wait(self, deadline: float) -> bool: """After ``stop``, wait for a cycle in progress until ``deadline``. @@ -78,6 +88,6 @@ def _run(self) -> None: if self._stopping.is_set(): return try: - self.cycle() + self.cycle(self.stop_requested) except Exception: self.on_error() From 7f62a2ad6bccc920ec795a73188c0a96ce1377be Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:20:02 +0400 Subject: [PATCH 18/23] Test that the commit spacing holds across the host's thread and a worker The #23 fix review found that the 10-second spacing between this machine's commits to one repository is checked only when a cycle builds its list of due repositories. Since bundles moved to their own worker, the host's thread can commit to a repository on the worker's list (or the other way round) while the worker uploads elsewhere, and the second commit then follows at once. Two tests, each with the clock standing still, make the other thread's commit from inside the first commit of a cycle: - a bundle cycle lists the private and staging repositories, and the staging run files are committed during the private upload: the staging bundle must wait for the spacing; - a run-file cycle lists both repositories, and the staging bundle is committed during the private run-file commit: the staging run files must wait for the spacing. Red: 2 failed, 5 passed (`AssertionError: the staging bundle was committed less than 10 seconds after its run files`; `AssertionError: the staging run files were committed less than 10 seconds after the bundle`). Co-Authored-By: Claude Opus 5.5 --- tests/huggingface/test_delivery_workers.py | 92 ++++++++++++++++++++++ 1 file changed, 92 insertions(+) diff --git a/tests/huggingface/test_delivery_workers.py b/tests/huggingface/test_delivery_workers.py index 22c4625..1781f8b 100644 --- a/tests/huggingface/test_delivery_workers.py +++ b/tests/huggingface/test_delivery_workers.py @@ -355,3 +355,95 @@ def create_commit_holding_the_second_repository(repo_id, operations, **options): finally: first_upload_released.set() test_finished.set() + + +# The 10-second spacing between commits to one repository, across threads. +# +# The host's thread (run files) and the bundle worker each build a list of due +# repositories, then deliver them in turn. Another thread can commit to a +# repository on that list meanwhile, so the spacing must be checked again +# once the repository is leased. Each test makes the other thread's commit +# from inside the first commit of the cycle, with the clock standing still. + + +def hand_over_run_file(controls, repository: str, run_id: str) -> None: + local = controls.tmp_path / run_id / "progress.json" + local.parent.mkdir(parents=True, exist_ok=True) + local.write_bytes(b'{"stage":"calibrating"}') + receipt = controls.module.query( + StageRunFiles( + repository=repository, + run_id=run_id, + files=(RunFile("progress.json", local),), + flush=True, + ).message() + )["receipt"] + assert receipt["status"] == "recorded", receipt + + +def hand_over_bundle(controls, repository: str, run_id: str) -> None: + receipt = controls.module.query( + bundle_message(controls.tmp_path / f"bundle-{run_id}", repository, run_id) + )["receipt"] + assert receipt["status"] == "pending", receipt + + +def bundle_cycle(controls) -> int: + from microcosm_emitter.huggingface.service.delivery import DATASET_BUNDLE + + return controls.module.delivery.job_cycle(frozenset({DATASET_BUNDLE})) + + +def test_a_bundle_waits_for_the_spacing_after_a_run_file_commit_during_its_cycle( + hugging_face_module, +): + from microcosm_emitter.huggingface.service.queue import ( + MIN_COMMIT_SPACING_SECONDS, + ) + + controls = hugging_face_module + hand_over_bundle(controls, PRIVATE_REPOSITORY, "uk-run-1") + hand_over_bundle(controls, STAGING_REPOSITORY, "uk-run-2") + hand_over_run_file(controls, STAGING_REPOSITORY, "staging-run-1") + # While the bundle cycle uploads to the private repository, the host's + # thread commits the staging repository's run files. + controls.hub.run_before_next_commit(controls.module.delivery.run_file_cycle) + + bundle_cycle(controls) + + staging_commits = controls.hub.commits(STAGING_REPOSITORY) + assert [list(commit.added) for commit in staging_commits] == [ + ["runs/staging-run-1/progress.json"] + ], "the staging bundle was committed less than 10 seconds after its run files" + assert len(controls.hub.commits(PRIVATE_REPOSITORY)) == 1 + controls.clock.advance(MIN_COMMIT_SPACING_SECONDS) + bundle_cycle(controls) + assert len(controls.hub.commits(STAGING_REPOSITORY)) == 2 + + +def test_run_files_wait_for_the_spacing_after_a_bundle_commit_during_their_cycle( + hugging_face_module, +): + from microcosm_emitter.huggingface.service.queue import ( + MIN_COMMIT_SPACING_SECONDS, + ) + + controls = hugging_face_module + hand_over_run_file(controls, PRIVATE_REPOSITORY, "private-run-1") + hand_over_run_file(controls, STAGING_REPOSITORY, "staging-run-1") + hand_over_bundle(controls, STAGING_REPOSITORY, "uk-run-2") + # While the run-file cycle commits to the private repository, the bundle + # worker commits the staging repository's bundle. + controls.hub.run_before_next_commit(lambda: bundle_cycle(controls)) + + controls.module.delivery.run_file_cycle() + + staging_commits = controls.hub.commits(STAGING_REPOSITORY) + assert [commit.message for commit in staging_commits] == [ + "Stage uk-run-2 under staged/uk-run-2" + ], "the staging run files were committed less than 10 seconds after the bundle" + controls.clock.advance(MIN_COMMIT_SPACING_SECONDS) + controls.module.delivery.run_file_cycle() + assert "runs/staging-run-1/progress.json" in ( + controls.hub.commits(STAGING_REPOSITORY)[-1].added + ) From 2f840bbb325f2f7a3723991944368ee6eff78136 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:20:56 +0400 Subject: [PATCH 19/23] Check the commit spacing again once a repository is leased The bundle worker and the host's thread each list due repositories and then deliver them in turn, so a list can be minutes old when a delivery starts. After the lease is taken, `_claim_job` now refuses to claim a job, and `run_files_due` reports no run files due, until `MIN_COMMIT_SPACING_SECONDS` have passed since this machine's last commit to the repository. The commit at shutdown still ignores the spacing. The helpers `_spacing_elapsed`, `commit_spacing_elapsed` and `_run_files_allowed` match the ones #24 uses, which also exempts release changes from the spacing. Green: 7 passed; every test file this branch adds or changes: 199 passed. Co-Authored-By: Claude Opus 5.5 --- .../huggingface/service/delivery.py | 10 ++++- .../huggingface/service/queue.py | 37 +++++++++++++++++-- 2 files changed, 42 insertions(+), 5 deletions(-) diff --git a/src/microcosm_emitter/huggingface/service/delivery.py b/src/microcosm_emitter/huggingface/service/delivery.py index 8e5f496..432a3af 100644 --- a/src/microcosm_emitter/huggingface/service/delivery.py +++ b/src/microcosm_emitter/huggingface/service/delivery.py @@ -277,9 +277,17 @@ def _run_files_to_commit( ) def _claim_job(self, repository: str, kinds: frozenset[str]) -> ClaimedJob | None: + """The next due job of ``kinds``, once the commit spacing has elapsed. + + Checked here, after the repository is leased: another thread may have + committed to it since this cycle listed it. + """ + now = self.clock() + if not self.queue.commit_spacing_elapsed(repository, now): + return None return self.queue.claim_job( repository, - now=self.clock(), + now=now, lease_seconds=REPOSITORY_LEASE_SECONDS, kinds=kinds, ) diff --git a/src/microcosm_emitter/huggingface/service/queue.py b/src/microcosm_emitter/huggingface/service/queue.py index f01927a..d556e2d 100644 --- a/src/microcosm_emitter/huggingface/service/queue.py +++ b/src/microcosm_emitter/huggingface/service/queue.py @@ -47,8 +47,10 @@ #: The shortest time between two commits to one repository from this machine, #: of run files or jobs, even for hand-overs that ask to be committed at the -#: next cycle. Only the drain at shutdown ignores it; other machines' commits -#: are not counted. +#: next cycle. A cycle checks it when it lists the repository and again once +#: it has leased it, because another thread may have committed in between. +#: Only the drain at shutdown ignores it; other machines' commits are not +#: counted. MIN_COMMIT_SPACING_SECONDS = 10.0 MAX_RETRY_SECONDS = 300.0 DESTINATION_HUGGING_FACE = "huggingface" @@ -249,6 +251,33 @@ def due_repositories( due.append(repository) return due + @staticmethod + def _spacing_elapsed(schedule: RepositoryCommitSchedule | None, now: float) -> bool: + """Whether this machine's last commit to the repository is far enough back. + + The spacing keeps bursts of commits, each starting the dataset + repository's webhook, from reaching the Hub. + """ + return ( + schedule is None + or schedule.last_commit_at + MIN_COMMIT_SPACING_SECONDS <= now + ) + + def commit_spacing_elapsed(self, repository: str, now: float) -> bool: + with self._sessions() as session: + return self._spacing_elapsed( + session.get(RepositoryCommitSchedule, repository), now + ) + + @classmethod + def _run_files_allowed( + cls, schedule: RepositoryCommitSchedule | None, now: float + ) -> bool: + """Run files wait out failure backoff and the spacing between commits.""" + return cls._spacing_elapsed(schedule, now) and ( + schedule is None or schedule.retry_after <= now + ) + # Staging run files --------------------------------------------------- @staticmethod @@ -351,10 +380,10 @@ def stage_run_files( delivery.updated_at = now def run_files_due(self, repository: str, now: float) -> bool: - """Whether pending run files are due and past any failure backoff.""" + """Whether pending run files are due, past backoff and commit spacing.""" with self._sessions() as session: schedule = session.get(RepositoryCommitSchedule, repository) - if schedule is not None and schedule.retry_after > now: + if not self._run_files_allowed(schedule, now): return False due_file = session.scalar( select(StagedRunFile.path_in_repo) From a72457f7f1075b1d376a0fd4506de1cd6f82e583 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:21:37 +0400 Subject: [PATCH 20/23] Say how to repair a spool database an earlier build moved to 20261020_03 A developer machine that ran an earlier, unreleased build of this branch has its spool database at revision `20261020_03`, which no migration history contains now, so every emitter there refuses to start. The README now names the file to delete (and its rollback journal), says what deleting it loses, and gives the revision reset that keeps the queued telemetry events and graph publication jobs. No repair code is added for an unreleased revision. Documentation only; no red step. Co-Authored-By: Claude Opus 5.5 --- README.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/README.md b/README.md index 0003d45..987ff13 100644 --- a/README.md +++ b/README.md @@ -60,6 +60,8 @@ The `huggingface` module takes over the uploads Microcosm made itself. Microcosm **Where hand-overs are stored.** The module keeps its queue in its own SQLite file, the hand-over database, with its own Alembic history (`core/handover_alembic`). `HuggingFaceSettings` names it: by default it sits next to the spool database (`events.sqlite3` → `events.handovers.sqlite3`), and `handover_database_path` chooses another file, never the spool database itself. The spool database stays at revision `20261008_02`, because every emitter on a machine shares it, including the released emitters (0.2.0 to 0.4.0) that Microcosm builds pin, and those refuse a revision they don't know. Each history also refuses the other's file. +**Repair for a machine that ran an earlier, unreleased build of the hand-over queue.** Before the hand-over database existed, an unreleased build kept the hand-over tables in the spool database and moved it to revision `20261020_03`. No migration history contains that revision now, so on such a machine every emitter, this one and the released ones, refuses the spool database at startup (`Can't locate revision identified by '20261020_03'`), and builds run without telemetry or hand-overs. Nothing repairs this automatically. To repair it, stop every Microcosm build and emitter on the machine, then delete the spool database: `$XDG_CACHE_HOME/microcosm/telemetry/events.sqlite3`, or `~/.cache/microcosm/telemetry/events.sqlite3` when `XDG_CACHE_HOME` is unset (or the `spool_path` the build set), together with `events.sqlite3-journal` beside it if that file exists. Deleting the file loses everything queued in it and not yet delivered: telemetry events not yet sent to the collector, graph publication jobs not yet published, and the run files and dataset bundles that the earlier build handed over and had not committed. The next emitter creates a new spool database at `20261008_02`. To keep the queued telemetry events and graph publication jobs instead, reset the file's revision with `sqlite3 "UPDATE alembic_version SET version_num = '20261008_02'"`: the earlier build added only the five hand-over tables, which every emitter then ignores, so their undelivered hand-overs are lost either way. + **Staging run files** (Microcosm's `staging.py` and `staging_v2.py` uploads). A build hands over its current run files; the module copies their bytes into the hand-over database before answering, so the build may rewrite them at once. - Files land at `//` (by default `runs//...`), exactly as the build's writer lays them out. Staging contract version 1 also writes `latest_staging.json` at the repository root, and its `runs.json` index: at commit time the module merges the run's summary into the repository's copy at the head it commits on, keeping other runs, with Microcosm's ordering and JSON formatting. It starts a new index only when the Hub answers that the repository has no `runs.json`; any other download error (a connection error, a timeout, a 429 or 5xx response) fails the commit and retries it after the backoff. Microcosm started from an empty index on any download error, which could drop every other run. From a7e0ea5ba3c987108566c971b9fe68504e68e168 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 23:10:35 +0400 Subject: [PATCH 21/23] Test that a job claimed as shutdown begins starts no upload The round-2 review of #23 found one window left after the stop checks: the bundle cycle checks for shutdown before each repository and once the repository is leased, but not after it claims the job. A shutdown that begins between the post-lease check and the claim still starts that upload. The new test begins shutdown from inside the job claim and expects no upload, and the job left pending with its lease released and its attempt count unchanged. Red: 1 failed, 7 passed. The failure is the stated reason: `AssertionError: a bundle upload started after shutdown began, once its job was claimed`. Co-Authored-By: Claude Opus 5.5 --- tests/huggingface/test_delivery_workers.py | 48 ++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/tests/huggingface/test_delivery_workers.py b/tests/huggingface/test_delivery_workers.py index 1781f8b..19b1f53 100644 --- a/tests/huggingface/test_delivery_workers.py +++ b/tests/huggingface/test_delivery_workers.py @@ -357,6 +357,54 @@ def create_commit_holding_the_second_repository(repo_id, operations, **options): test_finished.set() +def test_a_job_claimed_as_shutdown_begins_is_released_without_an_upload( + hugging_face_module, monkeypatch +): + """Shutdown between a job's claim and its upload: no upload starts. + + The bundle cycle checks for shutdown before each repository and once the + repository is leased. Here shutdown begins after both checks, while the + job is being claimed. The cycle must start no upload, and must leave the + job pending, with its lease released and its attempt count unchanged, so + the next emitter can deliver it at once. + """ + from microcosm_emitter.huggingface.service.delivery import ( + DATASET_BUNDLE, + REPOSITORY_LEASE_SECONDS, + ) + + controls = hugging_face_module + hand_over_bundle(controls, PRIVATE_REPOSITORY, "uk-run-1") + queue = controls.module.queue + shutdown_begun = threading.Event() + claim_job_in_the_queue = queue.claim_job + + def claim_job_as_shutdown_begins(*arguments, **options): + claimed = claim_job_in_the_queue(*arguments, **options) + shutdown_begun.set() + return claimed + + monkeypatch.setattr(queue, "claim_job", claim_job_as_shutdown_begins) + + controls.module.delivery.job_cycle( + frozenset({DATASET_BUNDLE}), stop_requested=shutdown_begun.is_set + ) + + assert shutdown_begun.is_set(), "the cycle never claimed the job" + assert controls.hub.commits(PRIVATE_REPOSITORY) == [], ( + "a bundle upload started after shutdown began, once its job was claimed" + ) + monkeypatch.undo() + reclaimed = queue.claim_job( + PRIVATE_REPOSITORY, + now=controls.clock(), + lease_seconds=REPOSITORY_LEASE_SECONDS, + kinds=[DATASET_BUNDLE], + ) + assert reclaimed is not None, "the job's lease was kept after shutdown began" + assert reclaimed.attempts == 0 + + # The 10-second spacing between commits to one repository, across threads. # # The host's thread (run files) and the bundle worker each build a list of due From b8215238a804d454101680679f061eddbf18eea4 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 23:11:45 +0400 Subject: [PATCH 22/23] Give back a job claimed after shutdown began instead of uploading it `_deliver_repository` now checks for shutdown once more after it claims a job, immediately before the upload begins. A job claimed after shutdown began is given back with the new `HandoverQueue.release_job`, which clears its lease without counting an attempt, so the job stays pending and due for the next emitter. The README, the `workers` module docstring and the delivery docstrings now state the exact guarantee: the worker checks before each repository, once the repository is leased, and once the job is claimed. During the commit at shutdown a repository can be held only by an upload that passed that last check before shutdown began, or, for the few database operations it takes to give it back, by a lease taken just before. Green: 8 passed. Co-Authored-By: Claude Opus 5.5 --- README.md | 2 +- .../huggingface/service/delivery.py | 16 ++++++++++++---- .../huggingface/service/queue.py | 17 +++++++++++++++++ .../huggingface/service/workers.py | 9 ++++++--- 4 files changed, 36 insertions(+), 8 deletions(-) diff --git a/README.md b/README.md index 987ff13..a4d8ac7 100644 --- a/README.md +++ b/README.md @@ -79,7 +79,7 @@ The `huggingface` module takes over the uploads Microcosm made itself. Microcosm - Unlike Microcosm's single attempt, `REPOSITORY_UNAVAILABLE` and `UPLOAD_FAILED` are retried with backoff, and a commit that loses a race to another writer is checked and tried again. The receipt keeps the last error code while the job is pending. - `check_write_access` is Microcosm's staged-dataset preflight: whether a credential is present, the repository is reachable, and the Hub reports a write scope for it. -**Bundle uploads run on their own thread.** A bundle upload can take minutes, so the module uploads bundles on a worker thread of its own, while the host's worker thread commits run files and makes the commit at shutdown. A bundle upload to one repository therefore never delays run-file commits to another repository, the commit at shutdown, or a read-back. Once shutdown begins, the bundle worker starts no new upload, even to a repository left in its current pass, so only an upload already under way can hold a repository during the commit at shutdown. Each delivery leases its repository first, so this machine never commits to one repository twice at once: run files for the repository a bundle is uploading to wait for that upload, and then for the 10-second spacing. +**Bundle uploads run on their own thread.** A bundle upload can take minutes, so the module uploads bundles on a worker thread of its own, while the host's worker thread commits run files and makes the commit at shutdown. A bundle upload to one repository therefore never delays run-file commits to another repository, the commit at shutdown, or a read-back. Once shutdown begins, the bundle worker starts no further upload, even to a repository left in its current pass: it checks for shutdown before each repository, again once it has leased the repository, and again once it has claimed a job, immediately before that upload begins, and it gives back a job claimed after shutdown began without attempting it. So during the commit at shutdown a repository can be held only by an upload that passed that last check before shutdown began, or, for the few database operations it takes to give it back, by a lease the worker took just before; the second commit at shutdown covers that repository if the worker stops within the shutdown period. Each delivery leases its repository first, so this machine never commits to one repository twice at once: run files for the repository a bundle is uploading to wait for that upload, and then for the 10-second spacing. ## Development diff --git a/src/microcosm_emitter/huggingface/service/delivery.py b/src/microcosm_emitter/huggingface/service/delivery.py index 432a3af..7cc6c0f 100644 --- a/src/microcosm_emitter/huggingface/service/delivery.py +++ b/src/microcosm_emitter/huggingface/service/delivery.py @@ -167,8 +167,10 @@ def job_cycle( ) -> int: """Deliver one due job of ``kinds`` per repository; return how many. - Once ``stop_requested`` turns true (shutdown has begun), the cycle - starts no new upload and returns. + ``stop_requested`` is checked before each repository, once its lease + is taken, and once its job is claimed, immediately before the upload + begins. Once it is true (shutdown has begun), the cycle starts no + upload and returns; an upload that began before that continues. """ return self._cycle( job_kinds=kinds & SUPPORTED_JOB_KINDS, @@ -230,8 +232,9 @@ def _deliver_repository( """One commit to ``repository``: a job of ``job_kinds``, else run files. ``at_shutdown`` commits pending run files whether or not they are due. - ``stop_requested`` is checked again once the repository is leased, so - a shutdown that began while the lease was being taken starts no upload. + ``stop_requested`` is checked again once the repository is leased, and + once more after a job is claimed, immediately before its upload + begins; a job claimed after shutdown began is released unattempted. """ lease = self.queue.claim_repository( repository, now=self.clock(), lease_seconds=REPOSITORY_LEASE_SECONDS @@ -260,6 +263,11 @@ def _deliver_repository( hub = DatasetRepository(repository, api) job = self._claim_job(repository, job_kinds) if job_kinds else None if job is not None: + if stop_requested(): + # Shutdown began while the job was being claimed: give it + # back untouched instead of starting its upload. + self.queue.release_job(job) + return 0 held.hold_job(job) return int(self._deliver_job(hub, held)) if not self._run_files_to_commit(repository, run_files, at_shutdown): diff --git a/src/microcosm_emitter/huggingface/service/queue.py b/src/microcosm_emitter/huggingface/service/queue.py index d556e2d..edd883d 100644 --- a/src/microcosm_emitter/huggingface/service/queue.py +++ b/src/microcosm_emitter/huggingface/service/queue.py @@ -712,6 +712,23 @@ def renew_job(self, job: ClaimedJob, renewed: float) -> bool: is not None ) + def release_job(self, job: ClaimedJob) -> None: + """Give back ``job``'s lease without recording an attempt. + + The job stays pending and due, so any emitter may claim it at once. + """ + with self._sessions.begin() as session: + session.execute( + update(HandoverJob) + .where( + HandoverJob.handover_id == job.handover_id, + HandoverJob.destination == job.destination, + HandoverJob.status == JOB_PENDING, + HandoverJob.lease_until == job.lease_until, + ) + .values(lease_until=0.0) + ) + def finish_job( self, job: ClaimedJob, diff --git a/src/microcosm_emitter/huggingface/service/workers.py b/src/microcosm_emitter/huggingface/service/workers.py index a5e42a6..0893ce5 100644 --- a/src/microcosm_emitter/huggingface/service/workers.py +++ b/src/microcosm_emitter/huggingface/service/workers.py @@ -13,9 +13,12 @@ repository at once: each delivery first leases its repository in the hand-over database, and one that finds it leased tries again next cycle. -Once shutdown begins (``stop``), a worker starts no new upload: its cycle -receives ``stop_requested`` and checks it before each repository, so only -an upload already under way when ``stop`` was called can still be running. +Once shutdown begins (``stop``), a worker starts no further upload: its +cycle receives ``stop_requested`` and checks it before each repository, once +that repository is leased, and once a job is claimed, immediately before the +upload begins; a job claimed after ``stop`` is given back unattempted. So +only an upload that passed that last check before ``stop`` was called can +still be running. """ from __future__ import annotations From 8fa62ebaaa1ad256b89a45ee28ca004f686997cc Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 23:14:29 +0400 Subject: [PATCH 23/23] Guard the spool repair to revision 20261020_03 and drop its tables The round-2 review of #23 asked for three changes to the README's repair note for a spool database an earlier, unreleased build moved to `20261020_03`: - The revision reset now runs in one transaction that first inserts the file's revision into a temporary table whose check accepts only `20261020_03`. With `sqlite3 -bail`, any other revision (a later spool revision, or the hand-over database, whose tables have the same names) stops the command with nothing changed. The `UPDATE` also carries `WHERE version_num = '20261020_03'`. - The same transaction drops the earlier build's tables, as that migration's own `downgrade()` did. Checked against the migration file at every earlier head (`f438369`, `473b576`, `04a612a`): it only created `staged_run_files`, `staged_run_deliveries`, `repository_commit_schedules`, `handover_jobs` and, from `04a612a`'s form on, `handover_files`, with no index. No spool revision at any branch of the stack or any release uses these names. - The note says four or five tables, and that an emitter started with only the `huggingface` module never opens the spool database. Documentation only; no red step. Co-Authored-By: Claude Opus 5.5 --- README.md | 23 ++++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index a4d8ac7..a1342d7 100644 --- a/README.md +++ b/README.md @@ -60,7 +60,28 @@ The `huggingface` module takes over the uploads Microcosm made itself. Microcosm **Where hand-overs are stored.** The module keeps its queue in its own SQLite file, the hand-over database, with its own Alembic history (`core/handover_alembic`). `HuggingFaceSettings` names it: by default it sits next to the spool database (`events.sqlite3` → `events.handovers.sqlite3`), and `handover_database_path` chooses another file, never the spool database itself. The spool database stays at revision `20261008_02`, because every emitter on a machine shares it, including the released emitters (0.2.0 to 0.4.0) that Microcosm builds pin, and those refuse a revision they don't know. Each history also refuses the other's file. -**Repair for a machine that ran an earlier, unreleased build of the hand-over queue.** Before the hand-over database existed, an unreleased build kept the hand-over tables in the spool database and moved it to revision `20261020_03`. No migration history contains that revision now, so on such a machine every emitter, this one and the released ones, refuses the spool database at startup (`Can't locate revision identified by '20261020_03'`), and builds run without telemetry or hand-overs. Nothing repairs this automatically. To repair it, stop every Microcosm build and emitter on the machine, then delete the spool database: `$XDG_CACHE_HOME/microcosm/telemetry/events.sqlite3`, or `~/.cache/microcosm/telemetry/events.sqlite3` when `XDG_CACHE_HOME` is unset (or the `spool_path` the build set), together with `events.sqlite3-journal` beside it if that file exists. Deleting the file loses everything queued in it and not yet delivered: telemetry events not yet sent to the collector, graph publication jobs not yet published, and the run files and dataset bundles that the earlier build handed over and had not committed. The next emitter creates a new spool database at `20261008_02`. To keep the queued telemetry events and graph publication jobs instead, reset the file's revision with `sqlite3 "UPDATE alembic_version SET version_num = '20261008_02'"`: the earlier build added only the five hand-over tables, which every emitter then ignores, so their undelivered hand-overs are lost either way. +**Repair for a machine that ran an earlier, unreleased build of the hand-over queue.** Before the hand-over database existed, an unreleased build kept four or five hand-over tables in the spool database (`staged_run_files`, `staged_run_deliveries`, `repository_commit_schedules`, `handover_jobs`, and in its later form `handover_files`) and moved the spool database to revision `20261020_03`. No migration history contains that revision now. So on such a machine an emitter started with the `telemetry` or `orrery` module, this one or a released one, refuses the spool database at startup (`Can't locate revision identified by '20261020_03'`), and the build runs without that emitter's telemetry or hand-overs. An emitter started with only the `huggingface` module never opens the spool database and is not affected. Nothing repairs this automatically. To repair it, first stop every Microcosm build and emitter on the machine. The spool database is `$XDG_CACHE_HOME/microcosm/telemetry/events.sqlite3`, or `~/.cache/microcosm/telemetry/events.sqlite3` when `XDG_CACHE_HOME` is unset (or the `spool_path` the build set). Then do one of these: + +- **Delete it,** together with `events.sqlite3-journal` beside it if that file exists. This loses everything queued in it and not yet delivered: telemetry events not yet sent to the collector, graph publication jobs not yet published, and the run files and dataset bundles that the earlier build handed over and had not committed. The next emitter creates a new spool database at `20261008_02`. +- **Keep the queued telemetry events and graph publication jobs** by undoing what the earlier build's migration did, as its own downgrade would have: drop its tables and set the revision back to `20261008_02`. + + ```sh + sqlite3 -bail <<'SQL' + BEGIN; + -- Stop here, changing nothing, unless the file is at the earlier build's revision. + CREATE TEMP TABLE repair_guard (revision TEXT NOT NULL CHECK (revision = '20261020_03')); + INSERT INTO repair_guard VALUES ((SELECT version_num FROM alembic_version)); + DROP TABLE IF EXISTS handover_files; + DROP TABLE IF EXISTS handover_jobs; + DROP TABLE IF EXISTS repository_commit_schedules; + DROP TABLE IF EXISTS staged_run_deliveries; + DROP TABLE IF EXISTS staged_run_files; + UPDATE alembic_version SET version_num = '20261008_02' WHERE version_num = '20261020_03'; + COMMIT; + SQL + ``` + + Unless the file's revision is `20261020_03`, the guard row is refused, `-bail` stops `sqlite3` at that error, and the open transaction is rolled back. So the command changes nothing in a spool database at any other revision, nor in the hand-over database (`events.handovers.sqlite3`), whose tables have the same names. No spool revision uses these table names. The dropped tables hold only the earlier build's undelivered hand-overs, which are lost either way. **Staging run files** (Microcosm's `staging.py` and `staging_v2.py` uploads). A build hands over its current run files; the module copies their bytes into the hand-over database before answering, so the build may rewrite them at once.