diff --git a/README.md b/README.md index cf20166..a1342d7 100644 --- a/README.md +++ b/README.md @@ -54,6 +54,54 @@ 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. + +**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 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. + +- 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 (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`. + +**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 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. +- `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 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 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..1db66a8 --- /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`), 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/changelog.d/huggingface-handover-queue.added.md b/changelog.d/huggingface-handover-queue.added.md new file mode 100644 index 0000000..c6fe36e --- /dev/null +++ b/changelog.d/huggingface-handover-queue.added.md @@ -0,0 +1 @@ +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/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/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 c02750d..4583df3 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, @@ -11,18 +12,32 @@ ) 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, + StageRunFiles, +) from microcosm_emitter_interface.telemetry import TelemetryRun, TelemetrySettings __all__ = [ + "BundleFile", "COMMAND_ENVIRONMENT_VARIABLE", "DEFAULT_COMMAND", "PROTOCOL_VERSION", "EmitterClient", "EmitterStartupError", + "HuggingFaceSettings", "OrrerySettings", "RequestFailure", + "RunFile", + "StageDatasetBundle", + "StageRunFiles", "TelemetryRun", "TelemetrySettings", + "default_handover_database_path", "default_spool_path", ] 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..66e7cb9 --- /dev/null +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/bundles.py @@ -0,0 +1,236 @@ +"""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``. + +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``. +""" + +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, NotRequired, 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 + #: The files the emitter copied at hand-over; the build may rewrite them. + copied_files: NotRequired[list[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: + """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": handover_id, + "status": "unavailable", + "error_code": None, + "revision": None, + "attempts": 0, + "repository": request.repository, + "prefix": request.remote_prefix, + "copied_files": [], + } + + +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 1ba3f3e..04d6d56 100644 --- a/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/client.py @@ -17,8 +17,16 @@ 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.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, @@ -41,11 +49,25 @@ from microcosm_emitter_interface.protocol import ( ACKNOWLEDGEMENT_TIMEOUT_SECONDS, 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, ) 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 +131,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 +144,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 +348,144 @@ 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. + """ + response = self._query( + HUGGING_FACE_MODULE, + request.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: + response = self._query( + HUGGING_FACE_MODULE, + lambda: run_reference_message("run_files_receipt", repository, run_id), + 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: + self._fail("staging", "emitter_unavailable") + return unavailable_read_back() + 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"] + and time.monotonic() < deadline + ): + time.sleep(0.1) + receipt = self.run_files_receipt(repository, run_id) + response = self._query( + HUGGING_FACE_MODULE, + lambda: run_reference_message("read_back_run_files", repository, run_id), + 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 _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]: + 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 + and receipt.get("status") != "unavailable" + and time.monotonic() < deadline + ): + time.sleep(0.1) + latest = self._handover_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``. + """ + response = self._query( + HUGGING_FACE_MODULE, + request.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``.""" + response = self._query( + HUGGING_FACE_MODULE, + lambda: { + "protocol_version": PROTOCOL_VERSION, + "command": "check_write_access", + "repository": repository_identifier(repository), + }, + 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-configuration.json b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/fixtures/protocol-v1/huggingface-configuration.json new file mode 100644 index 0000000..4a9948e --- /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": { + "handover_database_path": "{directory}/events.handovers.sqlite3" + } +} 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-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-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/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..2c8022e 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,21 @@ 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 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/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 new file mode 100644 index 0000000..0ee0e34 --- /dev/null +++ b/packages/microcosm-emitter-interface/src/microcosm_emitter_interface/staging.py @@ -0,0 +1,219 @@ +"""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 +from microcosm_emitter_interface.spool import default_handover_database_path + +DEFAULT_RUN_PREFIX = "runs" +DEFAULT_COMMIT_INTERVAL_SECONDS = 30.0 + + +@dataclass(frozen=True) +class HuggingFaceSettings: + """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]: + 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) +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 + 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 + 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. + + ``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 + 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 eddfeab..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" @@ -43,6 +44,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/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/handover_alembic/versions/v20261020_01_handover_queue.py b/src/microcosm_emitter/core/handover_alembic/versions/v20261020_01_handover_queue.py new file mode 100644 index 0000000..d69da1b --- /dev/null +++ b/src/microcosm_emitter/core/handover_alembic/versions/v20261020_01_handover_queue.py @@ -0,0 +1,80 @@ +"""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 = "handovers_20261020_01" +down_revision = None +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), + ) + 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") + op.drop_table("staged_run_files") 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/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..04c8be6 --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/__init__.py @@ -0,0 +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/bundles.py b/src/microcosm_emitter/huggingface/service/bundles.py new file mode 100644 index 0000000..1242259 --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/bundles.py @@ -0,0 +1,147 @@ +"""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], + copied: Mapping[str, bytes] | None = None, +) -> BundleOutcome: + """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( + [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']}" + 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']}", + copied.get(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 new file mode 100644 index 0000000..7cc6c0f --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/delivery.py @@ -0,0 +1,415 @@ +"""Delivery cycles: at most one Hugging Face commit per repository per cycle. + +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 + +import sys +import threading +import time +from collections.abc import Callable +from dataclasses import replace +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 ClaimedJob, HandoverQueue +from microcosm_emitter.huggingface.service.run_files import commit_run_files + +#: 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}) +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 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() + + +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.""" + + def __init__( + self, + queue: HandoverQueue, + *, + 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() + + 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: + """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], + *, + stop_requested: Callable[[], bool] = never_stop, + ) -> int: + """Deliver one due job of ``kinds`` per repository; return how many. + + ``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, + 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() + 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, + ) + return delivered + + 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: 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, job_kinds=frozenset(), run_files=True, at_shutdown=True + ) + + def _deliver_repository( + self, + repository: str, + api: Any | None, + *, + 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, 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 + ) + 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 stop_requested(): + return 0 + if api is None: + self._credential_missing( + held, + job_kinds=job_kinds, + run_files=run_files, + at_shutdown=at_shutdown, + ) + return 0 + 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): + return 0 + return int(self._deliver_run_files(hub)) + finally: + held.stop() + self.queue.release_repository(repository, held.repository_lease) + + 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: + """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=now, + lease_seconds=REPOSITORY_LEASE_SECONDS, + kinds=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']}", + "copied_files": sorted(job.files), + }, + retry=retry, + now=self.clock(), + committed=committed, + ) + + 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( + held.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, + 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 = self._claim_job(repository, job_kinds) if job_kinds else None + if job is not None: + held.hold_job(job) + held.stop() + assert held.job is not None + self._finish_job( + held.job, + status="pending", + error_code=CREDENTIAL_MISSING, + revision=None, + retry=True, + ) + return + 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) + 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/hub.py b/src/microcosm_emitter/huggingface/service/hub.py new file mode 100644 index 0000000..6db29e6 --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/hub.py @@ -0,0 +1,220 @@ +"""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}) +#: 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: + """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 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: + """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/module.py b/src/microcosm_emitter/huggingface/service/module.py new file mode 100644 index 0000000..48c6fb8 --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/module.py @@ -0,0 +1,335 @@ +"""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. 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 + +import time +from collections.abc import Callable, Mapping +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, + DATASET_BUNDLE, + 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_dataset_bundle, + parse_stage_run_files, + 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. +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." +) + + +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 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, + configuration: JsonObject, + context: ModuleContext, + *, + 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: + # 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 + ) + 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 stop_requested: delivery.job_cycle( + frozenset({DATASET_BUNDLE}), stop_requested=stop_requested + ), + 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.") + + 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)} + 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": "", + "copied_files": [], + } + 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, + "copied_files": sorted(hand_over.copied), + }, + now=self.clock(), + files=hand_over.copied, + ) + + 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: + 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"} + 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} + 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: + 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._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. 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 + self._closed = True + for worker in self.workers.values(): + worker.stop() + try: + 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: + self._worker_error() + workers_stopped = False + finally: + # 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 new file mode 100644 index 0000000..edd883d --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/queue.py @@ -0,0 +1,789 @@ +"""Durable hand-overs, kept in the Hugging Face module's hand-over database. + +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 + 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 json +import time +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +from sqlalchemy import delete, select, update +from sqlalchemy.orm import Session + +from microcosm_emitter.core.database import ( + create_spool_engine, + create_spool_session_factory, +) +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 +#: 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" +JOB_PENDING = "pending" + + +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: + """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) +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 + #: The bytes of the files copied at hand-over, by name. + files: dict[str, bytes] = field(default_factory=dict) + + +class HandoverQueue: + """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 + (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) + self.path.parent.mkdir(parents=True, exist_ok=True) + self._engine = create_spool_engine(self.path) + upgrade_handover_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 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( + 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, + *, + 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; 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[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, + ) + ) + ) + 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) + 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 + + @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 + 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. 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 + 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 = _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) + 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 run_files_due(self, repository: str, now: float) -> bool: + """Whether pending run files are due, past backoff and commit spacing.""" + with self._sessions() as session: + schedule = session.get(RepositoryCommitSchedule, repository) + if not self._run_files_allowed(schedule, 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( + 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( + 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, _entry(row.run_index_entry)) 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"" + 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): + 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 == _entry_text(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, + files: Mapping[str, bytes] | None = None, + ) -> dict[str, Any]: + """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: + 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, + ) + ) + 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, + *, + 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``. + + 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: + 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( + statement.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) + copied = session.scalars( + select(HandoverFile).where( + HandoverFile.handover_id == candidate.handover_id + ) + ) + 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, + 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 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, + *, + 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"]) + 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 + + 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/src/microcosm_emitter/huggingface/service/requests.py b/src/microcosm_emitter/huggingface/service/requests.py new file mode 100644 index 0000000..8799ccf --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/requests.py @@ -0,0 +1,242 @@ +"""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 hashlib +import math +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path +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, + 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" +#: 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): + """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), + ) + + +@dataclass(frozen=True) +class BundleHandOver: + handover_id: str + 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. + + 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") + 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 + 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/src/microcosm_emitter/huggingface/service/run_files.py b/src/microcosm_emitter/huggingface/service/run_files.py new file mode 100644 index 0000000..3204a97 --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/run_files.py @@ -0,0 +1,157 @@ +"""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: + # 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 + 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/src/microcosm_emitter/huggingface/service/workers.py b/src/microcosm_emitter/huggingface/service/workers.py new file mode 100644 index 0000000..0893ce5 --- /dev/null +++ b/src/microcosm_emitter/huggingface/service/workers.py @@ -0,0 +1,96 @@ +"""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. + +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 + +import threading +import time +from collections.abc import Callable + + +class DeliveryWorker: + """Runs ``cycle`` on its own thread every ``interval_seconds`` until stopped. + + ``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[[Callable[[], bool]], 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 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``. + + ``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(self.stop_requested) + except Exception: + self.on_error() diff --git a/tests/huggingface/conftest.py b/tests/huggingface/conftest.py new file mode 100644 index 0000000..00c0345 --- /dev/null +++ b/tests/huggingface/conftest.py @@ -0,0 +1,63 @@ +"""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( + {"handover_database_path": str(tmp_path / "events.handovers.sqlite3")}, + 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( + 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_bundles.py b/tests/huggingface/test_bundles.py new file mode 100644 index 0000000..ff04080 --- /dev/null +++ b/tests/huggingface/test_bundles.py @@ -0,0 +1,519 @@ +"""Staged dataset bundles: Microcosm's paths, idempotency and error codes.""" + +import hashlib +import json +from pathlib import Path + +import pytest +from fake_hugging_face import ( + PRIVATE_REPOSITORY, + FakeHubHTTPError, + FakeHuggingFaceApi, +) + +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 + +DATASET = b"synthetic h5 bytes" +DIAGNOSTICS = b'{"schema_version":8}' + + +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), + "microcosm_uk_2024_25.diagnostics.json": ( + "calibration_diagnostics", + 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() + } + 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 + + +@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, nothing_copied +): + 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, nothing_copied +): + 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" + + +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) + 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( + {"handover_database_path": str(tmp_path / "events.handovers.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/huggingface/test_delivery_workers.py b/tests/huggingface/test_delivery_workers.py new file mode 100644 index 0000000..19b1f53 --- /dev/null +++ b/tests/huggingface/test_delivery_workers.py @@ -0,0 +1,497 @@ +"""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 + ) + + +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() + + +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 +# 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 + ) diff --git a/tests/huggingface/test_hub.py b/tests/huggingface/test_hub.py new file mode 100644 index 0000000..c60ba71 --- /dev/null +++ b/tests/huggingface/test_hub.py @@ -0,0 +1,219 @@ +"""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) + # 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"}) + 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_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) diff --git a/tests/huggingface/test_queue.py b/tests/huggingface/test_queue.py new file mode 100644 index 0000000..2568b49 --- /dev/null +++ b/tests/huggingface/test_queue.py @@ -0,0 +1,369 @@ +"""The durable hand-over queue, in its own database next to the spool database.""" + +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.handovers.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_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(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 ( + session.scalar(select(func.count(TelemetryEventRecord.event_id))) + == events_before + ) + 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) + (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.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 + 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_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, + 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.handovers.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/huggingface/test_run_files.py b/tests/huggingface/test_run_files.py new file mode 100644 index 0000000..6b540b0 --- /dev/null +++ b/tests/huggingface/test_run_files.py @@ -0,0 +1,494 @@ +"""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] + + +@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"} + 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, + } + + +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"), + [ + ("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/conftest.py b/tests/interface/conftest.py index 8d1f57a..c067537 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( + {"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, + ) + 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_errors.py b/tests/interface/test_client_errors.py index d0ec57e..e8de8cb 100644 --- a/tests/interface/test_client_errors.py +++ b/tests/interface/test_client_errors.py @@ -252,3 +252,111 @@ 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" + + +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"} diff --git a/tests/interface/test_client_messages.py b/tests/interface/test_client_messages.py index 71680e3..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 @@ -14,8 +15,36 @@ 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" +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: """Records requests and answers queries like an emitter that queued nothing.""" @@ -28,6 +57,27 @@ 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, + "handover_id": "fixture-handover", + }, + "access": { + "status": "checked", + "credential": "present", + "reachable": True, + "can_write": True, + "error_type": None, + }, + "read_back": { + "status": "passed", + "files": 0, + "mismatched_paths": [], + "error_code": None, + }, + } if message.get("command") == "inventory": return {"inventory": None} return { @@ -52,7 +102,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 +183,46 @@ 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-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 + ) + ), } CONFIGURED_BY: dict[str, Callable[[Path], dict[str, Any]]] = { @@ -142,6 +235,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 +265,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..10aff54 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,30 @@ "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": ( + "0ecea1a8e7ba6c4458c530b58fa90079207ee3f48da3e522479e9afde581551a" + ), + "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 +201,38 @@ 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(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()), + 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 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()) + + 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..ef205eb --- /dev/null +++ b/tests/interface/test_staging_client.py @@ -0,0 +1,175 @@ +"""Typed staging-run-file hand-overs, through a real socket to the module.""" + +from pathlib import Path + +import pytest +from fake_hugging_face import STAGING_REPOSITORY + +from microcosm_emitter_interface.client import EmitterClient +from microcosm_emitter_interface.staging import ( + HuggingFaceSettings, + RunFile, + StageRunFiles, +) +from microcosm_emitter_interface.transport import SocketTransport + + +@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_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( + 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, in_thread_emitter +): + client = EmitterClient( + in_thread_emitter( + 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, in_thread_emitter): + client = EmitterClient( + in_thread_emitter( + 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 new file mode 100644 index 0000000..a80ad07 --- /dev/null +++ b/tests/support/fake_hugging_face.py @@ -0,0 +1,287 @@ +"""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 + +STAGING_REPOSITORY = "policyengine/fixture-staging" +PRIVATE_REPOSITORY = "policyengine/fixture-private" + + +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/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/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() diff --git a/tools/installed-interface-smoke.py b/tools/installed-interface-smoke.py index 0e2b9ad..348cd11 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,20 @@ 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 + # 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()