From 0cda403a3c53ac4c30cfb2b4287d76a18809d23f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Thu, 17 Sep 2026 18:37:26 +0100 Subject: [PATCH 1/7] Stage UK rowwise candidate builds: telemetry to runs/ and the dataset bundle to staged/ Wire the shared staging telemetry (#896) into tools/build_uk_rowwise_candidate.py the way the national and spine commands carry it: the shared --staging-* options, a StagingTelemetryV2 run opened with the attempt (run_kind calibration, non-release), stage events at every phase boundary, per-epoch losses through a new progress_events fan-out on solve_uk_rowwise_weights_under_doctrine (thinned in the driver because a 2,000-epoch size run emits ~24,000 epochs against the contract's 5 MiB cap), staging_delivery evidence in the manifest, and a failed run marked on any exception. Add the staged-dataset lane: the finished bundle every manifest-registered output, the manifest, staged_manifest.json and sha256sums.txt is verified from disk and written in one commit to staged// of the private UK dataset repository, keyed by the telemetry run id, with staged_dataset evidence in the manifest and reviewed staged_dataset/fit_summary artifacts on the telemetry run. Staging is not publication: releases/ and latest.json are untouched and no release contract is consulted, so a releasable:false size run stages like any other. A remote stage is refused up front without a Hub credential that can see the repository; an upload failure is recorded and never changes the exit code. tools/stage_uk_rowwise_candidate.py re-stages a finished directory (idempotent on the outputs' digests); tools/fetch_uk_staged_dataset.py fetches a bundle by run id with digest checks. The country-neutral pieces live in microcosm.build.staging_dataset and on HuggingFaceDatasetStorage (commit, head_revision, file_exists, download_file); the UK repository defaults sit beside UK_STAGING_REPOSITORY. The dense release assembler adopts the national rule and requires the run's staging_delivery receipt, copying it into build_manifest.json as staging. Co-Authored-By: Claude Fable 5.1 --- CLAUDE.md | 5 +- README.md | 12 + ...ense-assembler-staging-evidence.changed.md | 1 + changelog.d/uk-rowwise-staging.added.md | 1 + changelog.d/uk-staged-dataset-lane.added.md | 1 + docs/uk-dataset-size-plan-355.md | 15 + docs/uk-dense-release-assembly-runbook-762.md | 7 +- docs/uk-staging-operations.md | 72 +- experiments/355-uk-dataset-size-receipts.md | 45 ++ .../src/microcosm/build/staging_cli.py | 58 +- .../src/microcosm/build/staging_dataset.py | 654 ++++++++++++++++++ .../src/microcosm/build/staging_storage.py | 72 +- .../build/uk_runtime/local_rowwise.py | 42 +- .../src/microcosm/build/uk_runtime/staging.py | 23 +- .../tests/test_staging_dataset.py | 377 ++++++++++ .../microcosm-build/tests/test_staging_v2.py | 114 ++- .../tests/test_uk_dense_release_assembler.py | 59 ++ .../tests/test_uk_local_rowwise.py | 42 ++ .../tests/test_uk_rowwise_candidate.py | 544 +++++++++++++++ .../tests/test_uk_size_evaluation.py | 51 ++ tools/assemble_uk_dense_release_dir.py | 24 + tools/build_uk_rowwise_candidate.py | 632 ++++++++++++++++- tools/fetch_uk_staged_dataset.py | 84 +++ tools/stage_uk_rowwise_candidate.py | 161 +++++ 24 files changed, 3062 insertions(+), 34 deletions(-) create mode 100644 changelog.d/uk-dense-assembler-staging-evidence.changed.md create mode 100644 changelog.d/uk-rowwise-staging.added.md create mode 100644 changelog.d/uk-staged-dataset-lane.added.md create mode 100644 packages/microcosm-build/src/microcosm/build/staging_dataset.py create mode 100644 packages/microcosm-build/tests/test_staging_dataset.py create mode 100644 tools/fetch_uk_staged_dataset.py create mode 100644 tools/stage_uk_rowwise_candidate.py diff --git a/CLAUDE.md b/CLAUDE.md index 41490df9ed..6dd458959f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -100,7 +100,10 @@ see README "Releasing & alerts". Publication also refuses a release whose build recorded staging telemetry that never reached its repo (`--allow-missing-staging` overrides); a build that declared `--no-staging` publishes without the flag. Never publish or promote artifacts as a side -effect of another task. +effect of another task. A UK rowwise run's **staged** bundle +(`staged//` in the private repository, written by the build itself) +is inspection evidence, not a release: it never moves `releases/` or +`latest.json` and is not loadable through the certified loader. The US native-SPM-role source-enrichment lane is a separate release type: `tools/build_us_spm_role_enrichment.py` creates a local candidate from the exact diff --git a/README.md b/README.md index 12c9c1bde6..6fc241c316 100644 --- a/README.md +++ b/README.md @@ -63,6 +63,18 @@ This writes `progress.json`, `events.ndjson`, `calibration_progress.json`, and final candidate diagnostics under `runs//` without updating production `latest.json`. +The UK commands (`tools/build_uk_frs_spine.py`, +`tools/calibrate_uk_national_dataset.py`, `tools/build_uk_rowwise_candidate.py`) +stage version 2 telemetry to `policyengine/populace-uk-staging` under the same +switch. The rowwise candidate command also **stages the finished dataset +bundle** it built, dense or exact-count, under `staged//` in the +private `policyengine/populace-uk-private` repository so the team can inspect +it without publishing it: `releases/` and `latest.json` are untouched, the +release contract is not consulted, and a `releasable: false` size run stages +like any other. Fetch a bundle with `tools/fetch_uk_staged_dataset.py`; +re-stage a finished run directory with `tools/stage_uk_rowwise_candidate.py`. +See [docs/uk-staging-operations.md](docs/uk-staging-operations.md). + See [SYSTEM_REQUIREMENTS.md](SYSTEM_REQUIREMENTS.md) for the measured memory, disk, and CPU footprint of developing and building locally (and what to budget on a build machine — RAM is the binding constraint). diff --git a/changelog.d/uk-dense-assembler-staging-evidence.changed.md b/changelog.d/uk-dense-assembler-staging-evidence.changed.md new file mode 100644 index 0000000000..bd841f598f --- /dev/null +++ b/changelog.d/uk-dense-assembler-staging-evidence.changed.md @@ -0,0 +1 @@ +`tools/assemble_uk_dense_release_dir.py` now requires the run manifest's `staging_delivery` receipt and copies it into `build_manifest.json` as `staging`, the national assembler's rule, so `microcosm-publish-release` applies its undelivered-staging refusal to dense cuts. diff --git a/changelog.d/uk-rowwise-staging.added.md b/changelog.d/uk-rowwise-staging.added.md new file mode 100644 index 0000000000..3db838e14c --- /dev/null +++ b/changelog.d/uk-rowwise-staging.added.md @@ -0,0 +1 @@ +Stage UK rowwise candidate builds (dense and `--dataset-households` runs) like the national command: version 2 telemetry to `runs//` with stage events, thinned per-epoch losses and delivery evidence in the manifest; the driver adopts the shared `--staging-*` options. diff --git a/changelog.d/uk-staged-dataset-lane.added.md b/changelog.d/uk-staged-dataset-lane.added.md new file mode 100644 index 0000000000..c7beaee835 --- /dev/null +++ b/changelog.d/uk-staged-dataset-lane.added.md @@ -0,0 +1 @@ +Add the staged-dataset lane: the finished rowwise bundle (every manifest-registered output, the manifest, `staged_manifest.json`, `sha256sums.txt`) is uploaded in one commit to `staged//` of the private UK dataset repository, keyed by the telemetry run id, with `staged_dataset` evidence in the manifest and reviewed `staged_dataset`/`fit_summary` artifacts in the telemetry run; `--staged-dataset-repo-id`, `--no-staged-dataset`, `tools/stage_uk_rowwise_candidate.py` and `tools/fetch_uk_staged_dataset.py`. Staging is not publication: `releases/` and `latest.json` are untouched. diff --git a/docs/uk-dataset-size-plan-355.md b/docs/uk-dataset-size-plan-355.md index 4b2ce4b213..e657f4137e 100644 --- a/docs/uk-dataset-size-plan-355.md +++ b/docs/uk-dataset-size-plan-355.md @@ -93,6 +93,21 @@ Never reduce `--n-clones` to request a smaller output. Full builds still need the dense build's peak memory and add L0/refit work; the reduction is in the exported dataset's storage and downstream loading/simulation footprint. +Every run stages by default (see [UK staging operations](uk-staging-operations.md)): +telemetry to `runs//` in `policyengine/populace-uk-staging` while it +runs, and at the end the bundle the manifest registers, under +`staged//` in `policyengine/populace-uk-private`, keyed by the same +run id (the Logbook build id unless `--staging-run-id` is given). A size run +stays `releasable=false`; staging does not change that, it makes the +candidate inspectable. The run needs an ambient `HF_TOKEN` that can see the +private repository, or `--staging-local-only` to keep everything on disk +(`--no-staging` records a deliberate opt-out, `--no-staged-dataset` keeps only +the bundle local). The command ends with a `staged dataset:` line naming the +repository, prefix and revision, and the manifest carries `staging_delivery` +and `staged_dataset`. Fetch a staged bundle with +`tools/fetch_uk_staged_dataset.py --run-id `; re-stage a finished +directory with `tools/stage_uk_rowwise_candidate.py --run-dir `. + ### The search stops on the draw's own feasibility; the solve is checkpointed before the draw (2026-09-09) S2 on spine-p (55,000 at `--selection-pi-hi 0.95`, 2,000 epochs) was refused at the diff --git a/docs/uk-dense-release-assembly-runbook-762.md b/docs/uk-dense-release-assembly-runbook-762.md index fc2fb830f1..f8bedd93fc 100644 --- a/docs/uk-dense-release-assembly-runbook-762.md +++ b/docs/uk-dense-release-assembly-runbook-762.md @@ -105,8 +105,11 @@ uv run --no-sync python tools/assemble_uk_dense_release_dir.py \ ``` Assembly verifies the hash join (every manifest output against its bytes, the -spine against its pin, the gate report against the Logbook build id), re-runs -the candidate pre-flight, mints the cut tag +spine against its pin, the gate report against the Logbook build id), requires +the manifest's `staging_delivery` receipt (the run's version 2 staging +telemetry evidence, copied into `build_manifest.json` as `staging` so +publication can apply its undelivered-staging refusal, as on the national +lane), re-runs the candidate pre-flight, mints the cut tag `microcosm-uk-2024-25-dense--` from the run's attempt id, clones the H5 beside itself as `microcosm_uk_2025_dense.h5`, stages `build_manifest.json`, `release_manifest.json`, `calibration_diagnostics.json`, diff --git a/docs/uk-staging-operations.md b/docs/uk-staging-operations.md index 6c26abf94b..b18ef13b9c 100644 --- a/docs/uk-staging-operations.md +++ b/docs/uk-staging-operations.md @@ -115,7 +115,9 @@ read-only credential. ## Command modes and files -Both UK commands support these staging modes: +The three UK commands (`tools/build_uk_frs_spine.py`, +`tools/calibrate_uk_national_dataset.py` and +`tools/build_uk_rowwise_candidate.py`) support these staging modes: - Default: local version 2 files plus best-effort delivery to `policyengine/populace-uk-staging`. @@ -149,6 +151,74 @@ artifacts are permitted; population H5 files, NumPy archives, source survey tables, row-level extracts, archives, credentials, and environment data are rejected before remote storage is called. +## Staged datasets + +Telemetry is not the dataset. The rowwise candidate command (dense K=15 and +exact-count `--dataset-households` runs alike) also stages the bundle its +manifest vouches for, so a run can be inspected by the team without being +published. The two destinations share one run id: + +```text +policyengine/populace-uk-staging runs//... telemetry (above) +policyengine/populace-uk-private staged//... the bundle +``` + +The bundle is every file `rowwise_candidate_manifest.json` registers under +`outputs` (the H5, `calibration_diagnostics.json`, the gate report, the CSVs +and the registry), the manifest as built, `staged_manifest.json` (inventory +with digests, summary fields and the telemetry run reference) and +`sha256sums.txt`. Nothing else in the run directory is eligible: logs, the +Logbook spool, checkpoints and evaluation trees stay local. The bundle is +verified from disk against the manifest's own digests and written in one +commit; the commit is recorded as the bundle's revision. `releases/` and +`latest.json` are never touched, no tag is created, and the release contract is +not consulted: a staged bundle is not a release and cannot be loaded as one. + +Modes follow the staging switch. The default uploads telemetry and the bundle; +`--staging-local-only` keeps both on disk (the sidecars are still written); +`--no-staging` disables both and records the opt-out; `--no-staged-dataset` +keeps telemetry remote and the bundle local. The repository is +`--staged-dataset-repo-id` (environment `POPULACE_UK_STAGED_DATASET_REPO_ID`). +Because the upload closes a multi-hour run, the command refuses to start a +remote dataset stage without an ambient Hub credential (`HF_TOKEN`) that can +see the repository; the telemetry stays best-effort with no pre-flight. + +The manifest gains two evidence blocks after the bundle is on disk: +`staging_delivery`, the validated version 2 telemetry receipt, and +`staged_dataset` (`mode`, `repository`, `prefix`, `run_id`, `revision`, +`status` in `uploaded`, `already_staged`, `failed`, `skipped`, a reviewed +`error_code`, and the `files` inventory as a mapping). A failed upload is +recorded with its code, warned on stderr, and never changes the build's exit +code or Logbook disposition. The telemetry run declares two reviewed +artifacts, `artifacts/staged_dataset.json` (the same block) and +`artifacts/fit_summary.json` (loss, fit by family, gate verdicts, the size +receipt without its per-row arrays), so a run in the dashboard points at its +bundle. + +Re-stage a finished directory, including runs built before this lane existed +or whose upload failed, with: + +```bash +uv run python tools/stage_uk_rowwise_candidate.py --run-dir +``` + +It is idempotent on the outputs' digests: an identical remote bundle is +`already_staged`; a different one under the same run id is refused +(`REMOTE_DIFFERS`). Fetch a bundle for a scorecard, an evaluation leg or the +dashboard's local-directory mode with: + +```bash +uv run python tools/fetch_uk_staged_dataset.py --run-id --dest +``` + +Every file is checked against `sha256sums.txt`; `--h5-only` fetches the +dataset alone. Rollback is configuration-first here too: `--no-staged-dataset` +or `--staging-local-only` for new runs, and deleting `staged//` on the +Hub for a bundle that must not remain (its manifest keeps the record). The +dense release assembler applies the national rule to this command's runs: it +requires the manifest's `staging_delivery` receipt and copies it into +`build_manifest.json` as `staging`, where publication reads it. + ## Smoke verification The UK spine command keeps fractional input sampling for scale tests. The diff --git a/experiments/355-uk-dataset-size-receipts.md b/experiments/355-uk-dataset-size-receipts.md index 49b7bec2af..6c848f0209 100644 --- a/experiments/355-uk-dataset-size-receipts.md +++ b/experiments/355-uk-dataset-size-receipts.md @@ -941,3 +941,48 @@ evaluation page's T6 tables). > **Erratum (2026-09-15, microcosm#929).** The "+3 to +12 %" on the VOA council-tax bands is the same keying artifact > as above (the Wales rollup summed into the England-pinned target); corrected, no VOA row diverges and the maximum > divergence is 6.2 % (`isc.private_school_students`); see the erratum under the D2 frozen-versus-recomputed paragraph. + +## Staging for the full build (María, 2026-09-17; plan `repos/uk-rowwise-staging-plan.md`) + +Branch `uk-rowwise-staging` (worktree `repos/populace-staging-rowwise`, cut from main d1196af1). +Vocabulary settled with María: **staging** is telemetry to `runs//` in +`policyengine/populace-uk-staging`; **publishing** is `publish_cli` into `releases/` of +`policyengine/populace-uk-private` with or without moving `latest.json`; the **staged dataset** is +the new lane, the run's outputs under `staged//` of the private repository, same run id, +no release contract. + +### Code landed + +- `tools/build_uk_rowwise_candidate.py` adopts the shared `--staging-*` options and creates the + telemetry with the attempt (spine-builder shape; `run_kind=calibration`, + `operation_id=uk_rowwise_candidate`, `pipeline_id=uk-local-candidate`, `non_release=true`). + Stage events: `input_pinning`, `target_compilation`, `cloning`, `surface_resolution`, + `calibration`, `gate_battery`, `holdout`, `output_bundle`, `dataset_staging`. Epoch events reach + `calibration_progress` through a new `progress_events` fan-out on + `solve_uk_rowwise_weights_under_doctrine`, thinned in the driver to every tenth epoch and each + phase's last epoch (`_STAGING_EPOCH_EVERY = 10`): a 2,000-epoch size run emits up to 24,000 + epochs (dense + ten full-length probes + refit) and the contract caps each file at 5 MiB. The + kernel's `budget_search: True` flag becomes `1` for the contract's integer field. +- New country-neutral `microcosm.build.staging_dataset`: bundle from the manifest's `outputs` + (digests re-verified from disk), `staged_manifest.json` + `sha256sums.txt` sidecars, one + `create_commit` under `staged//`, idempotence on the outputs' digests + (`already_staged` / `REMOTE_DIFFERS`), best-effort with reviewed error codes, `fetch_bundle`. + `HuggingFaceDatasetStorage` gains `commit`, `head_revision`, `file_exists`, `download_file`; + `uk_runtime/staging.py` gains `UK_STAGED_DATASET_REPOSITORY` + (`policyengine/populace-uk-private`, env `POPULACE_UK_STAGED_DATASET_REPO_ID`) and the prefix + `staged`; `staging_cli` gains `--staged-dataset-repo-id` / `--no-staged-dataset`. +- The manifest gains `staging_delivery` (v2) and `staged_dataset` (v1) after the bundle is on + disk (atomic rewrite; the uploaded manifest copy predates them, `staged_manifest.json` describes + the remote side). Two reviewed artifacts on the telemetry run: `staged_dataset.json` and + `fit_summary.json`. A remote dataset stage is refused up front without an ambient Hub credential + that can see the repository; an upload failure is recorded and never changes the exit code. +- `tools/stage_uk_rowwise_candidate.py` (re-stage a finished directory) and + `tools/fetch_uk_staged_dataset.py` (digest-verified fetch by run id). +- `tools/assemble_uk_dense_release_dir.py` requires `staging_delivery` and copies it into + `build_manifest.json` as `staging` (the national assembler's rule). +- Tests: `test_staging_dataset.py` (20), storage/CLI additions in `test_staging_v2.py`, the solver + fan-out in `test_uk_local_rowwise.py`, seven driver tests in `test_uk_rowwise_candidate.py` + (local-only, size phases, `--no-staging`, remote upload + re-stage + fetch, recorded upload + failure, `--no-staged-dataset`, up-front refusals and dry run), dense assembler evidence tests, + size-evaluation tolerance. Every existing driver test passes `--staging-local-only` through the + shared flag helper. diff --git a/packages/microcosm-build/src/microcosm/build/staging_cli.py b/packages/microcosm-build/src/microcosm/build/staging_cli.py index 17493f36f9..834c61a76f 100644 --- a/packages/microcosm-build/src/microcosm/build/staging_cli.py +++ b/packages/microcosm-build/src/microcosm/build/staging_cli.py @@ -9,7 +9,12 @@ from microcosm.build.staging_storage import StagingRepositoryConfig -__all__ = ["add_staging_arguments", "validate_staging_arguments"] +__all__ = [ + "add_staged_dataset_arguments", + "add_staging_arguments", + "validate_staged_dataset_arguments", + "validate_staging_arguments", +] def add_staging_arguments( @@ -82,3 +87,54 @@ def validate_staging_arguments( or not args.staging_repo_id.strip() ): parser.error("remote staging requires a non-empty --staging-repo-id.") + + +def add_staged_dataset_arguments( + parser: argparse.ArgumentParser, + *, + repository: StagingRepositoryConfig, +) -> None: + """Add the staged-dataset options beside the shared staging options. + + The finished bundle follows the staging mode switch: ``--no-staging`` + keeps nothing, ``--staging-local-only`` keeps the bundle and its sidecars + on disk, and the default uploads it to ``repository`` under + ``staged//``. ``--no-staged-dataset`` keeps telemetry remote while + the bundle stays local. + """ + + parser.add_argument( + "--staged-dataset-repo-id", + default=repository.repo_id(os.environ), + help=( + "Private Hugging Face dataset repository that receives the finished " + "run bundle under staged// (default from " + f"{repository.repo_id_environment_variable}, else " + f"{repository.default_repo_id})." + ), + ) + parser.add_argument( + "--no-staged-dataset", + action="store_true", + help=( + "Keep the finished bundle local while staging telemetry still runs; " + "--no-staging already disables both." + ), + ) + + +def validate_staged_dataset_arguments( + parser: argparse.ArgumentParser, args: argparse.Namespace +) -> None: + """Reject staged-dataset options that contradict the staging mode.""" + + if args.no_staged_dataset and args.no_staging: + parser.error( + "--no-staging already disables the staged dataset; " + "drop --no-staged-dataset." + ) + remote = not (args.no_staging or args.staging_local_only or args.no_staged_dataset) + if remote and not str(args.staged_dataset_repo_id or "").strip(): + parser.error( + "remote dataset staging requires a non-empty --staged-dataset-repo-id." + ) diff --git a/packages/microcosm-build/src/microcosm/build/staging_dataset.py b/packages/microcosm-build/src/microcosm/build/staging_dataset.py new file mode 100644 index 0000000000..10cd355b07 --- /dev/null +++ b/packages/microcosm-build/src/microcosm/build/staging_dataset.py @@ -0,0 +1,654 @@ +"""Country-neutral staging of a finished build's dataset bundle. + +Staging telemetry (:mod:`microcosm.build.staging_v2`) carries reviewed +aggregate JSON only; population files never pass its content policy. The +dataset a build produced is staged separately: every file the build's manifest +registers as an output, plus the manifest and two sidecars, goes into a private +dataset repository under ``//`` in one commit, keyed by the same +run id as the telemetry. A staged bundle is inspectable, never published: +nothing here reads or writes a release contract, ``releases/`` or a latest +pointer. + +The module knows nothing about a country's manifest beyond the ``outputs`` +mapping shape ``{name: {"path", "sha256", "bytes"}}``; callers pass the manifest +filename and any summary fields. Transport failures are recorded, never +raised, and no exception text enters the evidence. +""" + +from __future__ import annotations + +import hashlib +import json +import re +import shutil +import sys +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path, PurePosixPath +from typing import Any + +from microcosm.build.staging_storage import HuggingFaceDatasetStorage + +__all__ = [ + "DEFAULT_STAGED_DATASET_PREFIX", + "SHA256SUMS_FILENAME", + "STAGED_DATASET_CONTRACT_VERSION", + "STAGED_MANIFEST_FILENAME", + "STAGED_MANIFEST_SCHEMA", + "StagedDatasetBundle", + "StagedDatasetError", + "StagedFile", + "disabled_staged_dataset", + "fetch_bundle", + "local_only_staged_dataset", + "parse_sha256sums", + "sha256_file", + "stage_bundle", + "validate_staged_dataset_delivery", + "write_sidecars", +] + +STAGED_DATASET_CONTRACT_VERSION = 1 +STAGED_MANIFEST_SCHEMA = "microcosm.staged-dataset.manifest" +STAGED_MANIFEST_FILENAME = "staged_manifest.json" +SHA256SUMS_FILENAME = "sha256sums.txt" +DEFAULT_STAGED_DATASET_PREFIX = "staged" + +_MODES = ("local_and_remote", "local_only", "disabled") +_STATUSES = ("uploaded", "already_staged", "failed", "skipped") +_ERROR_CODES = ("UPLOAD_FAILED", "REMOTE_DIFFERS", "REPOSITORY_UNAVAILABLE") +_DELIVERY_KEYS = frozenset( + { + "contract_version", + "mode", + "repository", + "prefix", + "run_id", + "revision", + "status", + "error_code", + "opt_out_reason", + "files", + } +) +_DEFAULT_SUMMARY_KEYS = ( + "build_kind", + "created_at", + "git_commit", + "git_dirty", + "releasable", + "failing_gate_ids", +) +_SAFE_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,127}$") +_SHA256 = re.compile(r"^[0-9a-f]{64}$") +_SUMS_LINE = re.compile(r"^(?P[0-9a-f]{64}) (?P[^/\\\s][^/\\]*)$") + + +class StagedDatasetError(ValueError): + """A staged-dataset bundle or delivery record is malformed.""" + + +def sha256_file(path: Path | str) -> str: + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _safe_identifier(value: object, label: str) -> str: + if not isinstance(value, str) or not _SAFE_ID.fullmatch(value): + raise StagedDatasetError( + f"{label} must match {_SAFE_ID.pattern}; got {value!r}." + ) + return value + + +def _safe_prefix(prefix: object) -> str: + if not isinstance(prefix, str) or not prefix.strip(): + raise StagedDatasetError("prefix must be a non-empty relative path.") + parts = prefix.strip("/").split("/") + for part in parts: + _safe_identifier(part, "prefix segment") + return "/".join(parts) + + +def _json_bytes(payload: Any) -> bytes: + return ( + json.dumps(payload, indent=2, sort_keys=True, allow_nan=False) + "\n" + ).encode("utf-8") + + +def _jsonable(value: Any) -> Any: + return json.loads(json.dumps(value, default=str, allow_nan=False)) + + +@dataclass(frozen=True) +class StagedFile: + """One manifest-registered output, verified against its recorded digest.""" + + name: str + path: Path + sha256: str + bytes: int + + +@dataclass(frozen=True) +class StagedDatasetBundle: + """The files a build's manifest vouches for, ready to stage.""" + + run_dir: Path + run_id: str + manifest_name: str + manifest_sha256: str + manifest_bytes: int + files: tuple[StagedFile, ...] + summary: dict[str, Any] + + @classmethod + def from_manifest( + cls, + run_dir: Path | str, + *, + run_id: str, + manifest_name: str, + outputs_key: str = "outputs", + summary_keys: Sequence[str] = _DEFAULT_SUMMARY_KEYS, + extra_summary: Mapping[str, Any] | None = None, + ) -> StagedDatasetBundle: + """Read the manifest and verify every registered output on disk. + + Only files the manifest registers are eligible; each must sit in + ``run_dir`` and match its recorded SHA-256 and byte count. Anything + else in the directory (logs, spools, checkpoints) is never staged. + """ + + directory = Path(run_dir).expanduser().resolve() + run_id = _safe_identifier(run_id, "run_id") + manifest_path = directory / manifest_name + if not manifest_path.is_file(): + raise StagedDatasetError(f"manifest missing: {manifest_path}") + try: + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + except ValueError as exc: + raise StagedDatasetError( + f"manifest is not valid JSON: {manifest_path}" + ) from exc + if not isinstance(manifest, Mapping): + raise StagedDatasetError("manifest must be a JSON object.") + outputs = manifest.get(outputs_key) + if not isinstance(outputs, Mapping) or not outputs: + raise StagedDatasetError( + f"manifest registers no {outputs_key!r}; nothing to stage." + ) + reserved = {manifest_name, STAGED_MANIFEST_FILENAME, SHA256SUMS_FILENAME} + files: list[StagedFile] = [] + seen: set[str] = set() + for key in sorted(outputs): + entry = outputs[key] + if not isinstance(entry, Mapping): + raise StagedDatasetError(f"{outputs_key}.{key} must be an object.") + name = PurePosixPath(str(entry.get("path") or "")).name + if not name or name in reserved: + raise StagedDatasetError( + f"{outputs_key}.{key} names an ineligible file: {name!r}." + ) + if name in seen: + raise StagedDatasetError(f"{outputs_key} registers {name!r} twice.") + seen.add(name) + path = directory / name + if not path.is_file(): + raise StagedDatasetError( + f"{outputs_key}.{key} is missing from {directory}: {name}" + ) + recorded = entry.get("sha256") + if not isinstance(recorded, str) or not _SHA256.fullmatch(recorded): + raise StagedDatasetError( + f"{outputs_key}.{key} carries no valid sha256 digest." + ) + measured = sha256_file(path) + if measured != recorded: + raise StagedDatasetError( + f"{outputs_key}.{key} ({name}) does not match its manifest " + "digest; refusing to stage mixed bytes." + ) + size = int(path.stat().st_size) + recorded_bytes = entry.get("bytes") + if ( + isinstance(recorded_bytes, int) + and not isinstance(recorded_bytes, bool) + and recorded_bytes != size + ): + raise StagedDatasetError( + f"{outputs_key}.{key} ({name}) byte count differs from the " + "manifest." + ) + files.append(StagedFile(name=name, path=path, sha256=measured, bytes=size)) + summary = {key: manifest[key] for key in summary_keys if key in manifest} + if extra_summary: + summary.update(dict(extra_summary)) + return cls( + run_dir=directory, + run_id=run_id, + manifest_name=manifest_name, + manifest_sha256=sha256_file(manifest_path), + manifest_bytes=int(manifest_path.stat().st_size), + files=tuple(files), + summary=_jsonable(summary), + ) + + @property + def manifest_path(self) -> Path: + return self.run_dir / self.manifest_name + + def digests(self) -> dict[str, dict[str, Any]]: + return { + item.name: {"sha256": item.sha256, "bytes": item.bytes} + for item in self.files + } + + def remote_prefix(self, prefix: str = DEFAULT_STAGED_DATASET_PREFIX) -> str: + return f"{_safe_prefix(prefix)}/{self.run_id}" + + def staged_manifest( + self, + *, + repository: str | None, + prefix: str = DEFAULT_STAGED_DATASET_PREFIX, + telemetry: Mapping[str, Any] | None = None, + ) -> dict[str, Any]: + """The self-describing inventory written beside the bundle.""" + + return { + "schema_name": STAGED_MANIFEST_SCHEMA, + "schema_version": STAGED_DATASET_CONTRACT_VERSION, + "run_id": self.run_id, + "repository": repository, + "prefix": self.remote_prefix(prefix), + "manifest": { + "name": self.manifest_name, + "sha256": self.manifest_sha256, + "bytes": self.manifest_bytes, + }, + "files": self.digests(), + "summary": dict(self.summary), + "telemetry": None if telemetry is None else _jsonable(dict(telemetry)), + } + + +def write_sidecars( + bundle: StagedDatasetBundle, + *, + repository: str | None, + prefix: str = DEFAULT_STAGED_DATASET_PREFIX, + telemetry: Mapping[str, Any] | None = None, +) -> tuple[Path, Path]: + """Write ``staged_manifest.json`` then ``sha256sums.txt`` into the run dir. + + The sums file lists every output, the manifest as it will be uploaded, and + the staged manifest, so a fetched bundle verifies itself. + """ + + staged_manifest_path = bundle.run_dir / STAGED_MANIFEST_FILENAME + staged_manifest_path.write_bytes( + _json_bytes( + bundle.staged_manifest( + repository=repository, prefix=prefix, telemetry=telemetry + ) + ) + ) + lines = [(item.sha256, item.name) for item in bundle.files] + lines.append((bundle.manifest_sha256, bundle.manifest_name)) + lines.append((sha256_file(staged_manifest_path), STAGED_MANIFEST_FILENAME)) + text = "".join( + f"{sha} {name}\n" for sha, name in sorted(lines, key=lambda item: item[1]) + ) + sums_path = bundle.run_dir / SHA256SUMS_FILENAME + sums_path.write_text(text, encoding="utf-8") + return sums_path, staged_manifest_path + + +def parse_sha256sums(text: str) -> list[tuple[str, str]]: + """Parse ``sha256sum`` lines into ``(digest, name)`` pairs.""" + + entries: list[tuple[str, str]] = [] + seen: set[str] = set() + for raw in text.splitlines(): + line = raw.rstrip("\n") + if not line.strip(): + continue + match = _SUMS_LINE.fullmatch(line) + if match is None: + raise StagedDatasetError(f"malformed sha256sums line: {line!r}") + name = match.group("name") + if name in seen: + raise StagedDatasetError(f"sha256sums lists {name!r} twice.") + seen.add(name) + entries.append((match.group("sha"), name)) + if not entries: + raise StagedDatasetError("sha256sums lists no files.") + return entries + + +def validate_staged_dataset_delivery(payload: Mapping[str, Any]) -> dict[str, Any]: + """Validate and normalize a staged-dataset delivery record.""" + + if not isinstance(payload, Mapping): + raise StagedDatasetError("staged_dataset delivery must be an object.") + record = _jsonable(dict(payload)) + keys = set(record) + if keys != _DELIVERY_KEYS: + missing = sorted(_DELIVERY_KEYS - keys) + extra = sorted(keys - _DELIVERY_KEYS) + raise StagedDatasetError( + f"staged_dataset delivery keys differ: missing {missing}, extra {extra}." + ) + if record["contract_version"] != STAGED_DATASET_CONTRACT_VERSION: + raise StagedDatasetError("staged_dataset contract_version must be 1.") + mode = record["mode"] + status = record["status"] + if mode not in _MODES: + raise StagedDatasetError(f"staged_dataset mode {mode!r} is unknown.") + if status not in _STATUSES: + raise StagedDatasetError(f"staged_dataset status {status!r} is unknown.") + repository = record["repository"] + if repository is not None and ( + not isinstance(repository, str) or not repository.strip() + ): + raise StagedDatasetError("staged_dataset repository must be null or non-empty.") + prefix = record["prefix"] + if prefix is not None: + _safe_prefix(prefix) + run_id = record["run_id"] + if run_id is not None: + _safe_identifier(run_id, "staged_dataset run_id") + revision = record["revision"] + if revision is not None and (not isinstance(revision, str) or not revision.strip()): + raise StagedDatasetError("staged_dataset revision must be null or non-empty.") + error_code = record["error_code"] + if error_code is not None and error_code not in _ERROR_CODES: + raise StagedDatasetError(f"staged_dataset error_code {error_code!r} unknown.") + reason = record["opt_out_reason"] + if reason is not None and (not isinstance(reason, str) or not reason.strip()): + raise StagedDatasetError("staged_dataset opt_out_reason must be non-empty.") + files = record["files"] + if not isinstance(files, Mapping): + raise StagedDatasetError("staged_dataset files must be an object.") + for name, entry in files.items(): + if not isinstance(name, str) or not name or "/" in name: + raise StagedDatasetError(f"staged_dataset files name {name!r} invalid.") + if ( + not isinstance(entry, Mapping) + or set(entry) != {"sha256", "bytes"} + or not isinstance(entry["sha256"], str) + or not _SHA256.fullmatch(entry["sha256"]) + or isinstance(entry["bytes"], bool) + or not isinstance(entry["bytes"], int) + or entry["bytes"] < 0 + ): + raise StagedDatasetError(f"staged_dataset files entry {name!r} invalid.") + if mode == "disabled": + if status != "skipped" or reason is None: + raise StagedDatasetError( + "a disabled staged dataset is skipped with an opt-out reason." + ) + if ( + any( + record[key] is not None + for key in ("repository", "prefix", "run_id", "revision", "error_code") + ) + or files + ): + raise StagedDatasetError("a disabled staged dataset records nothing else.") + return record + if reason is not None: + raise StagedDatasetError( + "only a disabled staged dataset carries an opt-out reason." + ) + if run_id is None or prefix is None or not files: + raise StagedDatasetError( + "an enabled staged dataset names its run_id, prefix and files." + ) + if mode == "local_only": + if status != "skipped" or repository is not None or revision is not None: + raise StagedDatasetError( + "a local-only staged dataset is skipped without a repository or revision." + ) + if error_code is not None: + raise StagedDatasetError("a local-only staged dataset has no error code.") + return record + if repository is None: + raise StagedDatasetError("a remote staged dataset names its repository.") + if status == "skipped": + raise StagedDatasetError( + "a remote staged dataset is uploaded, already staged or failed." + ) + if status == "failed": + if error_code is None or revision is not None: + raise StagedDatasetError( + "a failed staged dataset carries an error code and no revision." + ) + return record + if error_code is not None: + raise StagedDatasetError("a delivered staged dataset carries no error code.") + if status == "uploaded" and revision is None: + raise StagedDatasetError("an uploaded staged dataset records its revision.") + return record + + +def disabled_staged_dataset(reason: str) -> dict[str, Any]: + """Explicit evidence for a deliberate staged-dataset opt-out.""" + + if not isinstance(reason, str) or not reason.strip(): + raise StagedDatasetError("A staged-dataset opt-out requires a reason.") + return validate_staged_dataset_delivery( + { + "contract_version": STAGED_DATASET_CONTRACT_VERSION, + "mode": "disabled", + "repository": None, + "prefix": None, + "run_id": None, + "revision": None, + "status": "skipped", + "error_code": None, + "opt_out_reason": reason.strip(), + "files": {}, + } + ) + + +def local_only_staged_dataset( + bundle: StagedDatasetBundle, *, prefix: str = DEFAULT_STAGED_DATASET_PREFIX +) -> dict[str, Any]: + """Evidence for a bundle verified and inventoried on disk but not uploaded.""" + + return validate_staged_dataset_delivery( + { + "contract_version": STAGED_DATASET_CONTRACT_VERSION, + "mode": "local_only", + "repository": None, + "prefix": bundle.remote_prefix(prefix), + "run_id": bundle.run_id, + "revision": None, + "status": "skipped", + "error_code": None, + "opt_out_reason": None, + "files": bundle.digests(), + } + ) + + +def _warn(message: str) -> None: + print(f"warning: {message}", file=sys.stderr, flush=True) + + +def stage_bundle( + bundle: StagedDatasetBundle, + *, + storage: HuggingFaceDatasetStorage, + prefix: str = DEFAULT_STAGED_DATASET_PREFIX, + message: str | None = None, +) -> dict[str, Any]: + """Upload the bundle and its sidecars in one commit; record, never raise. + + Idempotent on the outputs' digests: a remote ``staged_manifest.json`` whose + ``files`` equal the local digests is ``already_staged``; a differing one + is refused as ``REMOTE_DIFFERS`` so a run id never silently changes + meaning. Transport failures return ``status: failed`` with a reviewed + error code; the exception text goes to stderr only. + """ + + remote_prefix = bundle.remote_prefix(prefix) + sidecars = ( + bundle.run_dir / STAGED_MANIFEST_FILENAME, + bundle.run_dir / SHA256SUMS_FILENAME, + ) + for sidecar in sidecars: + if not sidecar.is_file(): + raise StagedDatasetError( + f"{sidecar.name} missing; call write_sidecars before stage_bundle." + ) + record: dict[str, Any] = { + "contract_version": STAGED_DATASET_CONTRACT_VERSION, + "mode": "local_and_remote", + "repository": storage.repo_id, + "prefix": remote_prefix, + "run_id": bundle.run_id, + "revision": None, + "status": "failed", + "error_code": None, + "opt_out_reason": None, + "files": bundle.digests(), + } + remote_manifest = f"{remote_prefix}/{STAGED_MANIFEST_FILENAME}" + try: + exists = storage.file_exists(remote_manifest) + except Exception as error: + _warn( + f"staged dataset repository {storage.repo_id} is unavailable " + f"({type(error).__name__}); the bundle stays local." + ) + return validate_staged_dataset_delivery( + {**record, "error_code": "REPOSITORY_UNAVAILABLE"} + ) + if exists: + try: + remote = json.loads(storage.download(remote_manifest)) + head = storage.head_revision() + except Exception as error: + _warn( + f"staged dataset at {remote_prefix} could not be read back " + f"({type(error).__name__}); the bundle stays local." + ) + return validate_staged_dataset_delivery( + {**record, "error_code": "REPOSITORY_UNAVAILABLE"} + ) + remote_files = remote.get("files") if isinstance(remote, Mapping) else None + if _same_digests(remote_files, bundle.digests()): + return validate_staged_dataset_delivery( + {**record, "status": "already_staged", "revision": head} + ) + _warn( + f"{remote_prefix} already holds a different bundle; refusing to " + "overwrite it. Re-stage under another run id if that is intended." + ) + return validate_staged_dataset_delivery( + {**record, "error_code": "REMOTE_DIFFERS"} + ) + try: + from huggingface_hub import CommitOperationAdd + except ImportError as exc: # pragma: no cover - declared dependency + raise ImportError( + "microcosm-build needs huggingface_hub to stage datasets; " + "reinstall microcosm-build with its dependencies." + ) from exc + operations = [ + CommitOperationAdd( + path_in_repo=f"{remote_prefix}/{item.name}", + path_or_fileobj=str(item.path), + ) + for item in bundle.files + ] + operations.append( + CommitOperationAdd( + path_in_repo=f"{remote_prefix}/{bundle.manifest_name}", + path_or_fileobj=str(bundle.manifest_path), + ) + ) + operations.extend( + CommitOperationAdd( + path_in_repo=f"{remote_prefix}/{sidecar.name}", + path_or_fileobj=str(sidecar), + ) + for sidecar in sidecars + ) + try: + parent = storage.head_revision() + revision = storage.commit( + operations, + message=message or f"Stage {bundle.run_id} under {remote_prefix}", + parent_commit=parent, + ) + except Exception as error: + _warn( + f"staged dataset upload failed for {remote_prefix} " + f"({type(error).__name__}); the bundle stays local and can be " + "re-staged." + ) + return validate_staged_dataset_delivery( + {**record, "error_code": "UPLOAD_FAILED"} + ) + return validate_staged_dataset_delivery( + {**record, "status": "uploaded", "revision": revision} + ) + + +def _same_digests(remote: Any, local: Mapping[str, Mapping[str, Any]]) -> bool: + if not isinstance(remote, Mapping) or set(remote) != set(local): + return False + for name, entry in local.items(): + other = remote.get(name) + if not isinstance(other, Mapping) or other.get("sha256") != entry["sha256"]: + return False + return True + + +def fetch_bundle( + *, + storage: HuggingFaceDatasetStorage, + run_id: str, + dest: Path | str, + prefix: str = DEFAULT_STAGED_DATASET_PREFIX, + revision: str | None = None, + h5_only: bool = False, +) -> list[Path]: + """Download a staged bundle into ``dest`` and verify every digest. + + Files are copied out of the Hub cache so each lands under its own name + (a ``.h5`` path is what the simulation loaders require). A digest + mismatch removes the file and raises. + """ + + remote_prefix = f"{_safe_prefix(prefix)}/{_safe_identifier(run_id, 'run_id')}" + destination = Path(dest).expanduser().resolve() + destination.mkdir(parents=True, exist_ok=True) + sums_local = storage.download_file( + f"{remote_prefix}/{SHA256SUMS_FILENAME}", revision=revision + ) + entries = parse_sha256sums(Path(sums_local).read_text(encoding="utf-8")) + fetched: list[Path] = [] + for sha, name in entries: + if h5_only and not name.lower().endswith((".h5", ".hdf5")): + continue + local = storage.download_file(f"{remote_prefix}/{name}", revision=revision) + target = destination / name + shutil.copyfile(local, target) + measured = sha256_file(target) + if measured != sha: + target.unlink(missing_ok=True) + raise StagedDatasetError( + f"{name} does not match the staged digest after download." + ) + fetched.append(target) + shutil.copyfile(sums_local, destination / SHA256SUMS_FILENAME) + return fetched diff --git a/packages/microcosm-build/src/microcosm/build/staging_storage.py b/packages/microcosm-build/src/microcosm/build/staging_storage.py index 8626a22e53..93d570e4a1 100644 --- a/packages/microcosm-build/src/microcosm/build/staging_storage.py +++ b/packages/microcosm-build/src/microcosm/build/staging_storage.py @@ -2,7 +2,7 @@ from __future__ import annotations -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path from typing import Any @@ -84,6 +84,72 @@ def download(self, path_in_repo: str) -> bytes: ) return Path(local).read_bytes() + def download_file(self, path_in_repo: str, *, revision: str | None = None) -> Path: + """Fetch one file through the Hub cache and return its local path. + + Unlike :meth:`download` this never reads the bytes into memory, so it + serves population files as well as small documents. + """ + + api = self._api() + download = getattr(api, "hf_hub_download", None) + if download is None: + from huggingface_hub import hf_hub_download as download + + local = download( + repo_id=self.repo_id, + filename=path_in_repo, + repo_type="dataset", + revision=revision, + ) + return Path(local) + + def file_exists(self, path_in_repo: str) -> bool: + return bool( + self._api().file_exists( + repo_id=self.repo_id, + filename=path_in_repo, + repo_type="dataset", + ) + ) + + def head_revision(self) -> str | None: + """The default branch's current commit, when the backend reports one.""" + + info = self._api().repo_info(repo_id=self.repo_id, repo_type="dataset") + sha = ( + info.get("sha") if isinstance(info, Mapping) else getattr(info, "sha", None) + ) + return str(sha) if sha else None + + 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. + + ``operations`` are ``huggingface_hub.CommitOperation`` values built by + the caller; pinning ``parent_commit`` makes a concurrent write fail + instead of silently 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() + @dataclass(frozen=True) class UploadResult: @@ -136,9 +202,7 @@ def upload( self.storage.upload(local_path, path_in_repo) except Exception as error: self.consecutive_failures += 1 - became_disabled = ( - self.consecutive_failures >= self.max_consecutive_failures - ) + became_disabled = self.consecutive_failures >= self.max_consecutive_failures if became_disabled: self.enabled = False return UploadResult( diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/local_rowwise.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/local_rowwise.py index 0b105cc71b..cd9056208c 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/local_rowwise.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/local_rowwise.py @@ -1040,6 +1040,33 @@ def vector(frame: Frame) -> np.ndarray: return TargetSet(targets) +def _progress_fan_out( + progress: Callable[[str], None] | None, + progress_events: Callable[[dict[str, object]], None] | None, +) -> Callable[[dict[str, object]], None] | None: + """One calibrator callback that feeds the line sink and the event sink.""" + + sinks: list[Callable[[dict[str, object]], None]] = [] + if progress is not None: + from microcosm.build.uk_runtime.solve_progress import ( + uk_solve_progress_callback, + ) + + sinks.append(uk_solve_progress_callback(progress)) + if progress_events is not None: + sinks.append(progress_events) + if not sinks: + return None + if len(sinks) == 1: + return sinks[0] + + def callback(event: dict[str, object]) -> None: + for sink in sinks: + sink(event) + + return callback + + def solve_uk_rowwise_weights_under_doctrine( frame: Frame, problem: UKRowwiseLocalMatrix, @@ -1064,9 +1091,16 @@ def solve_uk_rowwise_weights_under_doctrine( checkpoint_identity: Mapping[str, Any] | None = None, checkpoint_provenance: Mapping[str, Any] | None = None, progress: Callable[[str], None] | None = None, + progress_events: Callable[[dict[str, object]], None] | None = None, ) -> UKRowwiseDoctrineSolve: """Solve rowwise household weights under the reviewed doctrine. + ``progress`` receives readable lines (every 100 epochs, each probe, the + search stop); ``progress_events`` receives every raw calibrator event as a + dict (``calibration_epoch``, ``budget_probe``, ``budget_search_done``), + phase-tagged by the size machinery, so a build driver can publish staging + telemetry without changing the lines a log reader follows. + ``selection_seed`` (default ``seed``) seeds only the size selection — the informed L0 search, the exact-count draw and the refit — so two selections can be compared on one pool and one dense reference. @@ -1206,13 +1240,7 @@ def solve_uk_rowwise_weights_under_doctrine( raise ValueError("size checkpoints apply to a dataset_households solve.") if size_checkpoint_dir is not None and resume_size_checkpoint is not None: raise ValueError("a resumed solve does not write a second checkpoint.") - progress_callback = None - if progress is not None: - from microcosm.build.uk_runtime.solve_progress import ( - uk_solve_progress_callback, - ) - - progress_callback = uk_solve_progress_callback(progress) + progress_callback = _progress_fan_out(progress, progress_events) restored = None if resume_size_checkpoint is not None: from microcosm.build.uk_runtime.size_checkpoint import load_uk_size_checkpoint diff --git a/packages/microcosm-build/src/microcosm/build/uk_runtime/staging.py b/packages/microcosm-build/src/microcosm/build/uk_runtime/staging.py index a98cd17ee8..c70bdc2a1b 100644 --- a/packages/microcosm-build/src/microcosm/build/uk_runtime/staging.py +++ b/packages/microcosm-build/src/microcosm/build/uk_runtime/staging.py @@ -1,10 +1,29 @@ -"""UK-owned configuration for the country-neutral staging implementation.""" +"""UK-owned configuration for the country-neutral staging implementation. + +Two destinations, one run id. Telemetry (reviewed aggregate JSON) goes to the +staging repository under ``runs//``; the finished dataset bundle a +build's manifest vouches for goes to the private artifact repository under +``staged//`` (:mod:`microcosm.build.staging_dataset`), where it is +inspectable without being published: ``releases/`` and ``latest.json`` are +never touched by staging. +""" from microcosm.build.staging_storage import StagingRepositoryConfig -__all__ = ["UK_STAGING_REPOSITORY"] +__all__ = [ + "UK_STAGED_DATASET_PREFIX", + "UK_STAGED_DATASET_REPOSITORY", + "UK_STAGING_REPOSITORY", +] UK_STAGING_REPOSITORY = StagingRepositoryConfig( default_repo_id="policyengine/populace-uk-staging", repo_id_environment_variable="POPULACE_UK_STAGING_REPO_ID", ) + +UK_STAGED_DATASET_REPOSITORY = StagingRepositoryConfig( + default_repo_id="policyengine/populace-uk-private", + repo_id_environment_variable="POPULACE_UK_STAGED_DATASET_REPO_ID", +) + +UK_STAGED_DATASET_PREFIX = "staged" diff --git a/packages/microcosm-build/tests/test_staging_dataset.py b/packages/microcosm-build/tests/test_staging_dataset.py new file mode 100644 index 0000000000..c5bc1399fa --- /dev/null +++ b/packages/microcosm-build/tests/test_staging_dataset.py @@ -0,0 +1,377 @@ +"""The staged-dataset lane: bundle from a manifest, one-commit upload, fetch.""" + +from __future__ import annotations + +import hashlib +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from microcosm.build.staging_dataset import ( + SHA256SUMS_FILENAME, + STAGED_MANIFEST_FILENAME, + StagedDatasetBundle, + StagedDatasetError, + disabled_staged_dataset, + fetch_bundle, + local_only_staged_dataset, + parse_sha256sums, + stage_bundle, + validate_staged_dataset_delivery, + write_sidecars, +) +from microcosm.build.staging_storage import HuggingFaceDatasetStorage + +MANIFEST = "rowwise_candidate_manifest.json" +RUN_ID = "uk-local-candidate-f100-s7-20260917T180000Z-0badcafe" + + +def _sha(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def _write_run(root: Path, *, manifest_extra: dict | None = None) -> Path: + run_dir = root / "run" + run_dir.mkdir() + files = { + "microcosm_uk_2025_local.h5": b"h5-bytes" * 100, + "solve_diagnostics.csv": b"a,b\n1,2\n", + "calibration_diagnostics.json": json.dumps({"targets": []}).encode(), + } + outputs = {} + for name, data in files.items(): + (run_dir / name).write_bytes(data) + outputs[name.split(".")[0]] = { + "path": str(run_dir / name), + "sha256": _sha(data), + "bytes": len(data), + } + (run_dir / "run.log").write_text("not staged\n") + manifest = { + "build_kind": "uk_rowwise_calibrated_candidate", + "releasable": False, + "git_commit": "b" * 40, + "outputs": outputs, + **(manifest_extra or {}), + } + (run_dir / MANIFEST).write_text(json.dumps(manifest, indent=2, sort_keys=True)) + return run_dir + + +class FakeHub: + """Records commits and serves files back, like a private dataset repo.""" + + def __init__(self, root: Path, *, fail_commit: bool = False) -> None: + self.root = root + self.files: dict[str, bytes] = {} + self.commits: list[dict] = [] + self.fail_commit = fail_commit + self.sha = "0" * 40 + + def file_exists(self, *, repo_id, filename, repo_type): + assert repo_type == "dataset" + return filename in self.files + + def repo_info(self, *, repo_id, repo_type): + assert repo_type == "dataset" + return SimpleNamespace(sha=self.sha) + + def create_commit( + self, *, repo_id, operations, commit_message, repo_type, parent_commit + ): + assert repo_type == "dataset" + if self.fail_commit: + raise RuntimeError("401 Unauthorized token=do-not-record") + assert parent_commit == self.sha + for operation in operations: + self.files[operation.path_in_repo] = Path( + operation.path_or_fileobj + ).read_bytes() + self.commits.append( + { + "message": commit_message, + "paths": sorted(self.files), + "parent": parent_commit, + } + ) + self.sha = _sha(commit_message.encode())[:40] + return SimpleNamespace(oid=self.sha) + + def hf_hub_download(self, *, repo_id, filename, repo_type, **kwargs): + if filename not in self.files: + raise FileNotFoundError(filename) + destination = self.root / "cache" / filename + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(self.files[filename]) + return str(destination) + + +def _storage(hub: FakeHub) -> HuggingFaceDatasetStorage: + return HuggingFaceDatasetStorage("policyengine/populace-uk-private", api=hub) + + +def test_bundle_takes_exactly_the_manifest_outputs(tmp_path): + run_dir = _write_run(tmp_path) + bundle = StagedDatasetBundle.from_manifest( + run_dir, run_id=RUN_ID, manifest_name=MANIFEST + ) + assert [item.name for item in bundle.files] == [ + "calibration_diagnostics.json", + "microcosm_uk_2025_local.h5", + "solve_diagnostics.csv", + ] + assert "run.log" not in bundle.digests() + assert bundle.summary == { + "build_kind": "uk_rowwise_calibrated_candidate", + "git_commit": "b" * 40, + "releasable": False, + } + assert bundle.remote_prefix() == f"staged/{RUN_ID}" + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + (lambda d: (d / "solve_diagnostics.csv").unlink(), "missing"), + (lambda d: (d / "solve_diagnostics.csv").write_bytes(b"x"), "digest"), + ], +) +def test_bundle_refuses_missing_or_changed_outputs(tmp_path, mutation, message): + run_dir = _write_run(tmp_path) + mutation(run_dir) + with pytest.raises(StagedDatasetError, match=message): + StagedDatasetBundle.from_manifest( + run_dir, run_id=RUN_ID, manifest_name=MANIFEST + ) + + +def test_bundle_refuses_an_unsafe_run_id_and_an_empty_manifest(tmp_path): + run_dir = _write_run(tmp_path) + with pytest.raises(StagedDatasetError, match="run_id"): + StagedDatasetBundle.from_manifest( + run_dir, run_id="../escape", manifest_name=MANIFEST + ) + (run_dir / MANIFEST).write_text(json.dumps({"outputs": {}})) + with pytest.raises(StagedDatasetError, match="nothing to stage"): + StagedDatasetBundle.from_manifest( + run_dir, run_id=RUN_ID, manifest_name=MANIFEST + ) + + +def test_sidecars_list_every_file_and_verify_themselves(tmp_path): + run_dir = _write_run(tmp_path) + bundle = StagedDatasetBundle.from_manifest( + run_dir, run_id=RUN_ID, manifest_name=MANIFEST + ) + sums, staged_manifest = write_sidecars( + bundle, repository="policyengine/populace-uk-private" + ) + entries = dict((name, sha) for sha, name in parse_sha256sums(sums.read_text())) + assert set(entries) == { + "calibration_diagnostics.json", + "microcosm_uk_2025_local.h5", + "solve_diagnostics.csv", + MANIFEST, + STAGED_MANIFEST_FILENAME, + } + for name, sha in entries.items(): + assert _sha((run_dir / name).read_bytes()) == sha + payload = json.loads(staged_manifest.read_text()) + assert payload["schema_name"] == "microcosm.staged-dataset.manifest" + assert payload["prefix"] == f"staged/{RUN_ID}" + assert payload["files"] == bundle.digests() + assert payload["manifest"]["sha256"] == _sha((run_dir / MANIFEST).read_bytes()) + # Byte-identical on a second write: no timestamps, no absolute paths. + before = staged_manifest.read_bytes() + write_sidecars(bundle, repository="policyengine/populace-uk-private") + assert staged_manifest.read_bytes() == before + assert str(run_dir) not in staged_manifest.read_text() + + +def test_stage_bundle_commits_everything_once_under_the_run_prefix(tmp_path): + run_dir = _write_run(tmp_path) + bundle = StagedDatasetBundle.from_manifest( + run_dir, run_id=RUN_ID, manifest_name=MANIFEST + ) + write_sidecars(bundle, repository="policyengine/populace-uk-private") + hub = FakeHub(tmp_path) + delivery = stage_bundle(bundle, storage=_storage(hub)) + assert delivery["status"] == "uploaded" + assert delivery["revision"] == hub.sha + assert delivery["repository"] == "policyengine/populace-uk-private" + assert delivery["prefix"] == f"staged/{RUN_ID}" + assert len(hub.commits) == 1 + assert hub.commits[0]["parent"] == "0" * 40 + assert hub.commits[0]["paths"] == sorted( + f"staged/{RUN_ID}/{name}" + for name in ( + "calibration_diagnostics.json", + "microcosm_uk_2025_local.h5", + "solve_diagnostics.csv", + MANIFEST, + STAGED_MANIFEST_FILENAME, + SHA256SUMS_FILENAME, + ) + ) + assert validate_staged_dataset_delivery(delivery) == delivery + + +def test_stage_bundle_is_idempotent_on_the_output_digests(tmp_path): + run_dir = _write_run(tmp_path) + bundle = StagedDatasetBundle.from_manifest( + run_dir, run_id=RUN_ID, manifest_name=MANIFEST + ) + write_sidecars(bundle, repository="policyengine/populace-uk-private") + hub = FakeHub(tmp_path) + first = stage_bundle(bundle, storage=_storage(hub)) + # The driver appends evidence to the local manifest after the upload; + # a re-stage must still recognise the same outputs. + manifest = json.loads((run_dir / MANIFEST).read_text()) + manifest["staged_dataset"] = first + (run_dir / MANIFEST).write_text(json.dumps(manifest, indent=2, sort_keys=True)) + again = StagedDatasetBundle.from_manifest( + run_dir, run_id=RUN_ID, manifest_name=MANIFEST + ) + write_sidecars(again, repository="policyengine/populace-uk-private") + second = stage_bundle(again, storage=_storage(hub)) + assert second["status"] == "already_staged" + assert second["revision"] == hub.sha + assert len(hub.commits) == 1 + + (run_dir / "solve_diagnostics.csv").write_bytes(b"a,b\n9,9\n") + manifest["outputs"]["solve_diagnostics"]["sha256"] = _sha(b"a,b\n9,9\n") + manifest["outputs"]["solve_diagnostics"]["bytes"] = len(b"a,b\n9,9\n") + (run_dir / MANIFEST).write_text(json.dumps(manifest, indent=2, sort_keys=True)) + changed = StagedDatasetBundle.from_manifest( + run_dir, run_id=RUN_ID, manifest_name=MANIFEST + ) + write_sidecars(changed, repository="policyengine/populace-uk-private") + refused = stage_bundle(changed, storage=_storage(hub)) + assert refused["status"] == "failed" + assert refused["error_code"] == "REMOTE_DIFFERS" + assert len(hub.commits) == 1 + + +def test_stage_bundle_records_transport_failure_without_raising(tmp_path, capsys): + run_dir = _write_run(tmp_path) + bundle = StagedDatasetBundle.from_manifest( + run_dir, run_id=RUN_ID, manifest_name=MANIFEST + ) + write_sidecars(bundle, repository="policyengine/populace-uk-private") + hub = FakeHub(tmp_path, fail_commit=True) + delivery = stage_bundle(bundle, storage=_storage(hub)) + assert delivery["status"] == "failed" + assert delivery["error_code"] == "UPLOAD_FAILED" + assert delivery["revision"] is None + assert delivery["files"] == bundle.digests() + err = capsys.readouterr().err + assert "staged dataset upload failed" in err + assert "do-not-record" not in err and "do-not-record" not in json.dumps(delivery) + + +def test_stage_bundle_requires_the_sidecars_first(tmp_path): + run_dir = _write_run(tmp_path) + bundle = StagedDatasetBundle.from_manifest( + run_dir, run_id=RUN_ID, manifest_name=MANIFEST + ) + with pytest.raises(StagedDatasetError, match="write_sidecars"): + stage_bundle(bundle, storage=_storage(FakeHub(tmp_path))) + + +def test_fetch_bundle_verifies_every_digest_and_copies_out_of_the_cache(tmp_path): + run_dir = _write_run(tmp_path) + bundle = StagedDatasetBundle.from_manifest( + run_dir, run_id=RUN_ID, manifest_name=MANIFEST + ) + write_sidecars(bundle, repository="policyengine/populace-uk-private") + hub = FakeHub(tmp_path) + stage_bundle(bundle, storage=_storage(hub)) + + dest = tmp_path / "fetched" + paths = fetch_bundle(storage=_storage(hub), run_id=RUN_ID, dest=dest) + assert {path.name for path in paths} == { + "calibration_diagnostics.json", + "microcosm_uk_2025_local.h5", + "solve_diagnostics.csv", + MANIFEST, + STAGED_MANIFEST_FILENAME, + } + assert (dest / "microcosm_uk_2025_local.h5").read_bytes() == b"h5-bytes" * 100 + assert (dest / SHA256SUMS_FILENAME).is_file() + + only = fetch_bundle( + storage=_storage(hub), run_id=RUN_ID, dest=tmp_path / "h5", h5_only=True + ) + assert [path.name for path in only] == ["microcosm_uk_2025_local.h5"] + + hub.files[f"staged/{RUN_ID}/solve_diagnostics.csv"] = b"tampered" + with pytest.raises(StagedDatasetError, match="does not match"): + fetch_bundle(storage=_storage(hub), run_id=RUN_ID, dest=tmp_path / "bad") + assert not (tmp_path / "bad" / "solve_diagnostics.csv").exists() + + +def test_delivery_records_follow_their_mode(tmp_path): + run_dir = _write_run(tmp_path) + bundle = StagedDatasetBundle.from_manifest( + run_dir, run_id=RUN_ID, manifest_name=MANIFEST + ) + disabled = disabled_staged_dataset("--no-staging") + assert disabled["mode"] == "disabled" and disabled["status"] == "skipped" + assert disabled["files"] == {} + local = local_only_staged_dataset(bundle) + assert local["mode"] == "local_only" and local["status"] == "skipped" + assert local["repository"] is None and local["files"] == bundle.digests() + with pytest.raises(StagedDatasetError): + disabled_staged_dataset("") + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"status": "uploaded", "revision": None}, "records its revision"), + ({"status": "failed", "error_code": None}, "error code"), + ({"status": "skipped"}, "uploaded, already staged or failed"), + ({"opt_out_reason": "x"}, "opt-out reason"), + ({"error_code": "BOOM", "status": "failed"}, "unknown"), + ( + { + "mode": "local_only", + "repository": None, + "revision": None, + "status": "uploaded", + }, + "local-only", + ), + ({"extra": 1}, "keys differ"), + ({"files": {"a.h5": {"sha256": "zz", "bytes": 1}}}, "invalid"), + ], +) +def test_delivery_validation_refuses_contradictions(overrides, message): + base = { + "contract_version": 1, + "mode": "local_and_remote", + "repository": "policyengine/populace-uk-private", + "prefix": f"staged/{RUN_ID}", + "run_id": RUN_ID, + "revision": "a" * 40, + "status": "uploaded", + "error_code": None, + "opt_out_reason": None, + "files": {"microcosm_uk_2025_local.h5": {"sha256": "0" * 64, "bytes": 3}}, + } + assert validate_staged_dataset_delivery(base) == base + with pytest.raises(StagedDatasetError, match=message): + validate_staged_dataset_delivery({**base, **overrides}) + + +def test_parse_sha256sums_rejects_paths_and_duplicates(): + good = f"{'a' * 64} file.h5\n{'b' * 64} other.json\n" + assert parse_sha256sums(good) == [("a" * 64, "file.h5"), ("b" * 64, "other.json")] + with pytest.raises(StagedDatasetError, match="malformed"): + parse_sha256sums(f"{'a' * 64} ../escape.h5\n") + with pytest.raises(StagedDatasetError, match="twice"): + parse_sha256sums(f"{'a' * 64} file.h5\n{'b' * 64} file.h5\n") + with pytest.raises(StagedDatasetError, match="no files"): + parse_sha256sums("\n") diff --git a/packages/microcosm-build/tests/test_staging_v2.py b/packages/microcosm-build/tests/test_staging_v2.py index 2607e51f6d..90419176c6 100644 --- a/packages/microcosm-build/tests/test_staging_v2.py +++ b/packages/microcosm-build/tests/test_staging_v2.py @@ -7,9 +7,15 @@ import pytest -from microcosm.build.staging_cli import add_staging_arguments +from microcosm.build.staging_cli import ( + add_staged_dataset_arguments, + add_staging_arguments, + validate_staged_dataset_arguments, + validate_staging_arguments, +) from microcosm.build.staging_storage import ( BestEffortUploadSession, + HuggingFaceDatasetStorage, StagingRepositoryConfig, ) from microcosm.build.staging_v2 import ( @@ -772,3 +778,109 @@ def test_typed_artifacts_reject_one_scalar_individual_record(tmp_path): artifact_kind="aggregate_diagnostics", classification="aggregate", ) + + +def test_staged_dataset_cli_follows_the_staging_mode_switch(monkeypatch): + telemetry_repository = StagingRepositoryConfig( + default_repo_id="example/default-staging", + repo_id_environment_variable="EXAMPLE_STAGING_REPO_ID", + ) + dataset_repository = StagingRepositoryConfig( + default_repo_id="example/default-private", + repo_id_environment_variable="EXAMPLE_STAGED_DATASET_REPO_ID", + ) + monkeypatch.setenv("EXAMPLE_STAGED_DATASET_REPO_ID", "example/configured-private") + + def parse(argv): + parser = argparse.ArgumentParser() + add_staging_arguments(parser, repository=telemetry_repository) + add_staged_dataset_arguments(parser, repository=dataset_repository) + args = parser.parse_args(argv) + validate_staging_arguments(parser, args) + validate_staged_dataset_arguments(parser, args) + return args + + args = parse([]) + assert args.staged_dataset_repo_id == "example/configured-private" + assert args.no_staged_dataset is False + assert parse(["--no-staged-dataset"]).no_staged_dataset is True + assert parse( + ["--staging-local-only", "--staged-dataset-repo-id", ""] + ).staging_local_only + with pytest.raises(SystemExit): + parse(["--no-staging", "--no-staged-dataset"]) + with pytest.raises(SystemExit): + parse(["--staged-dataset-repo-id", " "]) + + +class CommitApi(MemoryApi): + """MemoryApi plus the commit surface the staged-dataset lane uses.""" + + def __init__(self, root: Path, repo_id: str = "example/private") -> None: + super().__init__(root, repo_id) + self.sha = "1" * 40 + self.commits: list[tuple[str, str | None, list[str]]] = [] + + def file_exists(self, *, repo_id, filename, repo_type): + assert repo_id == self.repo_id and repo_type == "dataset" + return filename in self.files + + def repo_info(self, *, repo_id, repo_type): + assert repo_id == self.repo_id and repo_type == "dataset" + return {"sha": self.sha} + + def create_commit( + self, *, repo_id, operations, commit_message, repo_type, parent_commit + ): + assert repo_id == self.repo_id and repo_type == "dataset" + paths = [] + for operation in operations: + self.files[operation.path_in_repo] = Path( + operation.path_or_fileobj + ).read_bytes() + paths.append(operation.path_in_repo) + self.commits.append((commit_message, parent_commit, paths)) + self.sha = "2" * 40 + return {"oid": self.sha} + + +def test_storage_commits_several_files_at_once_and_reports_the_revision(tmp_path): + from huggingface_hub import CommitOperationAdd + + api = CommitApi(tmp_path) + storage = HuggingFaceDatasetStorage("example/private", api=api) + first = tmp_path / "a.json" + second = tmp_path / "b.h5" + first.write_text("{}") + second.write_bytes(b"binary") + + assert storage.head_revision() == "1" * 40 + assert storage.file_exists("staged/run/a.json") is False + revision = storage.commit( + [ + CommitOperationAdd( + path_in_repo="staged/run/a.json", path_or_fileobj=str(first) + ), + CommitOperationAdd( + path_in_repo="staged/run/b.h5", path_or_fileobj=str(second) + ), + ], + message="Stage run", + parent_commit="1" * 40, + ) + assert revision == "2" * 40 + assert api.commits == [ + ("Stage run", "1" * 40, ["staged/run/a.json", "staged/run/b.h5"]) + ] + assert storage.file_exists("staged/run/b.h5") is True + assert storage.download_file("staged/run/b.h5").read_bytes() == b"binary" + + +def test_storage_commit_refuses_a_revisionless_backend(tmp_path): + class NoRevision(CommitApi): + def create_commit(self, **kwargs): + return {"oid": ""} + + storage = HuggingFaceDatasetStorage("example/private", api=NoRevision(tmp_path)) + with pytest.raises(RuntimeError, match="no revision"): + storage.commit([], message="empty") diff --git a/packages/microcosm-build/tests/test_uk_dense_release_assembler.py b/packages/microcosm-build/tests/test_uk_dense_release_assembler.py index beb8828772..bde4d5bab4 100644 --- a/packages/microcosm-build/tests/test_uk_dense_release_assembler.py +++ b/packages/microcosm-build/tests/test_uk_dense_release_assembler.py @@ -94,6 +94,24 @@ def _signed_report(): return report +def _staging_delivery(**overrides) -> dict: + """A validated version 2 telemetry receipt the run carried (local-only).""" + + return { + "contract_version": 2, + "enabled": True, + "mode": "local_only", + "run_id": ATTEMPT, + "configured_repository": None, + "upload_attempts": 0, + "upload_successes": 0, + "read_back": "not_requested", + "last_error_code": None, + "opt_out_reason": None, + **overrides, + } + + def _candidate_dir(root: Path) -> tuple[Path, Path, Path]: candidate = root / "f100-k15-RC" (candidate / "logbook-spool").mkdir(parents=True) @@ -263,6 +281,7 @@ def _candidate_dir(root: Path) -> tuple[Path, Path, Path]: "adjudication": "synthetic decision", } }, + "staging_delivery": _staging_delivery(), "identity": { "spine": { "pin_verified": True, @@ -412,12 +431,52 @@ def test_assembler_stages_a_contract_valid_dense_release( "--no-latest" in summary["publish_command"] and "populace-uk-private" in summary["publish_command"] ) + # The national assembler's rule: the run's staging receipt rides into the + # build manifest unchanged, where publication reads it. + build_manifest = json.loads((release_dir / "build_manifest.json").read_text()) + assert build_manifest["staging"] == _staging_delivery() # The contract re-validates the finished directory, and so does the pre-flight. dc.validate_release_dir(release_dir) preflight = _load("preflight_uk_local_release_candidate") assert preflight.main(["--release-dir", str(release_dir)]) == 0 +@pytest.mark.parametrize( + ("staging_delivery", "message"), + [ + (None, "missing valid staging-delivery evidence"), + ({"contract_version": 2, "enabled": True}, "invalid staging-delivery evidence"), + ], +) +def test_assembler_refuses_a_run_without_staging_evidence( + tmp_path: Path, monkeypatch, staging_delivery, message +) -> None: + monkeypatch.setenv("MICROCOSM_UK_TERMINAL_GATE_SIGNING_KEY", KEY) + assembler = _load("assemble_uk_dense_release_dir") + candidate, spine, incumbent = _candidate_dir(tmp_path) + manifest_path = candidate / "rowwise_candidate_manifest.json" + manifest = json.loads(manifest_path.read_text()) + if staging_delivery is None: + del manifest["staging_delivery"] + else: + manifest["staging_delivery"] = staging_delivery + manifest_path.write_text(json.dumps(manifest)) + with pytest.raises(SystemExit, match=message): + assembler.main( + [ + "--candidate-dir", + str(candidate), + "--spine-h5", + str(spine), + "--incumbent-manifest", + str(incumbent), + "--out-dir", + str(tmp_path / "releases"), + ] + ) + assert not (tmp_path / "releases").exists() + + def test_assembler_refuses_a_dev_posture_run(tmp_path: Path, monkeypatch) -> None: monkeypatch.setenv("MICROCOSM_UK_TERMINAL_GATE_SIGNING_KEY", KEY) assembler = _load("assemble_uk_dense_release_dir") diff --git a/packages/microcosm-build/tests/test_uk_local_rowwise.py b/packages/microcosm-build/tests/test_uk_local_rowwise.py index e595249fc6..673149cc9f 100644 --- a/packages/microcosm-build/tests/test_uk_local_rowwise.py +++ b/packages/microcosm-build/tests/test_uk_local_rowwise.py @@ -1301,6 +1301,48 @@ def restore(full): assert result.selected_support.tolist() == [0, 2] +def test_progress_events_reach_the_event_sink_beside_the_line_sink(): + frame = _clone_frame() + metrics = pd.DataFrame({"households": [1.0, 1.0, 1.0]}, index=[101, 102, 103]) + problem = build_uk_rowwise_local_matrix( + metrics, + _assigned(), + pd.DataFrame({"code": ["E001", "S001"], "households": [2.0, 1.0]}), + ) + lines: list[str] = [] + events: list[dict[str, object]] = [] + common = dict(bound_families=["census_households/constituency"], epochs=2, seed=7) + dense = solve_uk_rowwise_weights_under_doctrine( + frame, problem, progress=lines.append, progress_events=events.append, **common + ) + epochs = [event for event in events if event.get("kind") == "calibration_epoch"] + # One raw dict per epoch, untagged on a dense solve; the line sink still + # reports the last epoch only. + assert [event["epoch"] for event in epochs] == [1, 2] + assert all(event["epochs"] == 2 and "phase" not in event for event in epochs) + assert len(lines) == 1 and "dense solve: epoch 2/2" in lines[0] + + only_events: list[dict[str, object]] = [] + sized = solve_uk_rowwise_weights_under_doctrine( + frame, + problem, + dataset_households=2, + progress_events=only_events.append, + **common, + ) + phases = { + str(event.get("phase")) + for event in only_events + if event.get("kind") == "calibration_epoch" + } + # The size machinery tags its own phases; the dense solve carries none. + assert {"size_search", "size_refit", "None"} <= phases + kinds = {str(event.get("kind")) for event in only_events} + assert {"calibration_epoch", "budget_probe", "budget_search_done"} <= kinds + # The event sink changes nothing about the solve itself. + np.testing.assert_array_equal(sized.dense_reference.weights, dense.weights) + + def test_size_solve_keeps_its_dense_reference_and_selection_seed_moves_only_the_draw(): frame = _clone_frame() metrics = pd.DataFrame({"households": [1.0, 1.0, 1.0]}, index=[101, 102, 103]) diff --git a/packages/microcosm-build/tests/test_uk_rowwise_candidate.py b/packages/microcosm-build/tests/test_uk_rowwise_candidate.py index c6fc30b472..0ab82ea147 100644 --- a/packages/microcosm-build/tests/test_uk_rowwise_candidate.py +++ b/packages/microcosm-build/tests/test_uk_rowwise_candidate.py @@ -368,6 +368,8 @@ def _mandatory_input_flags(input_h5: Path, ladder_path: Path) -> list[str]: "1" * 64, "--ledger-manifest-sha256", "2" * 64, + # Tests never reach the Hub: telemetry and the staged bundle stay local. + "--staging-local-only", ] @@ -2555,3 +2557,545 @@ def test_size_candidate_checkpoints_before_the_draw_and_resumes_from_it( == checkpoint["provenance"]["build_id"] ) assert "written_at" not in resumed_receipt + + +# --------------------------------------------------------------------------- +# Staging: telemetry to runs// and the staged dataset bundle. + + +def _load_tool(name: str): + root = Path(__file__).resolve().parents[3] + spec = importlib.util.spec_from_file_location(name, root / "tools" / f"{name}.py") + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +class _FakeHub: + """One fake Hub serving the telemetry repo and the private dataset repo.""" + + def __init__(self, *, fail_commit: bool = False) -> None: + self.files: dict[tuple[str, str], bytes] = {} + self.uploads: list[tuple[str, str]] = [] + self.commits: list[dict[str, object]] = [] + self.fail_commit = fail_commit + self.sha = "a" * 40 + + def paths(self, repo_id: str) -> list[str]: + return sorted(path for repo, path in self.files if repo == repo_id) + + def upload_file(self, *, path_or_fileobj, path_in_repo, repo_id, repo_type): + assert repo_type == "dataset" + self.files[(repo_id, path_in_repo)] = Path(path_or_fileobj).read_bytes() + self.uploads.append((repo_id, path_in_repo)) + + def hf_hub_download(self, *, repo_id, filename, repo_type, **kwargs): + assert repo_type == "dataset" + if (repo_id, filename) not in self.files: + raise FileNotFoundError(filename) + destination = ( + Path(kwargs.get("local_dir") or _FAKE_HUB_CACHE) / repo_id / filename + ) + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(self.files[(repo_id, filename)]) + return str(destination) + + def file_exists(self, *, repo_id, filename, repo_type): + assert repo_type == "dataset" + return (repo_id, filename) in self.files + + def repo_info(self, *, repo_id, repo_type): + assert repo_type == "dataset" + return SimpleNamespace(sha=self.sha) + + def create_commit( + self, *, repo_id, operations, commit_message, repo_type, parent_commit + ): + assert repo_type == "dataset" + if self.fail_commit: + raise RuntimeError("403 Forbidden token=do-not-record") + assert parent_commit == self.sha + for operation in operations: + self.files[(repo_id, operation.path_in_repo)] = Path( + operation.path_or_fileobj + ).read_bytes() + self.commits.append( + { + "repo_id": repo_id, + "message": commit_message, + "paths": sorted(op.path_in_repo for op in operations), + } + ) + self.sha = "b" * 40 + return SimpleNamespace(oid=self.sha) + + +_FAKE_HUB_CACHE = "/tmp/microcosm-fake-hub-cache" + + +def _staging_run_setup(builder, monkeypatch, tmp_path, *, remote: bool = False): + """Fixture inputs plus the flag list; ``remote`` drops the local-only switch.""" + + pytest.importorskip("tables") + pytest.importorskip("h5py") + input_h5 = tmp_path / "staging.h5" + ladder_path = tmp_path / "ladder.npz" + _write_staging_h5(input_h5, households_per_region=52) + ladder = _write_ladder(ladder_path) + flags = _configure_households_only_inputs( + builder, monkeypatch, input_h5=input_h5, ladder_path=ladder_path + ) + import microcosm.build.uk_runtime.battery_bindings as battery_bindings + + monkeypatch.setattr( + battery_bindings, + "_local_area_roster", + lambda _resource, levels: { + "constituency": tuple(sorted(set(ladder.constituency_code))), + "local_authority": tuple(sorted(set(ladder.local_authority_code))), + }, + ) + if remote: + flags = [flag for flag in flags if flag != "--staging-local-only"] + return input_h5, ladder_path, flags + + +def _build_args(input_h5, ladder_path, flags, out, *extra): + return [ + "--input-h5", + str(input_h5), + "--ladder", + str(ladder_path), + *flags, + "--out", + str(out), + "--n-clones", + "2", + "--seed", + "7", + "--epochs", + "2", + "--skip-holdout", + *extra, + ] + + +def _single_run_id(out: Path) -> str: + runs = sorted(path.name for path in (out / "staging" / "runs").iterdir()) + assert len(runs) == 1, runs + return runs[0] + + +def test_candidate_build_stages_telemetry_locally_and_inventories_the_bundle( + monkeypatch, tmp_path, capsys +): + from microcosm.build.staging_v2 import validate_v2_bundle + + builder = _load_builder_module() + input_h5, ladder_path, flags = _staging_run_setup(builder, monkeypatch, tmp_path) + out = tmp_path / "candidate" + status = builder.main(_build_args(input_h5, ladder_path, flags, out)) + assert status == 0 + captured = capsys.readouterr() + manifest = json.loads((out / builder.MANIFEST_FILENAME).read_text()) + # The stdout manifest is the on-disk manifest, evidence blocks included. + assert json.loads(captured.out)["staged_dataset"] == manifest["staged_dataset"] + assert "staged dataset: skipped (local_only)" in captured.err + + run_id = _single_run_id(out) + rows = load_spool_rows(out / "logbook-spool") + assert rows[0].build_id == run_id + bundle = validate_v2_bundle(out / "staging", run_id) + run_manifest = bundle["run_manifest"] + assert run_manifest["status"] == "completed" + assert run_manifest["run_kind"] == "calibration" + assert run_manifest["operation_id"] == "uk_rowwise_candidate" + assert run_manifest["pipeline"]["id"] == "uk-local-candidate" + assert run_manifest["non_release"] is True + assert run_manifest["sample"] == {"mode": "full"} + assert run_manifest["delivery"]["mode"] == "local_only" + assert run_manifest["delivery"]["upload_attempts"] == 0 + assert {a["logical_name"] for a in run_manifest["artifacts"]} == { + "fit_summary", + "staged_dataset", + } + run_dir = out / "staging" / "runs" / run_id + fit_summary = json.loads((run_dir / "artifacts" / "fit_summary.json").read_text()) + assert fit_summary["run_id"] == run_id + assert set(fit_summary["gates"]) == set(builder.UK_LOCAL_GATE_SCOPE) + assert fit_summary["loss"]["final"] == manifest["solve"]["final_loss"] + assert set(fit_summary["fit_by_family"]["local"]) == {"census_households"} + staged_artifact = json.loads( + (run_dir / "artifacts" / "staged_dataset.json").read_text() + ) + assert staged_artifact == manifest["staged_dataset"] + + # Every phase reports started then completed, in build order; the + # calibration progress keeps the last epoch only (thinning at 2 epochs). + events = bundle["events"] + completed = [e["stage_id"] for e in events if e["status"] == "completed"] + assert completed == [ + "input_pinning", + "target_compilation", + "cloning", + "surface_resolution", + "calibration", + "gate_battery", + "holdout", + "output_bundle", + "dataset_staging", + "complete", + ] + for stage in ("target_compilation", "cloning", "calibration", "gate_battery"): + transitions = [e["status"] for e in events if e["stage_id"] == stage] + assert transitions == ["started", "completed"], stage + calibration_rows = json.loads((run_dir / "calibration_progress.json").read_text())[ + "events" + ] + assert [ + (row["epoch"], row["epochs"], row["phase"]) for row in calibration_rows + ] == [(2, 2, None)] + calibration_done = next( + e + for e in events + if e["stage_id"] == "calibration" and e["status"] == "completed" + ) + assert calibration_done["details"]["final_loss"] == manifest["solve"]["final_loss"] + assert calibration_done["details"]["size_checkpoint"] is None + + # The manifest carries both receipts; the bundle inventory is the outputs. + assert manifest["staging_delivery"]["mode"] == "local_only" + assert manifest["staging_delivery"]["run_id"] == run_id + staged = manifest["staged_dataset"] + assert staged["mode"] == "local_only" and staged["status"] == "skipped" + assert staged["prefix"] == f"staged/{run_id}" + assert set(staged["files"]) == { + Path(entry["path"]).name for entry in manifest["outputs"].values() + } + for entry in manifest["outputs"].values(): + assert staged["files"][Path(entry["path"]).name]["sha256"] == entry["sha256"] + assert (out / "sha256sums.txt").is_file() + inventory = json.loads((out / "staged_manifest.json").read_text()) + assert inventory["run_id"] == run_id and inventory["files"] == staged["files"] + assert inventory["summary"]["releasable"] is True + assert inventory["telemetry"] == { + "repository": None, + "prefix": f"runs/{run_id}", + "mode": "local_only", + } + # The sidecars are evidence about the outputs, never outputs themselves. + assert "sha256sums" not in manifest["outputs"] + assert "staged_manifest" not in manifest["outputs"] + + +def test_size_candidate_stages_the_search_and_refit_phases(monkeypatch, tmp_path): + from microcosm.build.staging_v2 import validate_v2_bundle + + builder = _load_builder_module() + input_h5, ladder_path, flags = _staging_run_setup(builder, monkeypatch, tmp_path) + out = tmp_path / "k300" + status = builder.main( + _build_args( + input_h5, + ladder_path, + flags, + out, + "--dataset-households", + "300", + "--selection-seed", + "11", + ) + ) + assert status in (0, 1) + run_id = _single_run_id(out) + bundle = validate_v2_bundle(out / "staging", run_id) + assert bundle["run_manifest"]["status"] == "completed" + rows = json.loads( + (out / "staging" / "runs" / run_id / "calibration_progress.json").read_text() + )["events"] + phases = [row["phase"] for row in rows] + assert phases[0] is None and "size_search" in phases and phases[-1] == "size_refit" + probe_rows = [row for row in rows if row["budget_search"] == 1] + assert probe_rows and all(row["l0_lambda"] is not None for row in probe_rows) + assert all(row["epoch"] == 2 for row in rows) + calibration_done = next( + e + for e in bundle["events"] + if e["stage_id"] == "calibration" and e["status"] == "completed" + ) + assert calibration_done["details"]["size_checkpoint"] == "written" + assert calibration_done["details"]["realized_households"] == 300 + manifest = json.loads((out / builder.MANIFEST_FILENAME).read_text()) + assert manifest["releasable"] is False + assert manifest["staged_dataset"]["status"] == "skipped" + assert "dataset_size_selection.csv" in manifest["staged_dataset"]["files"] + fit_summary = json.loads( + ( + out / "staging" / "runs" / run_id / "artifacts" / "fit_summary.json" + ).read_text() + ) + assert fit_summary["dataset_size"]["requested_households"] == 300 + assert "pool_row_indices" not in fit_summary["dataset_size"] + assert fit_summary["releasable"] is False + + +def test_no_staging_records_both_opt_outs(monkeypatch, tmp_path): + builder = _load_builder_module() + input_h5, ladder_path, flags = _staging_run_setup( + builder, monkeypatch, tmp_path, remote=True + ) + out = tmp_path / "quiet" + status = builder.main( + _build_args(input_h5, ladder_path, flags, out, "--no-staging") + ) + assert status == 0 + assert not (out / "staging").exists() + assert not (out / "sha256sums.txt").exists() + manifest = json.loads((out / builder.MANIFEST_FILENAME).read_text()) + assert manifest["staging_delivery"]["mode"] == "disabled" + assert manifest["staging_delivery"]["opt_out_reason"] == "--no-staging" + assert manifest["staged_dataset"] == { + "contract_version": 1, + "mode": "disabled", + "repository": None, + "prefix": None, + "run_id": None, + "revision": None, + "status": "skipped", + "error_code": None, + "opt_out_reason": "--no-staging", + "files": {}, + } + rows = load_spool_rows(out / "logbook-spool") + assert "dataset_stage_skipped" in rows[0].phases_reached + + +def test_remote_staging_uploads_telemetry_and_the_bundle_in_one_commit( + monkeypatch, tmp_path, capsys +): + builder = _load_builder_module() + input_h5, ladder_path, flags = _staging_run_setup( + builder, monkeypatch, tmp_path, remote=True + ) + hub = _FakeHub() + monkeypatch.setattr(builder, "_hub_api", lambda: hub) + monkeypatch.setattr(builder, "_hub_token", lambda: "hf_test_token") + out = tmp_path / "remote" + status = builder.main( + _build_args( + input_h5, ladder_path, flags, out, "--staging-upload-interval-seconds", "0" + ) + ) + assert status == 0 + err = capsys.readouterr().err + run_id = _single_run_id(out) + manifest = json.loads((out / builder.MANIFEST_FILENAME).read_text()) + + # Telemetry went to runs// of the staging repository, artifacts + # included, under the fixed prefix and nothing else. + telemetry_paths = hub.paths("policyengine/populace-uk-staging") + assert telemetry_paths == sorted( + f"runs/{run_id}/{name}" + for name in ( + "run_manifest.json", + "progress.json", + "events.ndjson", + "calibration_progress.json", + "artifacts/fit_summary.json", + "artifacts/staged_dataset.json", + ) + ) + delivery = manifest["staging_delivery"] + assert delivery["mode"] == "local_and_remote" + assert delivery["configured_repository"] == "policyengine/populace-uk-staging" + assert delivery["upload_successes"] == delivery["upload_attempts"] > 0 + remote_progress = json.loads( + hub.files[("policyengine/populace-uk-staging", f"runs/{run_id}/progress.json")] + ) + assert remote_progress["status"] == "completed" + + # The bundle went to staged// of the private repository in one + # commit: every output, the manifest as built, and the two sidecars. + assert len(hub.commits) == 1 + commit = hub.commits[0] + assert commit["repo_id"] == "policyengine/populace-uk-private" + expected = {Path(e["path"]).name for e in manifest["outputs"].values()} | { + builder.MANIFEST_FILENAME, + "staged_manifest.json", + "sha256sums.txt", + } + assert commit["paths"] == sorted(f"staged/{run_id}/{name}" for name in expected) + assert hub.paths("policyengine/populace-uk-private") == commit["paths"] + staged = manifest["staged_dataset"] + assert staged["status"] == "uploaded" + assert staged["repository"] == "policyengine/populace-uk-private" + assert staged["prefix"] == f"staged/{run_id}" + assert staged["revision"] == hub.sha + assert ( + f"staged dataset: uploaded at policyengine/populace-uk-private/staged/{run_id}" + in err + ) + remote_h5 = hub.files[ + ( + "policyengine/populace-uk-private", + f"staged/{run_id}/{Path(manifest['outputs']['dataset']['path']).name}", + ) + ] + assert ( + remote_h5 + == (out / Path(manifest["outputs"]["dataset"]["path"]).name).read_bytes() + ) + # The uploaded manifest is the one the bundle was built from; the local + # copy gained the two evidence blocks afterwards. + remote_manifest = json.loads( + hub.files[ + ( + "policyengine/populace-uk-private", + f"staged/{run_id}/{builder.MANIFEST_FILENAME}", + ) + ] + ) + assert ( + "staged_dataset" not in remote_manifest + and "staging_delivery" not in remote_manifest + ) + assert remote_manifest["outputs"] == manifest["outputs"] + rows = load_spool_rows(out / "logbook-spool") + assert "dataset_staged" in rows[0].phases_reached + assert rows[0].disposition == "iterating" + + # Re-staging the same directory is a no-op on the output digests. + stager = _load_tool("stage_uk_rowwise_candidate") + monkeypatch.setattr(stager, "_hub_api", lambda: hub) + assert stager.main(["--run-dir", str(out)]) == 0 + restaged = json.loads((out / builder.MANIFEST_FILENAME).read_text())[ + "staged_dataset" + ] + assert restaged["status"] == "already_staged" + assert restaged["run_id"] == run_id and restaged["revision"] == hub.sha + assert len(hub.commits) == 1 + + # Consumers fetch by run id and get digest-verified local files. + fetcher = _load_tool("fetch_uk_staged_dataset") + monkeypatch.setattr(fetcher, "_hub_api", lambda: hub) + dest = tmp_path / "fetched" + capsys.readouterr() + assert fetcher.main(["--run-id", run_id, "--dest", str(dest)]) == 0 + listed = capsys.readouterr().out.splitlines() + assert str(dest / Path(manifest["outputs"]["dataset"]["path"]).name) in listed + assert (dest / "sha256sums.txt").is_file() + assert ( + dest / Path(manifest["outputs"]["dataset"]["path"]).name + ).read_bytes() == remote_h5 + + +def test_remote_staging_failure_is_recorded_and_the_build_still_succeeds( + monkeypatch, tmp_path, capsys +): + builder = _load_builder_module() + input_h5, ladder_path, flags = _staging_run_setup( + builder, monkeypatch, tmp_path, remote=True + ) + hub = _FakeHub(fail_commit=True) + monkeypatch.setattr(builder, "_hub_api", lambda: hub) + monkeypatch.setattr(builder, "_hub_token", lambda: "hf_test_token") + out = tmp_path / "failed-upload" + status = builder.main(_build_args(input_h5, ladder_path, flags, out)) + assert status == 0 + err = capsys.readouterr().err + assert "staged dataset upload failed" in err and "do-not-record" not in err + manifest = json.loads((out / builder.MANIFEST_FILENAME).read_text()) + staged = manifest["staged_dataset"] + assert staged["status"] == "failed" and staged["error_code"] == "UPLOAD_FAILED" + assert staged["revision"] is None and staged["files"] + assert "do-not-record" not in json.dumps(manifest) + assert hub.paths("policyengine/populace-uk-private") == [] + # Telemetry still completed and recorded the outcome. + run_id = _single_run_id(out) + progress = json.loads( + hub.files[("policyengine/populace-uk-staging", f"runs/{run_id}/progress.json")] + ) + assert progress["status"] == "completed" + events = [ + json.loads(line) + for line in hub.files[ + ("policyengine/populace-uk-staging", f"runs/{run_id}/events.ndjson") + ] + .decode() + .splitlines() + if line + ] + done = next( + e + for e in events + if e["stage_id"] == "dataset_staging" and e["status"] == "completed" + ) + assert done["details"]["status"] == "failed" + assert done["details"]["error_code"] == "UPLOAD_FAILED" + rows = load_spool_rows(out / "logbook-spool") + assert rows[0].disposition == "iterating" + assert "dataset_stage_failed" in rows[0].phases_reached + # The sidecars are in place for a later re-stage. + assert (out / "sha256sums.txt").is_file() and ( + out / "staged_manifest.json" + ).is_file() + + +def test_no_staged_dataset_keeps_telemetry_remote_and_the_bundle_local( + monkeypatch, tmp_path +): + builder = _load_builder_module() + input_h5, ladder_path, flags = _staging_run_setup( + builder, monkeypatch, tmp_path, remote=True + ) + hub = _FakeHub() + monkeypatch.setattr(builder, "_hub_api", lambda: hub) + monkeypatch.setattr(builder, "_hub_token", lambda: "hf_test_token") + out = tmp_path / "telemetry-only" + status = builder.main( + _build_args(input_h5, ladder_path, flags, out, "--no-staged-dataset") + ) + assert status == 0 + assert hub.paths("policyengine/populace-uk-private") == [] + assert hub.commits == [] + assert hub.paths("policyengine/populace-uk-staging") + manifest = json.loads((out / builder.MANIFEST_FILENAME).read_text()) + assert manifest["staging_delivery"]["mode"] == "local_and_remote" + assert manifest["staged_dataset"]["mode"] == "disabled" + assert manifest["staged_dataset"]["opt_out_reason"] == "--no-staged-dataset" + assert not (out / "sha256sums.txt").exists() + + +def test_remote_dataset_staging_is_refused_up_front_without_credential_or_repo( + monkeypatch, tmp_path, capsys +): + builder = _load_builder_module() + input_h5, ladder_path, flags = _staging_run_setup( + builder, monkeypatch, tmp_path, remote=True + ) + out = tmp_path / "refused" + monkeypatch.setattr(builder, "_hub_token", lambda: None) + with pytest.raises(ValueError, match="write credential"): + builder.main(_build_args(input_h5, ladder_path, flags, out)) + assert not out.exists() + + class Unreachable: + def repo_info(self, **kwargs): + raise RuntimeError("503 token=do-not-record") + + monkeypatch.setattr(builder, "_hub_token", lambda: "hf_test_token") + monkeypatch.setattr(builder, "_hub_api", lambda: Unreachable()) + with pytest.raises(ValueError, match="cannot reach") as info: + builder.main(_build_args(input_h5, ladder_path, flags, out)) + assert "do-not-record" not in str(info.value) + assert not out.exists() + + # A dry run plans without staging, so it needs neither credential nor repo. + monkeypatch.setattr(builder, "_hub_token", lambda: None) + assert ( + builder.main(_build_args(input_h5, ladder_path, flags, out, "--dry-run")) == 0 + ) + plan = json.loads(capsys.readouterr().out) + assert plan["parameters"]["dataset_households"] is None + assert not out.exists() diff --git a/packages/microcosm-build/tests/test_uk_size_evaluation.py b/packages/microcosm-build/tests/test_uk_size_evaluation.py index 5927842132..0a4d7bdc97 100644 --- a/packages/microcosm-build/tests/test_uk_size_evaluation.py +++ b/packages/microcosm-build/tests/test_uk_size_evaluation.py @@ -571,6 +571,57 @@ def test_area_support_and_gate_tables(tmp_path: Path) -> None: assert gate_table(run)[-1]["status"] == "absent" +def test_run_acceptance_ignores_staging_evidence_and_sidecars(tmp_path: Path) -> None: + """The staging lane appends manifest blocks and two sidecars; acceptance is unmoved.""" + + path = _write_run_dir(tmp_path, "candidate", size_run=True) + before = run_acceptance( + load_run(path, label="candidate"), + expected_households=6, + expected_pool=30, + expected_epochs=10, + ) + manifest_path = path / "rowwise_candidate_manifest.json" + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + manifest["staging_delivery"] = { + "contract_version": 2, + "enabled": True, + "mode": "local_only", + "run_id": "uk-local-candidate-f100-s7-20260917T180000Z-0badcafe", + "configured_repository": None, + "upload_attempts": 0, + "upload_successes": 0, + "read_back": "not_requested", + "last_error_code": None, + "opt_out_reason": None, + } + manifest["staged_dataset"] = { + "contract_version": 1, + "mode": "local_only", + "repository": None, + "prefix": "staged/uk-local-candidate-f100-s7-20260917T180000Z-0badcafe", + "run_id": "uk-local-candidate-f100-s7-20260917T180000Z-0badcafe", + "revision": None, + "status": "skipped", + "error_code": None, + "opt_out_reason": None, + "files": {"microcosm_uk_2025_local.h5": {"sha256": "0" * 64, "bytes": 1}}, + } + manifest_path.write_text(json.dumps(manifest), encoding="utf-8") + (path / "sha256sums.txt").write_text(f"{'0' * 64} microcosm_uk_2025_local.h5\n") + (path / "staged_manifest.json").write_text("{}") + after = run_acceptance( + load_run(path, label="candidate"), + expected_households=6, + expected_pool=30, + expected_epochs=10, + ) + assert after["passed"] == before["passed"] + assert [(row["id"], row["status"]) for row in after["checks"]] == [ + (row["id"], row["status"]) for row in before["checks"] + ] + + def test_run_acceptance_statuses_and_bad_digest(tmp_path: Path) -> None: path = _write_run_dir(tmp_path, "candidate", size_run=True) run = load_run(path, label="candidate") diff --git a/tools/assemble_uk_dense_release_dir.py b/tools/assemble_uk_dense_release_dir.py index 1049a41db6..7d629644a4 100644 --- a/tools/assemble_uk_dense_release_dir.py +++ b/tools/assemble_uk_dense_release_dir.py @@ -30,6 +30,10 @@ from datetime import UTC, datetime from pathlib import Path +from microcosm.build.staging_v2 import ( + StagingContractError, + validate_staging_delivery, +) from microcosm.build.uk_runtime.release_identity import UK_DENSE_RELEASE_ID from microcosm.data.contract import ( _check_uk_incumbent_surface_evaluation, @@ -192,6 +196,20 @@ def _assemble(args: argparse.Namespace) -> dict[str, object]: ) outputs = _mapping(manifest.get("outputs"), "manifest.outputs") identity = _mapping(manifest.get("identity"), "manifest.identity") + # The national assembler's rule: a release carries the staging telemetry + # receipt of the run it came from, and publication refuses a release + # whose run intended to stage and delivered nothing. + raw_staging_delivery = manifest.get("staging_delivery") + if not isinstance(raw_staging_delivery, Mapping): + raise SystemExit( + "error: build record is missing valid staging-delivery evidence" + ) + try: + validate_staging_delivery(raw_staging_delivery) + except StagingContractError as error: + raise SystemExit( + f"error: invalid staging-delivery evidence: {error}" + ) from error def output_path(key: str) -> Path: entry = _mapping(outputs.get(key), f"manifest.outputs.{key}") @@ -412,6 +430,11 @@ def _stage_and_finalize( shutil.copyfile(source, release_dir / name) identity = _mapping(manifest.get("identity"), "identity") parameters = _mapping(manifest.get("parameters"), "parameters") + # Validated once more here so the copy into build_manifest.json is the + # normalized version 2 object, whatever the caller handed over. + staging_delivery = validate_staging_delivery( + _mapping(manifest.get("staging_delivery"), "staging_delivery") + ) solve = _mapping(manifest.get("solve"), "solve") fit = _mapping(manifest.get("fit"), "fit") weights = _mapping(manifest.get("weights"), "weights") @@ -589,6 +612,7 @@ def _stage_and_finalize( "attempt_id": attempt_id, "cut_tag": cut_tag, "created_at": created_at, + "staging": dict(staging_delivery), } _write_json(release_dir / "build_manifest.json", build_manifest) diff --git a/tools/build_uk_rowwise_candidate.py b/tools/build_uk_rowwise_candidate.py index 38aa5dbd96..f7aac25ce8 100644 --- a/tools/build_uk_rowwise_candidate.py +++ b/tools/build_uk_rowwise_candidate.py @@ -22,8 +22,9 @@ import tempfile import time import uuid -from collections.abc import Mapping +from collections.abc import Callable, Mapping from datetime import UTC, date, datetime +from importlib import metadata from pathlib import Path from typing import Any @@ -54,6 +55,25 @@ sha256_argument, write_error_receipt, ) +from microcosm.build.staging_cli import ( + add_staged_dataset_arguments, + add_staging_arguments, + validate_staged_dataset_arguments, + validate_staging_arguments, +) +from microcosm.build.staging_dataset import ( + StagedDatasetBundle, + disabled_staged_dataset, + local_only_staged_dataset, + stage_bundle, + write_sidecars, +) +from microcosm.build.staging_storage import HuggingFaceDatasetStorage +from microcosm.build.staging_v2 import ( + StagingContractError, + StagingTelemetryV2, + disabled_staging_delivery, +) from microcosm.build.target_materialization import resolve_target_measures from microcosm.build.uk_runtime import ( UK_GATE_REGISTRY, @@ -122,6 +142,11 @@ UK_SAMPLE_SEED_DEFAULT, sample_uk_spine_frame, ) +from microcosm.build.uk_runtime.staging import ( + UK_STAGED_DATASET_PREFIX, + UK_STAGED_DATASET_REPOSITORY, + UK_STAGING_REPOSITORY, +) from microcosm.calibrate import TargetRegistry, TargetSpec from microcosm.frame import Frame, MassChangeRecord @@ -149,6 +174,20 @@ _BUDGET_ITERS = 10 _UK_CANDIDATE_PIPELINE = "uk-local-candidate" _LOCAL_GATE_POLICY_SUFFIX = "local_candidate" +_STAGING_OPERATION_ID = "uk_rowwise_candidate" +# Staging telemetry keeps one row per forwarded epoch in +# calibration_progress.json and one event in events.ndjson, both under the +# contract's 5 MiB remote cap. A size run at 2,000 epochs solves the dense +# pool, up to ten full-length L0 probes and the refit: about 24,000 epochs, +# which would breach the cap mid-run. Forwarding every tenth epoch and the +# last epoch of each phase keeps the loss curve and stays near 0.8 MB. +_STAGING_EPOCH_EVERY = 10 +_STAGED_DATASET_PHASES = { + "uploaded": "dataset_staged", + "already_staged": "dataset_staged", + "failed": "dataset_stage_failed", + "skipped": "dataset_stage_skipped", +} _REPOSITORY = Path(__file__).resolve().parents[1] _PAST_CAP_COUNT_KEYS = ( "n_targets", @@ -804,7 +843,12 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: "POPULACE_LOGBOOK_PREV_ROW_DIGEST is used, then genesis null." ), ) - return parser.parse_args(argv) + add_staging_arguments(parser, repository=UK_STAGING_REPOSITORY) + add_staged_dataset_arguments(parser, repository=UK_STAGED_DATASET_REPOSITORY) + args = parser.parse_args(argv) + validate_staging_arguments(parser, args) + validate_staged_dataset_arguments(parser, args) + return args def main(argv: list[str] | None = None) -> int: @@ -812,6 +856,7 @@ def main(argv: list[str] | None = None) -> int: args = _parse_args(argv) _validate_cli_args(args) + _preflight_staged_dataset(args) if args.candidate_clone_counts is not None and not args.dry_run: raise ValueError("--candidate-clone-counts is valid only with --dry-run.") if _CONSERVE_MASS: @@ -844,25 +889,35 @@ def main(argv: list[str] | None = None) -> int: } }, ) - return _run_candidate( - args, - attempt={ - "state": state, - "started_at": started_at, - "started_ts": started_ts, - "code_pin": "unresolved-local-git-code-pin", - # Logbook chain configuration is validated before any terminal - # work: a malformed or conflicting head refuses the run with no - # row and no side effects (#666 adversarial-review finding). - "predecessor": resolve_predecessor(args.logbook_prev_row_digest), - }, - ) + # Logbook chain configuration is validated before any terminal work: a + # malformed or conflicting head refuses the run with no row and no side + # effects (#666 adversarial-review finding). + predecessor = resolve_predecessor(args.logbook_prev_row_digest) + # Staging telemetry opens with the attempt, as in the spine builder, so an + # early refusal still leaves a failed run under runs//. + telemetry = _create_staging_telemetry(args, state=state) + try: + return _run_candidate( + args, + attempt={ + "state": state, + "started_at": started_at, + "started_ts": started_ts, + "code_pin": "unresolved-local-git-code-pin", + "predecessor": predecessor, + }, + telemetry=telemetry, + ) + except BaseException as error: + _fail_staging_telemetry(telemetry, error) + raise def _run_candidate( args: argparse.Namespace, *, attempt: dict[str, object] | None, + telemetry: StagingTelemetryV2 | None = None, ) -> int: """Build the candidate, recording every non-dry terminal outcome. @@ -888,6 +943,13 @@ def _run_candidate( "dataset": _pin_from_artifact(input_artifact), "ladder": _pin_from_artifact(ladder_artifact), } + _stage( + telemetry, + "input_pinning", + "completed", + dataset_sha256=input_artifact["sha256"], + ladder_sha256=ladder_artifact["sha256"], + ) state: AttemptState | None = None if attempt is not None: unpacked_state = attempt["state"] @@ -918,6 +980,10 @@ def _run_candidate( seed=args.sample_seed, ) args._sampling_receipt = sampling + if telemetry is not None and args.sample_fraction == 1.0: + # The contract's only sampling statement is "full"; a rung below + # f100 stages a null sample, as the spine builder does. + telemetry.set_sample({"mode": "full"}) source_year = _source_year( args.source_year, time_period=uk_time_period(national_frame), @@ -943,6 +1009,7 @@ def _run_candidate( _refuse_stale_size_checkpoint(args, out_dir) ladder = load_uk_oa_ladder(ladder_path) target_provenance = ladder_target_provenance(ladder) + _stage(telemetry, "target_compilation", "started") joint_inputs = _load_joint_target_inputs(args) facts = getattr(joint_inputs.get("artifact"), "facts", None) if facts is None: @@ -961,6 +1028,16 @@ def _run_candidate( joint_inputs["household_dispersion"] = ladder_vs_chronicle_household_dispersion( ladder, joint_inputs["local_registry"].specs ) + _stage( + telemetry, + "target_compilation", + "completed", + local_target_count=len(joint_inputs["local_registry"].specs), + national_target_count=len(joint_inputs["national_registry"].specs), + census_household_uprating_applied=bool( + joint_inputs["census_household_uprating"].get("applied") + ), + ) if args.release_candidate and not joint_inputs["census_household_uprating"].get( "applied" ): @@ -984,6 +1061,7 @@ def _run_candidate( ) print("cloning through the ladder route...", file=sys.stderr, flush=True) + _stage(telemetry, "cloning", "started", clone_count=int(args.n_clones)) assignment = _clone_with_ladder_binding( national_frame, ladder, @@ -1001,6 +1079,13 @@ def _run_candidate( raise ValueError( "--dataset-households exceeds the cloned pool; selection never clamps the request." ) + _stage( + telemetry, + "cloning", + "completed", + clone_count=int(args.n_clones), + pool_rows=int(clone.frame.n("household")), + ) if state is not None: append_phase(state, "cloned") @@ -1025,6 +1110,7 @@ def _run_candidate( print(_json_text(plan), end="") return 0 + _stage(telemetry, "surface_resolution", "started") if args.households_only: print( "binding Chronicle census household targets...", @@ -1087,6 +1173,14 @@ def _run_candidate( args._bound_families = tuple(bound_families) args._joint_inputs_receipt = joint_inputs args._measure_resolution = dict(measure_resolution) + _stage( + telemetry, + "surface_resolution", + "completed", + target_count=int(problem.matrix.shape[0]), + bound_family_count=len(bound_families), + engine_blocks=int(args.engine_blocks), + ) if state is not None: append_phase(state, "targets_bound") @@ -1150,6 +1244,16 @@ def _run_candidate( file=sys.stderr, flush=True, ) + _stage( + telemetry, + "calibration", + "started", + target_count=int(problem.matrix.shape[0]), + pool_rows=int(problem.matrix.shape[1]), + dataset_households=args.dataset_households, + epochs=int(args.epochs), + resumed_from_checkpoint=resume_checkpoint is not None, + ) solve = solve_uk_rowwise_weights_under_doctrine( solve_frame, problem, @@ -1173,6 +1277,11 @@ def _run_candidate( checkpoint_identity=checkpoint_identity, checkpoint_provenance={"code_pin": code_pin, "build_id": state.build_id}, progress=_stderr_progress, + progress_events=( + None + if telemetry is None + else _thinned_epochs(telemetry.calibration_progress) + ), ) _validate_solve_result(solve, problem=problem) if solve.size_receipt is not None and solve.size_receipt.get("checkpoint"): @@ -1187,6 +1296,15 @@ def _run_candidate( elif "resumed_from" in checkpoint: append_phase(state, "size_selection_resumed") append_phase(state, "solved") + _stage( + telemetry, + "calibration", + "completed", + final_loss=float(solve.final_loss), + n_nonzero=int(solve.n_nonzero), + realized_households=int(solve.frame.n("household")), + size_checkpoint=_size_checkpoint_state(solve), + ) # The kernel minted the calibration mass record inside calibrate() (the # CALIBRATED kind transition is enforced there too); the record names @@ -1211,6 +1329,7 @@ def _run_candidate( None if args.households_only else joint_inputs["national_registry"] ), ) + _stage(telemetry, "gate_battery", "started") try: gate_report, candidate_gate = _run_local_gate_battery( frame=solve.frame, @@ -1272,7 +1391,16 @@ def _run_candidate( append_phase( state, "candidate_gated" if not blocked_failures else "candidate_blocked" ) + _stage( + telemetry, + "gate_battery", + "completed", + gate_statuses=_gate_statuses(gate_report), + blocking_failure_count=len(blocked_failures), + diagnostic_failure_count=len(diagnostic_failures), + ) + _stage(telemetry, "holdout", "started", skipped=bool(args.skip_holdout)) if args.skip_holdout: rotated_holdout = {"skipped": True} else: @@ -1296,6 +1424,7 @@ def _run_candidate( baseline_pi_floor=args.baseline_pi_floor, ) args._rotated_holdout = rotated_holdout + _stage(telemetry, "holdout", "completed", skipped=bool(args.skip_holdout)) candidate = dataclasses.replace( clone, @@ -1321,6 +1450,7 @@ def _run_candidate( ladder_artifact=ladder_artifact, ) + _stage(telemetry, "output_bundle", "started") manifest = _write_output_bundle( args, candidate=candidate, @@ -1340,7 +1470,34 @@ def _run_candidate( target_provenance=target_provenance, cross_grain=cross_grain, ) + _stage( + telemetry, + "output_bundle", + "completed", + output_bytes={ + key: int(entry["bytes"]) for key, entry in manifest["outputs"].items() + }, + ) append_phase(state, "published") + # The staged dataset and the telemetry receipt are evidence about the + # published bundle, so they are appended to the manifest after it is + # on disk (the national build record is rewritten the same way); the + # copy inside the staged bundle predates them and staged_manifest.json + # describes the remote side. + staged_dataset = _stage_dataset( + args, + manifest=manifest, + output_paths=output_paths, + run_id=state.build_id if telemetry is None else telemetry.run_id, + telemetry=telemetry, + ) + append_phase(state, _STAGED_DATASET_PHASES[staged_dataset["status"]]) + try: + _finalize_staging_telemetry(args, telemetry) + finally: + manifest["staging_delivery"] = _staging_delivery(telemetry) + manifest["staged_dataset"] = staged_dataset + _replace_manifest(output_paths["manifest"], manifest) state.artifact_location = local_artifact_reference( output_paths["dataset"], repository_hint=_REPOSITORY, @@ -1393,6 +1550,451 @@ def _run_candidate( raise +def _hub_api() -> Any: + """The Hub client used for telemetry and the staged dataset (test seam).""" + + from huggingface_hub import HfApi + + return HfApi() + + +def _hub_token() -> str | None: + """The ambient Hub credential, if any (test seam).""" + + from huggingface_hub import get_token + + return get_token() + + +def _staged_dataset_mode(args: argparse.Namespace) -> str: + if args.no_staging or args.no_staged_dataset: + return "disabled" + if args.staging_local_only: + return "local_only" + return "local_and_remote" + + +def _preflight_staged_dataset(args: argparse.Namespace) -> None: + """Refuse a remote dataset stage the run could not complete. + + The bundle upload is the last step of a multi-hour run, so the credential + and the repository are checked before the spine is read. Telemetry stays + best-effort with no pre-flight, as on the national command. + """ + + if args.dry_run or _staged_dataset_mode(args) != "local_and_remote": + return + repo_id = str(args.staged_dataset_repo_id).strip() + hint = "pass --staging-local-only or --no-staged-dataset to keep the bundle local" + if not _hub_token(): + raise ValueError( + f"remote dataset staging to {repo_id} needs a Hugging Face write " + f"credential (HF_TOKEN or `hf auth login`); {hint}." + ) + try: + _hub_api().repo_info(repo_id=repo_id, repo_type="dataset") + except Exception as error: + raise ValueError( + f"remote dataset staging cannot reach {repo_id} " + f"({type(error).__name__}); {hint}." + ) from error + + +def _create_staging_telemetry( + args: argparse.Namespace, *, state: AttemptState +) -> StagingTelemetryV2 | None: + if args.no_staging: + return None + local_only = bool(args.staging_local_only) + out_dir = args.out.expanduser().resolve() + return StagingTelemetryV2( + run_id=args.staging_run_id or state.build_id, + country_code="GB", + operation_id=_STAGING_OPERATION_ID, + pipeline_id=_UK_CANDIDATE_PIPELINE, + pipeline_version=metadata.version("microcosm-build"), + candidate_id=args.staging_candidate_id or state.build_id, + local_dir=args.staging_dir or out_dir / "staging", + run_kind="calibration", + delivery_mode="local_only" if local_only else "local_and_remote", + repo_id=None if local_only else args.staging_repo_id, + upload_interval_seconds=args.staging_upload_interval_seconds, + api=None if local_only else _hub_api(), + ) + + +def _stage( + telemetry: StagingTelemetryV2 | None, + stage_id: str, + event_status: str = "started", + **details: Any, +) -> None: + if telemetry is not None: + telemetry.stage(stage_id, event_status=event_status, **details) + + +def _thinned_epochs( + sink: Callable[[Mapping[str, Any]], None], *, every: int = _STAGING_EPOCH_EVERY +) -> Callable[[dict[str, object]], None]: + """Forward every ``every``-th epoch and each phase's last epoch to ``sink``. + + The kernel flags probe epochs with ``budget_search: True``; the staging + contract records that field as an integer or null, so the flag becomes 1 + (the national command never runs a budget search and never met this). + """ + + def callback(event: dict[str, object]) -> None: + if event.get("kind") != "calibration_epoch": + return + epoch = int(event["epoch"]) + epochs = int(event["epochs"]) + if epoch % every != 0 and epoch != epochs: + return + forwarded = dict(event) + budget_search = forwarded.get("budget_search") + if isinstance(budget_search, bool): + forwarded["budget_search"] = 1 if budget_search else None + sink(forwarded) + + return callback + + +def _size_checkpoint_state(solve: UKRowwiseDoctrineSolve) -> str | None: + if solve.size_receipt is None or not solve.size_receipt.get("checkpoint"): + return None + checkpoint = solve.size_receipt["checkpoint"] + if "written" in checkpoint: + return "written" + if "resumed_from" in checkpoint: + return "resumed" + return None + + +def _gate_statuses(gate_report: Mapping[str, Any]) -> dict[str, str]: + return { + str(gate_id): str(entry.get("status")) + for gate_id, entry in gate_report.get("gates", {}).items() + if isinstance(entry, Mapping) + } + + +def _fail_staging_telemetry( + telemetry: StagingTelemetryV2 | None, error: BaseException +) -> None: + if telemetry is None or telemetry.status != "running": + return + try: + telemetry.fail(error) + telemetry.validate_local_bundle() + except Exception: + pass + + +def _finalize_staging_telemetry( + args: argparse.Namespace, telemetry: StagingTelemetryV2 | None +) -> None: + if telemetry is None: + return + telemetry.complete(message="UK rowwise candidate staging run completed.") + try: + if args.staging_read_back: + telemetry.verify_remote() + finally: + telemetry.validate_local_bundle() + + +def _staging_delivery(telemetry: StagingTelemetryV2 | None) -> dict[str, Any]: + if telemetry is None: + return disabled_staging_delivery("--no-staging") + return telemetry.delivery_summary + + +def _stage_dataset( + args: argparse.Namespace, + *, + manifest: Mapping[str, Any], + output_paths: Mapping[str, Path], + run_id: str, + telemetry: StagingTelemetryV2 | None, +) -> dict[str, Any]: + """Stage the published bundle under ``staged//``; record, never raise. + + The bundle is every file the manifest registers as an output plus the + manifest and two sidecars, verified from disk against the manifest's own + digests. Nothing else in the run directory is eligible. + """ + + mode = _staged_dataset_mode(args) + if mode == "disabled": + return disabled_staged_dataset( + "--no-staging" if args.no_staging else "--no-staged-dataset" + ) + repository = ( + None if mode == "local_only" else str(args.staged_dataset_repo_id).strip() + ) + _stage(telemetry, "dataset_staging", "started", mode=mode, repository=repository) + gate_statuses = _gate_statuses(getattr(args, "_gate_report", {}) or {}) + bundle = StagedDatasetBundle.from_manifest( + output_paths["manifest"].parent, + run_id=run_id, + manifest_name=MANIFEST_FILENAME, + extra_summary={ + "dataset_households": manifest["parameters"]["dataset_households"], + "pool_rows": manifest["solve"]["pool_households"], + "realized_households": manifest["solve"]["n_households"], + "final_loss": manifest["solve"]["final_loss"], + "release_posture": manifest["release_posture"], + "gate_statuses": gate_statuses, + }, + ) + telemetry_reference = ( + None + if telemetry is None + else { + "repository": telemetry.repo_id, + "prefix": telemetry.repo_run_prefix, + "mode": telemetry.delivery_mode, + } + ) + write_sidecars( + bundle, + repository=repository, + prefix=UK_STAGED_DATASET_PREFIX, + telemetry=telemetry_reference, + ) + if mode == "local_only": + delivery = local_only_staged_dataset(bundle, prefix=UK_STAGED_DATASET_PREFIX) + else: + print( + f"staging the dataset bundle to {repository} under " + f"{bundle.remote_prefix(UK_STAGED_DATASET_PREFIX)}...", + file=sys.stderr, + flush=True, + ) + storage = HuggingFaceDatasetStorage(repository, api=_hub_api()) + delivery = stage_bundle( + bundle, storage=storage, prefix=UK_STAGED_DATASET_PREFIX + ) + print(_staged_dataset_line(delivery), file=sys.stderr, flush=True) + _stage( + telemetry, + "dataset_staging", + "completed", + status=delivery["status"], + repository=delivery["repository"], + revision=delivery["revision"], + error_code=delivery["error_code"], + file_count=len(delivery["files"]), + ) + _add_staging_artifact( + telemetry, + "staged_dataset", + delivery, + artifact_kind="build_metadata", + classification="non_row_level", + ) + _add_staging_artifact( + telemetry, + "fit_summary", + _fit_summary( + manifest, + run_id=run_id, + gate_statuses=gate_statuses, + staged_dataset=delivery, + ), + artifact_kind="aggregate_diagnostics", + classification="aggregate", + ) + return delivery + + +def _staged_dataset_line(delivery: Mapping[str, Any]) -> str: + status = delivery["status"] + if status in ("uploaded", "already_staged"): + return ( + f"staged dataset: {status} at {delivery['repository']}/" + f"{delivery['prefix']} (revision {delivery['revision']})" + ) + if status == "failed": + return ( + f"staged dataset: failed ({delivery['error_code']}); the bundle and " + "its sidecars stay local and can be re-staged with " + "tools/stage_uk_rowwise_candidate.py" + ) + return ( + f"staged dataset: skipped ({delivery['mode']}); sidecars written beside " + "the bundle" + ) + + +def _add_staging_artifact( + telemetry: StagingTelemetryV2 | None, + logical_name: str, + payload: Mapping[str, Any], + *, + artifact_kind: str, + classification: str, +) -> None: + """Attach a reviewed aggregate JSON artifact to the telemetry run. + + A content-policy refusal is reported and skipped: the telemetry is + best-effort and must never fail a finished build. + """ + + if telemetry is None: + return + with tempfile.TemporaryDirectory(prefix=".staging-artifact.") as scratch: + source = Path(scratch) / f"{logical_name}.json" + source.write_text(_json_text(payload), encoding="utf-8") + try: + telemetry.add_artifact( + logical_name, + source, + artifact_kind=artifact_kind, + classification=classification, + ) + except StagingContractError as error: + print( + f"warning: staging artifact {logical_name} was refused by the " + f"content policy and is not staged: {error}", + file=sys.stderr, + flush=True, + ) + + +def _fit_summary( + manifest: Mapping[str, Any], + *, + run_id: str, + gate_statuses: Mapping[str, str], + staged_dataset: Mapping[str, Any], +) -> dict[str, Any]: + """Aggregate fit, gate and size evidence shaped for a reviewed artifact. + + The content policy admits JSON objects without arrays of objects, so the + per-family rows become mappings keyed by family and anything row-shaped + is dropped by :func:`_aggregate_only`. + """ + + solve = manifest["solve"] + fit = manifest.get("fit") or {} + size = solve.get("dataset_size") + parameters = manifest["parameters"] + return { + "schema_name": "microcosm.uk.rowwise-fit-summary", + "schema_version": 1, + "run_id": run_id, + "build_kind": manifest["build_kind"], + "releasable": manifest["releasable"], + "release_posture": _aggregate_only(manifest["release_posture"]), + "git_commit": manifest["git_commit"], + "git_dirty": manifest["git_dirty"], + "parameters": { + key: parameters.get(key) + for key in ( + "n_clones", + "dataset_households", + "seed", + "epochs", + "sample_fraction", + "release_candidate", + "skip_holdout", + "target_weight_rule", + ) + }, + "targets": { + "count": solve["n_targets"], + "by_kind": _aggregate_only(solve["n_targets_by_kind"]), + }, + "pool_rows": solve["pool_households"], + "realized_households": solve["n_households"], + "loss": { + "initial": solve["initial_loss"], + "final": solve["final_loss"], + "max_abs_relative_error": solve["max_abs_relative_error"], + "median_abs_relative_error": solve["median_abs_relative_error"], + }, + "fit_by_family": { + "local": _rows_by_key(fit.get("local_by_family"), key="family"), + "national": _rows_by_key(fit.get("national_by_family"), key="family"), + }, + "weakest_areas_by_fit": _aggregate_only(fit.get("weakest_areas_by_fit")), + "rotated_holdout": _aggregate_only(fit.get("rotated_holdout")), + "gates": dict(gate_statuses), + "failing_gate_ids": list(manifest.get("failing_gate_ids", [])), + "blocking_failure_count": len(manifest.get("blocking_failures", [])), + # The receipt's per-row arrays (pool_row_indices, inclusion + # probabilities) live in dataset_size_selection.csv and would push the + # artifact past the 5 MiB cap on a real run. + "dataset_size": None + if size is None + else _aggregate_only( + { + key: value + for key, value in size.items() + if key not in ("pool_row_indices", "inclusion_probabilities") + } + ), + "staged_dataset": { + key: staged_dataset[key] + for key in ("repository", "prefix", "revision", "status", "error_code") + }, + } + + +def _rows_by_key(rows: Any, *, key: str) -> dict[str, Any]: + """Turn a list of row mappings into a mapping keyed by ``row[key]``.""" + + if not isinstance(rows, list): + return {} + keyed: dict[str, Any] = {} + for row in rows: + if not isinstance(row, Mapping) or key not in row: + continue + keyed[str(row[key])] = _aggregate_only( + {name: value for name, value in row.items() if name != key} + ) + return keyed + + +def _aggregate_only(value: Any) -> Any: + """Drop row-shaped data (lists holding mappings) recursively.""" + + if isinstance(value, Mapping): + kept = {} + for name, item in value.items(): + cleaned = _aggregate_only(item) + if cleaned is not _DROPPED: + kept[str(name)] = cleaned + return kept + if isinstance(value, (list, tuple)): + if any(isinstance(item, Mapping) for item in value): + return _DROPPED + return [ + item + for item in (_aggregate_only(entry) for entry in value) + if item is not _DROPPED + ] + return value + + +_DROPPED = object() + + +def _replace_manifest(path: Path, manifest: Mapping[str, Any]) -> None: + """Rewrite the published manifest atomically with appended evidence.""" + + handle, temporary = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent) + temporary_path = Path(temporary) + try: + with open(handle, "w", encoding="utf-8") as stream: + stream.write(_json_text(manifest)) + temporary_path.replace(path) + except BaseException: + temporary_path.unlink(missing_ok=True) + raise + + def _clone_with_ladder_binding( dataset: Any, ladder: UkOaLadder, diff --git a/tools/fetch_uk_staged_dataset.py b/tools/fetch_uk_staged_dataset.py new file mode 100644 index 0000000000..23ad8375a5 --- /dev/null +++ b/tools/fetch_uk_staged_dataset.py @@ -0,0 +1,84 @@ +"""Fetch a staged UK rowwise candidate bundle and verify its digests. + +Downloads ``staged//`` from the private dataset repository into a +local directory, checking every file against the bundle's ``sha256sums.txt``, +and prints one local path per line. Scorecards, the evaluation legs and the +calibration dashboard's local-directory mode take those paths as they take +any run directory. ``--h5-only`` fetches just the dataset. +""" + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path + +from microcosm.build.staging_dataset import StagedDatasetError, fetch_bundle +from microcosm.build.staging_storage import HuggingFaceDatasetStorage +from microcosm.build.uk_runtime.staging import ( + UK_STAGED_DATASET_PREFIX, + UK_STAGED_DATASET_REPOSITORY, +) + + +def _hub_api(): + from huggingface_hub import HfApi + + return HfApi() + + +def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run-id", required=True, help="The staged run id.") + parser.add_argument( + "--repo-id", + default=UK_STAGED_DATASET_REPOSITORY.repo_id(os.environ), + help=( + "Private Hugging Face dataset repository (default from " + f"{UK_STAGED_DATASET_REPOSITORY.repo_id_environment_variable})." + ), + ) + parser.add_argument( + "--dest", + type=Path, + help="Destination directory (default: ./staged/).", + ) + parser.add_argument( + "--revision", + help="Repository revision to fetch (default: the current main branch).", + ) + parser.add_argument( + "--h5-only", + action="store_true", + help="Fetch only the dataset file(s).", + ) + return parser.parse_args(argv) + + +def main(argv: list[str] | None = None) -> int: + args = _parse_args(argv) + repo_id = str(args.repo_id).strip() + if not repo_id: + raise SystemExit("error: --repo-id must be non-empty") + dest = args.dest or Path("staged") / args.run_id + storage = HuggingFaceDatasetStorage(repo_id, api=_hub_api()) + try: + paths = fetch_bundle( + storage=storage, + run_id=args.run_id, + dest=dest, + prefix=UK_STAGED_DATASET_PREFIX, + revision=args.revision, + h5_only=args.h5_only, + ) + except StagedDatasetError as error: + print(f"error: {error}", file=sys.stderr) + return 1 + for path in paths: + print(path) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tools/stage_uk_rowwise_candidate.py b/tools/stage_uk_rowwise_candidate.py new file mode 100644 index 0000000000..529aaace8d --- /dev/null +++ b/tools/stage_uk_rowwise_candidate.py @@ -0,0 +1,161 @@ +"""Stage a finished UK rowwise candidate run on the private dataset repository. + +The build driver stages its bundle at the end of every run; this command +covers the two other cases: a run whose upload failed (recorded as +``staged_dataset.status == "failed"`` in its manifest) and a run built before +the staged-dataset lane existed. It uploads exactly what the driver would have +uploaded: every file the manifest registers as an output, the manifest itself, +``staged_manifest.json`` and ``sha256sums.txt``, in one commit under +``staged//``. Re-running on an already staged directory is a no-op; +a directory whose bytes differ from the remote bundle is refused. + +Staging is not publication: ``releases/`` and ``latest.json`` are untouched. +""" + +from __future__ import annotations + +import argparse +import glob +import json +import os +import sys +from pathlib import Path + +from microcosm.build.staging_dataset import ( + StagedDatasetBundle, + local_only_staged_dataset, + stage_bundle, + validate_staged_dataset_delivery, + write_sidecars, +) +from microcosm.build.staging_storage import HuggingFaceDatasetStorage +from microcosm.build.uk_runtime.staging import ( + UK_STAGED_DATASET_PREFIX, + UK_STAGED_DATASET_REPOSITORY, +) + +MANIFEST_FILENAME = "rowwise_candidate_manifest.json" + + +def _hub_api(): + from huggingface_hub import HfApi + + return HfApi() + + +def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--run-dir", + type=Path, + required=True, + help="A finished build_uk_rowwise_candidate.py output directory.", + ) + parser.add_argument( + "--repo-id", + default=UK_STAGED_DATASET_REPOSITORY.repo_id(os.environ), + help=( + "Private Hugging Face dataset repository (default from " + f"{UK_STAGED_DATASET_REPOSITORY.repo_id_environment_variable})." + ), + ) + parser.add_argument( + "--run-id", + help=( + "Staged run id; defaults to the manifest's staged_dataset.run_id, " + "then the Logbook build id, then the directory name." + ), + ) + parser.add_argument( + "--local-only", + action="store_true", + help="Verify the bundle and write the sidecars without uploading.", + ) + return parser.parse_args(argv) + + +def _default_run_id(run_dir: Path, manifest: dict) -> str: + staged = manifest.get("staged_dataset") + if isinstance(staged, dict) and staged.get("run_id"): + return str(staged["run_id"]) + rows = sorted(glob.glob(str(run_dir / "logbook-spool" / "*.json"))) + for row_path in reversed(rows): + try: + row = json.loads(Path(row_path).read_text(encoding="utf-8")) + except ValueError: + continue + if isinstance(row, dict) and row.get("build_id"): + return str(row["build_id"]) + return run_dir.name + + +def _rewrite_manifest(path: Path, manifest: dict) -> None: + text = json.dumps(manifest, allow_nan=False, indent=2, sort_keys=True) + "\n" + temporary = path.with_name(f".{path.name}.staging") + temporary.write_text(text, encoding="utf-8") + temporary.replace(path) + + +def main(argv: list[str] | None = None) -> int: + args = _parse_args(argv) + run_dir = args.run_dir.expanduser().resolve() + manifest_path = run_dir / MANIFEST_FILENAME + if not manifest_path.is_file(): + raise SystemExit(f"error: {manifest_path} is missing") + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + run_id = args.run_id or _default_run_id(run_dir, manifest) + repository = None if args.local_only else str(args.repo_id).strip() + if repository == "": + raise SystemExit("error: --repo-id must be non-empty unless --local-only") + bundle = StagedDatasetBundle.from_manifest( + run_dir, + run_id=run_id, + manifest_name=MANIFEST_FILENAME, + extra_summary={ + "dataset_households": (manifest.get("parameters") or {}).get( + "dataset_households" + ), + "release_posture": manifest.get("release_posture"), + }, + ) + telemetry = manifest.get("staging_delivery") + write_sidecars( + bundle, + repository=repository, + prefix=UK_STAGED_DATASET_PREFIX, + telemetry=( + None + if not isinstance(telemetry, dict) + else { + "repository": telemetry.get("configured_repository"), + "prefix": ( + None + if telemetry.get("run_id") is None + else f"runs/{telemetry['run_id']}" + ), + "mode": telemetry.get("mode"), + } + ), + ) + if repository is None: + delivery = local_only_staged_dataset(bundle, prefix=UK_STAGED_DATASET_PREFIX) + else: + print( + f"staging {run_dir} to {repository} under " + f"{bundle.remote_prefix(UK_STAGED_DATASET_PREFIX)}...", + file=sys.stderr, + flush=True, + ) + delivery = stage_bundle( + bundle, + storage=HuggingFaceDatasetStorage(repository, api=_hub_api()), + prefix=UK_STAGED_DATASET_PREFIX, + ) + manifest["staged_dataset"] = validate_staged_dataset_delivery(delivery) + _rewrite_manifest(manifest_path, manifest) + print(json.dumps(delivery, indent=2, sort_keys=True)) + return 0 if delivery["status"] != "failed" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) From 5aac4661313528ebeb9af0cf433d9116e0827a64 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Thu, 17 Sep 2026 18:50:27 +0100 Subject: [PATCH 2/7] Record the two input mismatches met on the staging rehearsal (feed ec20085, ladder bed3f13d) Co-Authored-By: Claude Fable 5.1 --- experiments/355-uk-dataset-size-receipts.md | 26 +++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/experiments/355-uk-dataset-size-receipts.md b/experiments/355-uk-dataset-size-receipts.md index 6c848f0209..9b63829e13 100644 --- a/experiments/355-uk-dataset-size-receipts.md +++ b/experiments/355-uk-dataset-size-receipts.md @@ -986,3 +986,29 @@ no release contract. failure, `--no-staged-dataset`, up-front refusals and dry run), dense assembler evidence tests, size-evaluation tolerance. Every existing driver test passes `--staging-local-only` through the shared flag helper. + +### Rehearsal on spine-q (2026-09-17, `data/ukds/acceptance/355-dataset-size/run_size_candidate_staging.sh`) + +The rehearsal is the size path at 100 epochs with local-only staging +(`--dataset-households 55000 --epochs 100 --skip-holdout --selection-pi-hi 0.5 --baseline-pi-floor 0.001 --staging-local-only`), +on the spine-q H5 (cf1f9dda…) from code 0cda403a (`git_dirty` 0, engine 2.98.0). Two input +mismatches surfaced first, both about running the driver on today's main rather than about staging; +each left the expected failure evidence. + +- **Feed.** With the `ec7169b` artifact the spine-q runs stood on, 631 national target references + failed to compile: main pins Chronicle `ec20085` (facts 47612c48…, manifest c91fa9ff…) since + #927/#937. The run refused in 130 s. Its telemetry closed as `failed` during + `target_compilation` with `error_type SystemExit` and the sanitised message only; because the + compile refusal is a `SystemExit`, the driver's `except Exception` wrote no Logbook row or error + receipt (pre-existing behaviour, worth a follow-up). Kept as + `…-staging-failed-compile-ec7169b/`. +- **Ladder.** On `ec20085` with the `9c6d56b9…` ladder the pins file names, the ladder-versus-Chronicle + dispersion check refused: "NI DZ-to-PARLCON24 household dispersion exceeds the publisher oracle: + mean absolute delta 197.889, max absolute delta 694.000" (285 s, 5.1 GB). #887 (merged 2026-09-10) + rebuilt the ladder with NISRA's Data Zone lookup; the artifact from that build is + `bed3f13d3a82eea2d1f39248b71c0abf5ba6960a446ddd9415ae1dbcb7ae07fd` (present in the #887/#905/#929 + worktrees; the ladder tool has not changed since). This `ValueError` took the full failure path: + telemetry `failed` at `target_compilation`, a `failed` Logbook row and an error receipt. Kept as + `…-staging-failed-ladder-9c6d56b9/`. + +Third attempt: feed `ec20085`, ladder `bed3f13d…`, same arguments; run id and results below. From 0ab7ac4f71b09f587bb0920ef47d3abadc785a2c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:00:47 +0100 Subject: [PATCH 3/7] Record the stale-spine refusal on the staging rehearsal and the spine-r build Co-Authored-By: Claude Fable 5.1 --- experiments/355-uk-dataset-size-receipts.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/experiments/355-uk-dataset-size-receipts.md b/experiments/355-uk-dataset-size-receipts.md index 9b63829e13..9db457963d 100644 --- a/experiments/355-uk-dataset-size-receipts.md +++ b/experiments/355-uk-dataset-size-receipts.md @@ -1012,3 +1012,13 @@ each left the expected failure evidence. `…-staging-failed-ladder-9c6d56b9/`. Third attempt: feed `ec20085`, ladder `bed3f13d…`, same arguments; run id and results below. + +- **Spine.** On `ec20085` with ladder `bed3f13d…`, the run compiled 20,885 local and 564 national + targets, cloned the 792,690-row pool, and refused in surface resolution: + `MeasureResolutionError: provider does not know household.ons_household_type` (550 s, 8.4 GB). + The spine-q H5 predates the `frs_relationships` stage (#903, merged 2026-09-11), whose column the + national registry on main now resolves. Telemetry closed as `failed` at `surface_resolution` + after recording the compile and clone details; Logbook row and error receipt written. Kept as + `…-staging-failed-spine-q-stale/`. A spine on main's stages (`spine-r`, code 0cda403a) is built + with `834-childcare-tfc/build_twin_passthrough.sh … --staging-local-only` into + `data/ukds/acceptance/spine-r-355/`. From db40e56cf349c16b55d94ed9241b497f0eefa8ff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Thu, 17 Sep 2026 19:40:42 +0100 Subject: [PATCH 4/7] Keep the local sha256sums self-verifying after the manifest gains its staging evidence The build appends staging_delivery and staged_dataset to the manifest after the sidecars are written and the bundle uploaded, so the remote sums list the manifest as uploaded while the local directory held a stale digest. The driver and the re-stage tool now refresh that one line after the rewrite; each directory verifies itself. Records the first end-to-end rehearsal on spine-r in the receipts. Co-Authored-By: Claude Fable 5.1 --- experiments/355-uk-dataset-size-receipts.md | 38 +++++++++++++++++++ .../src/microcosm/build/staging_dataset.py | 22 +++++++++++ .../tests/test_staging_dataset.py | 24 ++++++++++++ .../tests/test_uk_rowwise_candidate.py | 8 +++- tools/build_uk_rowwise_candidate.py | 8 ++++ tools/stage_uk_rowwise_candidate.py | 2 + 6 files changed, 101 insertions(+), 1 deletion(-) diff --git a/experiments/355-uk-dataset-size-receipts.md b/experiments/355-uk-dataset-size-receipts.md index 9db457963d..64ad41b232 100644 --- a/experiments/355-uk-dataset-size-receipts.md +++ b/experiments/355-uk-dataset-size-receipts.md @@ -1022,3 +1022,41 @@ Third attempt: feed `ec20085`, ladder `bed3f13d…`, same arguments; run id and `…-staging-failed-spine-q-stale/`. A spine on main's stages (`spine-r`, code 0cda403a) is built with `834-childcare-tfc/build_twin_passthrough.sh … --staging-local-only` into `data/ukds/acceptance/spine-r-355/`. + +### spine-r (2026-09-17 18:01Z, `data/ukds/acceptance/spine-r-355/`) + +Built from the worktree at 0ab7ac4f (code identical to 0cda403a; the two commits between are +receipts) with `build_twin_passthrough.sh … --staging-local-only`: 29 stages (`frs_spine`, +`age_tail`, `frs_relationships`, …), `rules_engine` policyengine-uk 2.98.0, all 18 spine gates +passed, `spine-r.h5` sha `3ce8756ac6ce070be4eabf44e6969e03698ce23531021b622eadb94ce83dda58` +(169 MB). The spine builder's own staging worked as on main: sidecar `staging_delivery` mode +`local_only`, run `uk-frs-spine-20260917T180117Z` under `spine-r-355/staging/runs/`. Pins for the +rehearsal in `355-dataset-size/pins-spine-r-ec20085.txt`; run root `355-dataset-size/spine-r/`. + +### Rehearsal r1 on spine-r (2026-09-17 18:07–18:35Z, run `uk-local-candidate-f100-s42-20260917T180755Z-a244651d`) + +End to end for the first time: 21,449 targets (564 national, 19,874 local, 1,011 ladder) on the +792,690-row pool; dense solve 100 epochs to loss 0.01956; the search settled after 4 probes +(λ 3.16e-06, open mass 54,674 for 55,000, 43,681 certainties); refit loss 0.04559 on 55,000 rows +from the floored Horvitz–Thompson baseline; gates blocked on `uk_local_area_support`, +`uk_local_target_fit`, `uk_local_weight_ratio` (866 blocking lines, the 55k pattern Q50f showed); +exit 1 by design, bundle written, 27.4 min wall, 12.0 GB peak RSS. + +Staging (local-only): the telemetry bundle validates; `run_manifest.json` 1.9 KiB, `events.ndjson` +32.6 KiB, `calibration_progress.json` 17.9 KiB with 60 rows (10 dense with `phase` null, 40 probe +rows with `budget_search` 1 and their λ, 10 refit rows tagged `size_refit`), two reviewed +artifacts (`fit_summary.json` 16.1 KiB, `staged_dataset.json` 1.6 KiB); every file far under the +5 MiB cap. Completed stages in order: input_pinning, target_compilation, cloning, +surface_resolution, calibration (details: final loss, 55,000 realized, checkpoint written), +gate_battery, holdout (skipped), output_bundle, dataset_staging (`skipped`, 9 files), complete. +The manifest carries `staging_delivery` (mode local_only, run id) and `staged_dataset` (mode +local_only, prefix `staged/`, 9 files with digests); `staged_manifest.json` and +`sha256sums.txt` written beside the bundle (374 MB: H5 237 MB, diagnostics 49 MB, registry 29 MB, +two CSVs 47 MB, manifest 6 MB). + +Two warts, both fixed before the second rehearsal: the local `sha256sums.txt` listed the manifest +as uploaded rather than as rewritten with the evidence blocks (the driver and the re-stage tool now +refresh that line after the rewrite, so the local and the remote directory each verify themselves); +and the manifest recorded `git_dirty: true` because the receipts draft was edited in the worktree +while the run was going (the code was unchanged; never touch the run tree during a run). r1 is kept +as `spine-r/…-staging-r1/` and is not the bundle staged on the Hub. diff --git a/packages/microcosm-build/src/microcosm/build/staging_dataset.py b/packages/microcosm-build/src/microcosm/build/staging_dataset.py index 10cd355b07..3539133a3e 100644 --- a/packages/microcosm-build/src/microcosm/build/staging_dataset.py +++ b/packages/microcosm-build/src/microcosm/build/staging_dataset.py @@ -42,6 +42,7 @@ "fetch_bundle", "local_only_staged_dataset", "parse_sha256sums", + "refresh_sha256sums_entry", "sha256_file", "stage_bundle", "validate_staged_dataset_delivery", @@ -308,6 +309,27 @@ def write_sidecars( return sums_path, staged_manifest_path +def refresh_sha256sums_entry(run_dir: Path | str, name: str) -> str: + """Re-digest one listed file so the local sums stay self-verifying. + + The build appends its evidence blocks to the manifest after the sidecars + were written and uploaded; the remote copy lists the manifest as + uploaded, the local copy must list the manifest as it now is. + """ + + directory = Path(run_dir) + sums_path = directory / SHA256SUMS_FILENAME + entries = parse_sha256sums(sums_path.read_text(encoding="utf-8")) + if name not in {listed for _, listed in entries}: + raise StagedDatasetError(f"{SHA256SUMS_FILENAME} does not list {name!r}.") + digest = sha256_file(directory / name) + text = "".join( + f"{digest if listed == name else sha} {listed}\n" for sha, listed in entries + ) + sums_path.write_text(text, encoding="utf-8") + return digest + + def parse_sha256sums(text: str) -> list[tuple[str, str]]: """Parse ``sha256sum`` lines into ``(digest, name)`` pairs.""" diff --git a/packages/microcosm-build/tests/test_staging_dataset.py b/packages/microcosm-build/tests/test_staging_dataset.py index c5bc1399fa..8a0d7cc1df 100644 --- a/packages/microcosm-build/tests/test_staging_dataset.py +++ b/packages/microcosm-build/tests/test_staging_dataset.py @@ -18,6 +18,7 @@ fetch_bundle, local_only_staged_dataset, parse_sha256sums, + refresh_sha256sums_entry, stage_bundle, validate_staged_dataset_delivery, write_sidecars, @@ -375,3 +376,26 @@ def test_parse_sha256sums_rejects_paths_and_duplicates(): parse_sha256sums(f"{'a' * 64} file.h5\n{'b' * 64} file.h5\n") with pytest.raises(StagedDatasetError, match="no files"): parse_sha256sums("\n") + + +def test_refresh_sha256sums_entry_redigests_one_listed_file(tmp_path): + run_dir = _write_run(tmp_path) + bundle = StagedDatasetBundle.from_manifest( + run_dir, run_id=RUN_ID, manifest_name=MANIFEST + ) + sums, _ = write_sidecars(bundle, repository=None) + before = dict((name, sha) for sha, name in parse_sha256sums(sums.read_text())) + (run_dir / MANIFEST).write_text( + json.dumps( + {**json.loads((run_dir / MANIFEST).read_text()), "staged_dataset": {}} + ) + ) + digest = refresh_sha256sums_entry(run_dir, MANIFEST) + after = dict((name, sha) for sha, name in parse_sha256sums(sums.read_text())) + assert digest == _sha((run_dir / MANIFEST).read_bytes()) == after[MANIFEST] + assert after[MANIFEST] != before[MANIFEST] + assert {k: v for k, v in after.items() if k != MANIFEST} == { + k: v for k, v in before.items() if k != MANIFEST + } + with pytest.raises(StagedDatasetError, match="does not list"): + refresh_sha256sums_entry(run_dir, "run.log") diff --git a/packages/microcosm-build/tests/test_uk_rowwise_candidate.py b/packages/microcosm-build/tests/test_uk_rowwise_candidate.py index 0ab82ea147..de2708e404 100644 --- a/packages/microcosm-build/tests/test_uk_rowwise_candidate.py +++ b/packages/microcosm-build/tests/test_uk_rowwise_candidate.py @@ -2775,7 +2775,10 @@ def test_candidate_build_stages_telemetry_locally_and_inventories_the_bundle( } for entry in manifest["outputs"].values(): assert staged["files"][Path(entry["path"]).name]["sha256"] == entry["sha256"] - assert (out / "sha256sums.txt").is_file() + # The local sums verify the directory as it is, evidence blocks included. + for line in (out / "sha256sums.txt").read_text().splitlines(): + digest, name = line.split(" ") + assert hashlib.sha256((out / name).read_bytes()).hexdigest() == digest, name inventory = json.loads((out / "staged_manifest.json").read_text()) assert inventory["run_id"] == run_id and inventory["files"] == staged["files"] assert inventory["summary"]["releasable"] is True @@ -2973,6 +2976,9 @@ def test_remote_staging_uploads_telemetry_and_the_bundle_in_one_commit( "staged_dataset" ] assert restaged["status"] == "already_staged" + for line in (out / "sha256sums.txt").read_text().splitlines(): + digest, name = line.split(" ") + assert hashlib.sha256((out / name).read_bytes()).hexdigest() == digest, name assert restaged["run_id"] == run_id and restaged["revision"] == hub.sha assert len(hub.commits) == 1 diff --git a/tools/build_uk_rowwise_candidate.py b/tools/build_uk_rowwise_candidate.py index f7aac25ce8..b85c3baf72 100644 --- a/tools/build_uk_rowwise_candidate.py +++ b/tools/build_uk_rowwise_candidate.py @@ -62,9 +62,11 @@ validate_staging_arguments, ) from microcosm.build.staging_dataset import ( + SHA256SUMS_FILENAME, StagedDatasetBundle, disabled_staged_dataset, local_only_staged_dataset, + refresh_sha256sums_entry, stage_bundle, write_sidecars, ) @@ -1498,6 +1500,12 @@ def _run_candidate( manifest["staging_delivery"] = _staging_delivery(telemetry) manifest["staged_dataset"] = staged_dataset _replace_manifest(output_paths["manifest"], manifest) + if (output_paths["manifest"].parent / SHA256SUMS_FILENAME).is_file(): + # The uploaded copy lists the manifest as uploaded; the local + # copy lists the manifest as it now is, evidence included. + refresh_sha256sums_entry( + output_paths["manifest"].parent, output_paths["manifest"].name + ) state.artifact_location = local_artifact_reference( output_paths["dataset"], repository_hint=_REPOSITORY, diff --git a/tools/stage_uk_rowwise_candidate.py b/tools/stage_uk_rowwise_candidate.py index 529aaace8d..fc58674e4a 100644 --- a/tools/stage_uk_rowwise_candidate.py +++ b/tools/stage_uk_rowwise_candidate.py @@ -24,6 +24,7 @@ from microcosm.build.staging_dataset import ( StagedDatasetBundle, local_only_staged_dataset, + refresh_sha256sums_entry, stage_bundle, validate_staged_dataset_delivery, write_sidecars, @@ -153,6 +154,7 @@ def main(argv: list[str] | None = None) -> int: ) manifest["staged_dataset"] = validate_staged_dataset_delivery(delivery) _rewrite_manifest(manifest_path, manifest) + refresh_sha256sums_entry(run_dir, MANIFEST_FILENAME) print(json.dumps(delivery, indent=2, sort_keys=True)) return 0 if delivery["status"] != "failed" else 1 From 15b350b0bb1c88fd600fa1dec4d435be8a736236 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Thu, 17 Sep 2026 20:13:14 +0100 Subject: [PATCH 5/7] Refuse a read-only Hub credential before a remote dataset stage A read token sees the private repository, so the reachability pre-flight passed and the upload at the end of the run was refused with 403. The driver's pre-flight and the re-stage tool now read the credential's role from the Hub and refuse `read` up front, naming the write requirement. Records the r2 rehearsal and the Hub refusal in the receipts. Co-Authored-By: Claude Fable 5.1 --- experiments/355-uk-dataset-size-receipts.md | 26 +++++++++++++++++++ .../src/microcosm/build/staging_storage.py | 16 ++++++++++++ .../microcosm-build/tests/test_staging_v2.py | 19 ++++++++++++++ .../tests/test_uk_rowwise_candidate.py | 25 +++++++++++++++++- tools/build_uk_rowwise_candidate.py | 14 +++++++++- tools/stage_uk_rowwise_candidate.py | 14 +++++++--- 6 files changed, 109 insertions(+), 5 deletions(-) diff --git a/experiments/355-uk-dataset-size-receipts.md b/experiments/355-uk-dataset-size-receipts.md index 64ad41b232..73b48521af 100644 --- a/experiments/355-uk-dataset-size-receipts.md +++ b/experiments/355-uk-dataset-size-receipts.md @@ -1060,3 +1060,29 @@ refresh that line after the rewrite, so the local and the remote directory each and the manifest recorded `git_dirty: true` because the receipts draft was edited in the worktree while the run was going (the code was unchanged; never touch the run tree during a run). r1 is kept as `spine-r/…-staging-r1/` and is not the bundle staged on the Hub. + +### Rehearsal r2 on spine-r (2026-09-17 18:41–19:08Z, run `uk-local-candidate-f100-s42-20260917T184110Z-4a5f5af3`) + +Same arguments and inputs as r1, code db40e56c, tree untouched during the run: `git_dirty` +false in the manifest. Deterministic against r1: identical `solve_diagnostics.csv` and +`dataset_size_selection.csv` digests, dense loss 0.019561, search settled after 4 probes on +λ 3.16e-06, refit loss 0.045593, gates blocked on the same three ids; 26.9 min wall, 13.2 GB peak. +Telemetry bundle valid (60 calibration rows, six files under the cap, ten stages completed in +order), `staging_delivery` and `staged_dataset` in the manifest, and the local `sha256sums.txt` +verifies every listed file after the evidence rewrite. This is the bundle staged on the Hub with +`tools/stage_uk_rowwise_candidate.py` (results below). + +### Staging r2 on the Hub (2026-09-17 19:1xZ) + +`tools/stage_uk_rowwise_candidate.py --run-dir …/spine-r/f100-k15-h55000-e100-p50-f001-s42-staging` +against `policyengine/populace-uk-private`: the repository was reachable (read), the remote +prefix was absent, and the single commit was refused by the Hub with 403 "you must use a write +token to upload to a repository". The ambient credential on this machine is the **read** token +`PEHF` (user juaristi22). The lane recorded `status failed`, `error_code UPLOAD_FAILED`, nine +files with digests, no revision, exit 1, no exception text in the evidence, the bundle and sidecars +intact for a re-stage; 11 s. Because a read token also passes the reachability pre-flight (it can +see the private repository), the pre-flight and the re-stage tool now refuse a credential whose +Hub role is `read` before any work, naming the write requirement. The real upload waits on a write +credential in María's environment (`HF_TOKEN` or `hf auth login`), then +`tools/stage_uk_rowwise_candidate.py --run-dir ` stages this bundle, and the fetch-back and +scorecard checks follow. diff --git a/packages/microcosm-build/src/microcosm/build/staging_storage.py b/packages/microcosm-build/src/microcosm/build/staging_storage.py index 93d570e4a1..6f4755fb8a 100644 --- a/packages/microcosm-build/src/microcosm/build/staging_storage.py +++ b/packages/microcosm-build/src/microcosm/build/staging_storage.py @@ -113,6 +113,22 @@ def file_exists(self, path_in_repo: str) -> bool: ) ) + def credential_role(self) -> str | None: + """The ambient credential's role as the Hub reports it (``read``, ``write``, …). + + ``None`` when the backend cannot say; a fine-grained token reports + ``fineGrained`` and its scopes are only tested by the upload itself. + """ + + whoami = getattr(self._api(), "whoami", None) + if not callable(whoami): + return None + info = whoami() + auth = info.get("auth") if isinstance(info, Mapping) else None + token = auth.get("accessToken") if isinstance(auth, Mapping) else None + role = token.get("role") if isinstance(token, Mapping) else None + return str(role) if role else None + def head_revision(self) -> str | None: """The default branch's current commit, when the backend reports one.""" diff --git a/packages/microcosm-build/tests/test_staging_v2.py b/packages/microcosm-build/tests/test_staging_v2.py index 90419176c6..bd8d4193d8 100644 --- a/packages/microcosm-build/tests/test_staging_v2.py +++ b/packages/microcosm-build/tests/test_staging_v2.py @@ -876,6 +876,25 @@ def test_storage_commits_several_files_at_once_and_reports_the_revision(tmp_path assert storage.download_file("staged/run/b.h5").read_bytes() == b"binary" +def test_storage_reports_the_credential_role_when_the_backend_can(tmp_path): + class RoleApi(CommitApi): + def whoami(self): + return {"name": "x", "auth": {"accessToken": {"role": "read"}}} + + assert ( + HuggingFaceDatasetStorage( + "example/private", api=RoleApi(tmp_path) + ).credential_role() + == "read" + ) + assert ( + HuggingFaceDatasetStorage( + "example/private", api=CommitApi(tmp_path) + ).credential_role() + is None + ) + + def test_storage_commit_refuses_a_revisionless_backend(tmp_path): class NoRevision(CommitApi): def create_commit(self, **kwargs): diff --git a/packages/microcosm-build/tests/test_uk_rowwise_candidate.py b/packages/microcosm-build/tests/test_uk_rowwise_candidate.py index de2708e404..ecb0ba4e9d 100644 --- a/packages/microcosm-build/tests/test_uk_rowwise_candidate.py +++ b/packages/microcosm-build/tests/test_uk_rowwise_candidate.py @@ -2575,7 +2575,8 @@ def _load_tool(name: str): class _FakeHub: """One fake Hub serving the telemetry repo and the private dataset repo.""" - def __init__(self, *, fail_commit: bool = False) -> None: + def __init__(self, *, fail_commit: bool = False, role: str = "write") -> None: + self.role = role self.files: dict[tuple[str, str], bytes] = {} self.uploads: list[tuple[str, str]] = [] self.commits: list[dict[str, object]] = [] @@ -2609,6 +2610,9 @@ def repo_info(self, *, repo_id, repo_type): assert repo_type == "dataset" return SimpleNamespace(sha=self.sha) + def whoami(self): + return {"name": "tester", "auth": {"accessToken": {"role": self.role}}} + def create_commit( self, *, repo_id, operations, commit_message, repo_type, parent_commit ): @@ -3097,7 +3101,26 @@ def repo_info(self, **kwargs): assert "do-not-record" not in str(info.value) assert not out.exists() + # A read token sees the private repository but cannot upload: refused + # before the spine is read, not after the solve (the Hub answers 403). + monkeypatch.setattr(builder, "_hub_api", lambda: _FakeHub(role="read")) + with pytest.raises(ValueError, match="read-only"): + builder.main(_build_args(input_h5, ladder_path, flags, out)) + assert not out.exists() + + # The re-stage tool refuses the same credential the same way. + stager = _load_tool("stage_uk_rowwise_candidate") + monkeypatch.setattr(stager, "_hub_api", lambda: _FakeHub(role="read")) + local_out = tmp_path / "local" + monkeypatch.setattr(builder, "_hub_token", lambda: None) + assert builder.main( + _build_args(input_h5, ladder_path, flags, local_out, "--staging-local-only") + ) in (0, 1) + with pytest.raises(SystemExit, match="read-only"): + stager.main(["--run-dir", str(local_out)]) + # A dry run plans without staging, so it needs neither credential nor repo. + capsys.readouterr() monkeypatch.setattr(builder, "_hub_token", lambda: None) assert ( builder.main(_build_args(input_h5, ladder_path, flags, out, "--dry-run")) == 0 diff --git a/tools/build_uk_rowwise_candidate.py b/tools/build_uk_rowwise_candidate.py index b85c3baf72..57f2d48797 100644 --- a/tools/build_uk_rowwise_candidate.py +++ b/tools/build_uk_rowwise_candidate.py @@ -1599,13 +1599,25 @@ def _preflight_staged_dataset(args: argparse.Namespace) -> None: f"remote dataset staging to {repo_id} needs a Hugging Face write " f"credential (HF_TOKEN or `hf auth login`); {hint}." ) + storage = HuggingFaceDatasetStorage(repo_id, api=_hub_api()) try: - _hub_api().repo_info(repo_id=repo_id, repo_type="dataset") + storage.head_revision() except Exception as error: raise ValueError( f"remote dataset staging cannot reach {repo_id} " f"({type(error).__name__}); {hint}." ) from error + try: + role = storage.credential_role() + except Exception: + role = None + if role == "read": + # A read token sees the private repository, so the reachability check + # passes; the upload at the end of the run would be refused (403). + raise ValueError( + f"remote dataset staging to {repo_id} needs a write credential; the " + f"ambient Hugging Face token is read-only; {hint}." + ) def _create_staging_telemetry( diff --git a/tools/stage_uk_rowwise_candidate.py b/tools/stage_uk_rowwise_candidate.py index fc58674e4a..f7de7612ac 100644 --- a/tools/stage_uk_rowwise_candidate.py +++ b/tools/stage_uk_rowwise_candidate.py @@ -147,10 +147,18 @@ def main(argv: list[str] | None = None) -> int: file=sys.stderr, flush=True, ) + storage = HuggingFaceDatasetStorage(repository, api=_hub_api()) + try: + role = storage.credential_role() + except Exception: + role = None + if role == "read": + raise SystemExit( + f"error: the ambient Hugging Face token is read-only; staging to " + f"{repository} needs a write credential (HF_TOKEN or `hf auth login`)." + ) delivery = stage_bundle( - bundle, - storage=HuggingFaceDatasetStorage(repository, api=_hub_api()), - prefix=UK_STAGED_DATASET_PREFIX, + bundle, storage=storage, prefix=UK_STAGED_DATASET_PREFIX ) manifest["staged_dataset"] = validate_staged_dataset_delivery(delivery) _rewrite_manifest(manifest_path, manifest) From a8d57940df60f28f763920c3d00928d47d9cf8d7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Thu, 17 Sep 2026 20:17:44 +0100 Subject: [PATCH 6/7] Run the credential pre-flight after the argument refusals Co-Authored-By: Claude Fable 5.1 --- tools/build_uk_rowwise_candidate.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tools/build_uk_rowwise_candidate.py b/tools/build_uk_rowwise_candidate.py index 57f2d48797..d3dcb571af 100644 --- a/tools/build_uk_rowwise_candidate.py +++ b/tools/build_uk_rowwise_candidate.py @@ -858,7 +858,6 @@ def main(argv: list[str] | None = None) -> int: args = _parse_args(argv) _validate_cli_args(args) - _preflight_staged_dataset(args) if args.candidate_clone_counts is not None and not args.dry_run: raise ValueError("--candidate-clone-counts is valid only with --dry-run.") if _CONSERVE_MASS: @@ -872,6 +871,9 @@ def main(argv: list[str] | None = None) -> int: # Dry runs plan without solving or writing and record no Logbook # row on any path, so they need no chain configuration. return _run_candidate(args, attempt=None) + # Argument refusals above cost nothing; the credential check reaches the + # Hub, so it runs last, still before any input is read. + _preflight_staged_dataset(args) started_at = time.perf_counter() started_ts = datetime.now(UTC) digest = preflight_digest(_UK_CANDIDATE_PIPELINE) From a4d30c9f7ec61da85937526648447d232b883144 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mar=C3=ADa=20Juaristi?= <127882282+juaristi22@users.noreply.github.com> Date: Fri, 18 Sep 2026 11:18:32 +0100 Subject: [PATCH 7/7] Address round 1 of #946: keep the driver's receipt on re-stage, bound the epoch telemetry, check the credential's write scope, let pre-lane dense runs assemble - The re-stage tool leaves a directory alone when its record already says the same outputs reached the repository (the driver's revision stands, no sidecar rewrite, no commit); an already staged bundle the record does not know about is recorded with its own commit (get_paths_info), not the repository head. - The forwarded-epoch stride is size-aware (at most 2,400 rows per run for any --epochs and probe count) and a telemetry content refusal is reported once and stops the forwarding instead of aborting the solve; an invalid local bundle at finalize is a warning, never the run's failure (a requested read-back still propagates, as on the national command). - The pre-flight and the tool read the token's scopes from the Hub and refuse a credential without repo.write on the repository or its owner up front (a fine-grained token scoped to a user passed the role check and met the Hub's 403 hours later); an unreadable scope is warned about and proven by the upload. - assemble_uk_dense_release_dir.py --allow-missing-staging assembles a run built before the lane with a recorded disabled-staging opt-out; invalid evidence is still refused. - Nits: no credential identifiers in the receipts, placeholder time and missing run id fixed, the BestEffortUploadSession reflow reverted, --no-staged-dataset help says telemetry alone, the reachability error no longer chains the transport message, _same_digests explains the manifest exclusion, CLAUDE.md tells agents to pass --staging-local-only unless asked, and the argument-refusal ordering has its test. Co-Authored-By: Claude Fable 5.1 --- CLAUDE.md | 6 +- ...ense-assembler-staging-evidence.changed.md | 2 +- changelog.d/uk-staged-dataset-lane.added.md | 2 +- docs/uk-dataset-size-plan-355.md | 4 +- docs/uk-staging-operations.md | 30 ++- experiments/355-uk-dataset-size-receipts.md | 20 +- .../src/microcosm/build/staging_cli.py | 8 +- .../src/microcosm/build/staging_dataset.py | 39 +++- .../src/microcosm/build/staging_storage.py | 91 ++++++++- .../tests/test_staging_dataset.py | 53 ++++++ .../microcosm-build/tests/test_staging_v2.py | 93 +++++++++ .../tests/test_uk_dense_release_assembler.py | 67 +++++-- .../tests/test_uk_rowwise_candidate.py | 176 +++++++++++++++++- tools/assemble_uk_dense_release_dir.py | 28 ++- tools/build_uk_rowwise_candidate.py | 103 ++++++++-- tools/stage_uk_rowwise_candidate.py | 25 ++- 16 files changed, 675 insertions(+), 72 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 6dd458959f..9e34670f65 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -103,7 +103,11 @@ publishes without the flag. Never publish or promote artifacts as a side effect of another task. A UK rowwise run's **staged** bundle (`staged//` in the private repository, written by the build itself) is inspection evidence, not a release: it never moves `releases/` or -`latest.json` and is not loadable through the certified loader. +`latest.json` and is not loadable through the certified loader. The build's +default is to upload that bundle (hundreds of megabytes of licensed microdata) +to the private repository; when you run `tools/build_uk_rowwise_candidate.py` +yourself, pass `--staging-local-only` unless the operator asked for a staged +upload. The US native-SPM-role source-enrichment lane is a separate release type: `tools/build_us_spm_role_enrichment.py` creates a local candidate from the exact diff --git a/changelog.d/uk-dense-assembler-staging-evidence.changed.md b/changelog.d/uk-dense-assembler-staging-evidence.changed.md index bd841f598f..e66c76dbc6 100644 --- a/changelog.d/uk-dense-assembler-staging-evidence.changed.md +++ b/changelog.d/uk-dense-assembler-staging-evidence.changed.md @@ -1 +1 @@ -`tools/assemble_uk_dense_release_dir.py` now requires the run manifest's `staging_delivery` receipt and copies it into `build_manifest.json` as `staging`, the national assembler's rule, so `microcosm-publish-release` applies its undelivered-staging refusal to dense cuts. +`tools/assemble_uk_dense_release_dir.py` now requires the run manifest's `staging_delivery` receipt and copies it into `build_manifest.json` as `staging`, the national assembler's rule, so `microcosm-publish-release` applies its undelivered-staging refusal to dense cuts; `--allow-missing-staging` assembles a run built before the staging lane with a recorded disabled-staging opt-out. diff --git a/changelog.d/uk-staged-dataset-lane.added.md b/changelog.d/uk-staged-dataset-lane.added.md index c7beaee835..4ddf7d659f 100644 --- a/changelog.d/uk-staged-dataset-lane.added.md +++ b/changelog.d/uk-staged-dataset-lane.added.md @@ -1 +1 @@ -Add the staged-dataset lane: the finished rowwise bundle (every manifest-registered output, the manifest, `staged_manifest.json`, `sha256sums.txt`) is uploaded in one commit to `staged//` of the private UK dataset repository, keyed by the telemetry run id, with `staged_dataset` evidence in the manifest and reviewed `staged_dataset`/`fit_summary` artifacts in the telemetry run; `--staged-dataset-repo-id`, `--no-staged-dataset`, `tools/stage_uk_rowwise_candidate.py` and `tools/fetch_uk_staged_dataset.py`. Staging is not publication: `releases/` and `latest.json` are untouched. +Add the staged-dataset lane: the finished rowwise bundle (every manifest-registered output, the manifest, `staged_manifest.json`, `sha256sums.txt`) is uploaded in one commit to `staged//` of the private UK dataset repository, keyed by the telemetry run id, with `staged_dataset` evidence in the manifest and reviewed `staged_dataset`/`fit_summary` artifacts in the telemetry run; `--staged-dataset-repo-id`, `--no-staged-dataset`, `tools/stage_uk_rowwise_candidate.py` and `tools/fetch_uk_staged_dataset.py`. A remote stage is refused up front without a Hub credential that can write the repository (read tokens and fine-grained tokens scoped elsewhere included); epoch telemetry is thinned to a bounded row count for any `--epochs`. Staging is not publication: `releases/` and `latest.json` are untouched. diff --git a/docs/uk-dataset-size-plan-355.md b/docs/uk-dataset-size-plan-355.md index e657f4137e..0b9c932d6f 100644 --- a/docs/uk-dataset-size-plan-355.md +++ b/docs/uk-dataset-size-plan-355.md @@ -101,8 +101,8 @@ run id (the Logbook build id unless `--staging-run-id` is given). A size run stays `releasable=false`; staging does not change that, it makes the candidate inspectable. The run needs an ambient `HF_TOKEN` that can see the private repository, or `--staging-local-only` to keep everything on disk -(`--no-staging` records a deliberate opt-out, `--no-staged-dataset` keeps only -the bundle local). The command ends with a `staged dataset:` line naming the +(`--no-staging` records a deliberate opt-out, `--no-staged-dataset` runs the +telemetry alone). The command ends with a `staged dataset:` line naming the repository, prefix and revision, and the manifest carries `staging_delivery` and `staged_dataset`. Fetch a staged bundle with `tools/fetch_uk_staged_dataset.py --run-id `; re-stage a finished diff --git a/docs/uk-staging-operations.md b/docs/uk-staging-operations.md index b18ef13b9c..4f91a74957 100644 --- a/docs/uk-staging-operations.md +++ b/docs/uk-staging-operations.md @@ -177,11 +177,18 @@ not consulted: a staged bundle is not a release and cannot be loaded as one. Modes follow the staging switch. The default uploads telemetry and the bundle; `--staging-local-only` keeps both on disk (the sidecars are still written); `--no-staging` disables both and records the opt-out; `--no-staged-dataset` -keeps telemetry remote and the bundle local. The repository is -`--staged-dataset-repo-id` (environment `POPULACE_UK_STAGED_DATASET_REPO_ID`). -Because the upload closes a multi-hour run, the command refuses to start a -remote dataset stage without an ambient Hub credential (`HF_TOKEN`) that can -see the repository; the telemetry stays best-effort with no pre-flight. +runs telemetry alone, with the bundle neither inventoried nor uploaded. The +repository is `--staged-dataset-repo-id` (environment +`POPULACE_UK_STAGED_DATASET_REPO_ID`). Because the upload closes a multi-hour +run, the command refuses to start a remote dataset stage without an ambient +Hub credential (`HF_TOKEN`) that can see the repository **and** write it: a +read token, or a fine-grained token scoped to another owner, is refused up +front rather than by the Hub's 403 hours later (a fine-grained token needs +`repo.write` on the repository or on the `policyengine` organisation). The +telemetry stays best-effort with no pre-flight. Forwarded epochs are thinned to +at most 2,400 rows per run, whatever `--epochs` says, so the telemetry files +stay under the contract's 5 MiB cap; a content refusal from the telemetry is +reported once and stops the forwarding without touching the solve. The manifest gains two evidence blocks after the bundle is on disk: `staging_delivery`, the validated version 2 telemetry receipt, and @@ -202,9 +209,11 @@ or whose upload failed, with: uv run python tools/stage_uk_rowwise_candidate.py --run-dir ``` -It is idempotent on the outputs' digests: an identical remote bundle is -`already_staged`; a different one under the same run id is refused -(`REMOTE_DIFFERS`). Fetch a bundle for a scorecard, an evaluation leg or the +It is idempotent on the outputs' digests: a directory whose manifest already +records the same outputs as uploaded is left untouched (the driver's revision +stands); an identical remote bundle the manifest does not know about is +recorded as `already_staged` with the bundle's own commit; a different bundle +under the same run id is refused (`REMOTE_DIFFERS`). Fetch a bundle for a scorecard, an evaluation leg or the dashboard's local-directory mode with: ```bash @@ -217,7 +226,10 @@ or `--staging-local-only` for new runs, and deleting `staged//` on the Hub for a bundle that must not remain (its manifest keeps the record). The dense release assembler applies the national rule to this command's runs: it requires the manifest's `staging_delivery` receipt and copies it into -`build_manifest.json` as `staging`, where publication reads it. +`build_manifest.json` as `staging`, where publication reads it. A dense run +built before this lane carries no receipt; `--allow-missing-staging` assembles +it with a recorded disabled-staging opt-out naming the override, the same +posture publication's `--allow-missing-staging` grants. ## Smoke verification diff --git a/experiments/355-uk-dataset-size-receipts.md b/experiments/355-uk-dataset-size-receipts.md index 73b48521af..628cd3b4ed 100644 --- a/experiments/355-uk-dataset-size-receipts.md +++ b/experiments/355-uk-dataset-size-receipts.md @@ -1018,7 +1018,8 @@ Third attempt: feed `ec20085`, ladder `bed3f13d…`, same arguments; run id and `MeasureResolutionError: provider does not know household.ons_household_type` (550 s, 8.4 GB). The spine-q H5 predates the `frs_relationships` stage (#903, merged 2026-09-11), whose column the national registry on main now resolves. Telemetry closed as `failed` at `surface_resolution` - after recording the compile and clone details; Logbook row and error receipt written. Kept as + after recording the compile and clone details; Logbook row and error receipt written. Run + `uk-local-candidate-f100-s42-20260917T175007Z-2acb7634`, kept as `…-staging-failed-spine-q-stale/`. A spine on main's stages (`spine-r`, code 0cda403a) is built with `834-childcare-tfc/build_twin_passthrough.sh … --staging-local-only` into `data/ukds/acceptance/spine-r-355/`. @@ -1072,13 +1073,13 @@ order), `staging_delivery` and `staged_dataset` in the manifest, and the local ` verifies every listed file after the evidence rewrite. This is the bundle staged on the Hub with `tools/stage_uk_rowwise_candidate.py` (results below). -### Staging r2 on the Hub (2026-09-17 19:1xZ) +### Staging r2 on the Hub (2026-09-17, immediately after r2 closed; 2026-09-18 second attempt) `tools/stage_uk_rowwise_candidate.py --run-dir …/spine-r/f100-k15-h55000-e100-p50-f001-s42-staging` against `policyengine/populace-uk-private`: the repository was reachable (read), the remote prefix was absent, and the single commit was refused by the Hub with 403 "you must use a write -token to upload to a repository". The ambient credential on this machine is the **read** token -`PEHF` (user juaristi22). The lane recorded `status failed`, `error_code UPLOAD_FAILED`, nine +token to upload to a repository". The credential cached on the build machine has the Hub role +`read`. The lane recorded `status failed`, `error_code UPLOAD_FAILED`, nine files with digests, no revision, exit 1, no exception text in the evidence, the bundle and sidecars intact for a re-stage; 11 s. Because a read token also passes the reachability pre-flight (it can see the private repository), the pre-flight and the re-stage tool now refuse a credential whose @@ -1086,3 +1087,14 @@ Hub role is `read` before any work, naming the write requirement. The real uploa credential in María's environment (`HF_TOKEN` or `hf auth login`), then `tools/stage_uk_rowwise_candidate.py --run-dir ` stages this bundle, and the fetch-back and scorecard checks follow. + +Second attempt (2026-09-18, María's shell) with a fine-grained write token: refused again with +403, recorded again as `UPLOAD_FAILED`. The token's `repo.write` scope covered only her own user +namespace, not the `policyengine` organisation that owns the repository, and a fine-grained token +reports the role `fineGrained`, so the read-only pre-flight passed it. The pre-flight and the +re-stage tool now read the token's scopes from the Hub and refuse a credential without +`repo.write` on the repository or its owner; a scope the Hub does not describe is warned about and +proven by the upload. Vahid's review also had the re-stage tool keep the driver's record when the +same outputs are already uploaded (the bundle's own commit, not the repository head), made the epoch +thinning size-aware with a content refusal reported rather than raised from inside the solve, and +gave the dense assembler `--allow-missing-staging` for runs built before the lane. diff --git a/packages/microcosm-build/src/microcosm/build/staging_cli.py b/packages/microcosm-build/src/microcosm/build/staging_cli.py index 834c61a76f..b083a331bb 100644 --- a/packages/microcosm-build/src/microcosm/build/staging_cli.py +++ b/packages/microcosm-build/src/microcosm/build/staging_cli.py @@ -99,8 +99,8 @@ def add_staged_dataset_arguments( The finished bundle follows the staging mode switch: ``--no-staging`` keeps nothing, ``--staging-local-only`` keeps the bundle and its sidecars on disk, and the default uploads it to ``repository`` under - ``staged//``. ``--no-staged-dataset`` keeps telemetry remote while - the bundle stays local. + ``staged//``. ``--no-staged-dataset`` runs telemetry alone: the + bundle is neither inventoried nor uploaded. """ parser.add_argument( @@ -117,8 +117,8 @@ def add_staged_dataset_arguments( "--no-staged-dataset", action="store_true", help=( - "Keep the finished bundle local while staging telemetry still runs; " - "--no-staging already disables both." + "Run staging telemetry alone: the finished bundle is neither " + "inventoried nor uploaded (--no-staging already disables both)." ), ) diff --git a/packages/microcosm-build/src/microcosm/build/staging_dataset.py b/packages/microcosm-build/src/microcosm/build/staging_dataset.py index 3539133a3e..37e7ae97ab 100644 --- a/packages/microcosm-build/src/microcosm/build/staging_dataset.py +++ b/packages/microcosm-build/src/microcosm/build/staging_dataset.py @@ -38,6 +38,7 @@ "StagedDatasetBundle", "StagedDatasetError", "StagedFile", + "delivery_covers_bundle", "disabled_staged_dataset", "fetch_bundle", "local_only_staged_dataset", @@ -458,6 +459,30 @@ def validate_staged_dataset_delivery(payload: Mapping[str, Any]) -> dict[str, An return record +def delivery_covers_bundle( + delivery: Any, bundle: StagedDatasetBundle, *, repository: str | None +) -> bool: + """Whether an existing delivery record already accounts for this bundle. + + True when the record says the same outputs reached ``repository`` (status + ``uploaded`` or ``already_staged``); a re-stage is then a no-op that must + keep the record's revision rather than overwrite it. + """ + + if not isinstance(delivery, Mapping) or repository is None: + return False + try: + record = validate_staged_dataset_delivery(delivery) + except StagedDatasetError: + return False + return ( + record["status"] in ("uploaded", "already_staged") + and record["repository"] == repository + and record["run_id"] == bundle.run_id + and _same_digests(record["files"], bundle.digests()) + ) + + def disabled_staged_dataset(reason: str) -> dict[str, Any]: """Explicit evidence for a deliberate staged-dataset opt-out.""" @@ -567,8 +592,13 @@ def stage_bundle( ) remote_files = remote.get("files") if isinstance(remote, Mapping) else None if _same_digests(remote_files, bundle.digests()): + # The bundle's own commit, not the repository head at re-run time. + try: + revision = storage.last_commit(remote_manifest) or head + except Exception: + revision = head return validate_staged_dataset_delivery( - {**record, "status": "already_staged", "revision": head} + {**record, "status": "already_staged", "revision": revision} ) _warn( f"{remote_prefix} already holds a different bundle; refusing to " @@ -626,6 +656,13 @@ def stage_bundle( def _same_digests(remote: Any, local: Mapping[str, Mapping[str, Any]]) -> bool: + """Compare the outputs' digests only. + + The manifest is deliberately excluded: the local copy gains its evidence + blocks after the upload, so its digest differs from the uploaded copy by + construction while the dataset it describes is unchanged. + """ + if not isinstance(remote, Mapping) or set(remote) != set(local): return False for name, entry in local.items(): diff --git a/packages/microcosm-build/src/microcosm/build/staging_storage.py b/packages/microcosm-build/src/microcosm/build/staging_storage.py index 6f4755fb8a..7f9ffc4ede 100644 --- a/packages/microcosm-build/src/microcosm/build/staging_storage.py +++ b/packages/microcosm-build/src/microcosm/build/staging_storage.py @@ -113,22 +113,93 @@ def file_exists(self, path_in_repo: str) -> bool: ) ) - def credential_role(self) -> str | None: - """The ambient credential's role as the Hub reports it (``read``, ``write``, …). - - ``None`` when the backend cannot say; a fine-grained token reports - ``fineGrained`` and its scopes are only tested by the upload itself. - """ - + def _access_token(self) -> Mapping[str, Any] | None: whoami = getattr(self._api(), "whoami", None) if not callable(whoami): return None info = whoami() auth = info.get("auth") if isinstance(info, Mapping) else None token = auth.get("accessToken") if isinstance(auth, Mapping) else None - role = token.get("role") if isinstance(token, Mapping) else None + return token if isinstance(token, Mapping) else None + + def credential_role(self) -> str | None: + """The ambient credential's role as the Hub reports it. + + ``read`` and ``write`` are the classic roles; a fine-grained token + reports ``fineGrained`` and carries its scopes separately (see + :meth:`credential_can_write`). ``None`` when the backend cannot say. + """ + + 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 ambient credential may write this repository. + + ``True`` for a classic write token or a fine-grained token whose + scopes grant ``repo.write`` on the repository, its owner (user or + organisation) or globally; ``False`` for a read token or a fine-grained + token scoped elsewhere; ``None`` when the backend cannot say, in which + case only the upload itself proves the scope. + """ + + 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 + + def last_commit(self, path_in_repo: str) -> str | None: + """The commit that last changed ``path_in_repo``, when the backend reports it.""" + + paths_info = getattr(self._api(), "get_paths_info", None) + if not callable(paths_info): + return None + entries = 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 head_revision(self) -> str | None: """The default branch's current commit, when the backend reports one.""" @@ -218,7 +289,9 @@ def upload( self.storage.upload(local_path, path_in_repo) except Exception as error: self.consecutive_failures += 1 - became_disabled = self.consecutive_failures >= self.max_consecutive_failures + became_disabled = ( + self.consecutive_failures >= self.max_consecutive_failures + ) if became_disabled: self.enabled = False return UploadResult( diff --git a/packages/microcosm-build/tests/test_staging_dataset.py b/packages/microcosm-build/tests/test_staging_dataset.py index 8a0d7cc1df..c8750acad9 100644 --- a/packages/microcosm-build/tests/test_staging_dataset.py +++ b/packages/microcosm-build/tests/test_staging_dataset.py @@ -399,3 +399,56 @@ def test_refresh_sha256sums_entry_redigests_one_listed_file(tmp_path): } with pytest.raises(StagedDatasetError, match="does not list"): refresh_sha256sums_entry(run_dir, "run.log") + + +class RecordingHub(FakeHub): + """FakeHub that also answers get_paths_info with each file's last commit.""" + + def __init__(self, root: Path, **kwargs) -> None: + super().__init__(root, **kwargs) + self.commit_of: dict[str, str] = {} + + def create_commit(self, **kwargs): + info = super().create_commit(**kwargs) + for operation in kwargs["operations"]: + self.commit_of[operation.path_in_repo] = self.sha + return info + + def get_paths_info(self, *, repo_id, paths, expand, repo_type): + assert expand and repo_type == "dataset" + return [ + SimpleNamespace(path=p, last_commit=SimpleNamespace(oid=self.commit_of[p])) + for p in paths + if p in self.commit_of + ] + + +def test_already_staged_records_the_bundles_own_commit_not_the_head(tmp_path): + from microcosm.build.staging_dataset import delivery_covers_bundle + + run_dir = _write_run(tmp_path) + bundle = StagedDatasetBundle.from_manifest( + run_dir, run_id=RUN_ID, manifest_name=MANIFEST + ) + write_sidecars(bundle, repository="policyengine/populace-uk-private") + hub = RecordingHub(tmp_path) + first = stage_bundle(bundle, storage=_storage(hub)) + staged_at = first["revision"] + # Another commit moves the repository head after the bundle landed. + hub.sha = "f" * 40 + again = stage_bundle(bundle, storage=_storage(hub)) + assert again["status"] == "already_staged" + assert again["revision"] == staged_at != hub.sha + assert delivery_covers_bundle( + first, bundle, repository="policyengine/populace-uk-private" + ) + assert delivery_covers_bundle( + again, bundle, repository="policyengine/populace-uk-private" + ) + assert not delivery_covers_bundle(first, bundle, repository="other/repo") + assert not delivery_covers_bundle( + {**first, "status": "failed", "revision": None, "error_code": "UPLOAD_FAILED"}, + bundle, + repository="policyengine/populace-uk-private", + ) + assert not delivery_covers_bundle(None, bundle, repository="x/y") diff --git a/packages/microcosm-build/tests/test_staging_v2.py b/packages/microcosm-build/tests/test_staging_v2.py index bd8d4193d8..4669082151 100644 --- a/packages/microcosm-build/tests/test_staging_v2.py +++ b/packages/microcosm-build/tests/test_staging_v2.py @@ -903,3 +903,96 @@ def create_commit(self, **kwargs): storage = HuggingFaceDatasetStorage("example/private", api=NoRevision(tmp_path)) with pytest.raises(RuntimeError, match="no revision"): storage.commit([], message="empty") + + +@pytest.mark.parametrize( + ("token", "expected"), + [ + ({"role": "read"}, False), + ({"role": "write"}, True), + ( + { + "role": "fineGrained", + "fineGrained": {"global": ["repo.write"], "scoped": []}, + }, + True, + ), + ( + { + "role": "fineGrained", + "fineGrained": { + "global": [], + "scoped": [ + { + "entity": {"type": "org", "name": "example"}, + "permissions": ["repo.write"], + } + ], + }, + }, + True, + ), + ( + { + "role": "fineGrained", + "fineGrained": { + "global": [], + "scoped": [ + { + "entity": {"type": "dataset", "name": "example/private"}, + "permissions": ["repo.content.read", "repo.write"], + } + ], + }, + }, + True, + ), + ( + { + "role": "fineGrained", + "fineGrained": { + "global": ["discussion.write"], + "scoped": [ + { + "entity": {"type": "user", "name": "someone"}, + "permissions": ["repo.write"], + } + ], + }, + }, + False, + ), + ( + { + "role": "fineGrained", + "fineGrained": { + "global": [], + "scoped": [ + { + "entity": {"type": "org", "name": "example"}, + "permissions": ["repo.content.read"], + } + ], + }, + }, + False, + ), + ({"role": "mystery"}, None), + ], +) +def test_storage_reads_whether_the_credential_can_write_this_repository( + tmp_path, token, expected +): + class TokenApi(CommitApi): + def whoami(self): + return {"name": "x", "auth": {"accessToken": token}} + + storage = HuggingFaceDatasetStorage("example/private", api=TokenApi(tmp_path)) + assert storage.credential_can_write() is expected + # No whoami on the backend: the scope is unknown, never assumed. + assert ( + HuggingFaceDatasetStorage( + "example/private", api=CommitApi(tmp_path) + ).credential_can_write() + is None + ) diff --git a/packages/microcosm-build/tests/test_uk_dense_release_assembler.py b/packages/microcosm-build/tests/test_uk_dense_release_assembler.py index bde4d5bab4..4fdaae9d04 100644 --- a/packages/microcosm-build/tests/test_uk_dense_release_assembler.py +++ b/packages/microcosm-build/tests/test_uk_dense_release_assembler.py @@ -441,6 +441,46 @@ def test_assembler_stages_a_contract_valid_dense_release( assert preflight.main(["--release-dir", str(release_dir)]) == 0 +def test_assembler_override_assembles_a_pre_lane_run_with_a_recorded_opt_out( + tmp_path: Path, monkeypatch, capsys +) -> None: + monkeypatch.setenv("MICROCOSM_UK_TERMINAL_GATE_SIGNING_KEY", KEY) + assembler = _load("assemble_uk_dense_release_dir") + candidate, spine, incumbent = _candidate_dir(tmp_path) + manifest_path = candidate / "rowwise_candidate_manifest.json" + manifest = json.loads(manifest_path.read_text()) + del manifest["staging_delivery"] + manifest_path.write_text(json.dumps(manifest)) + # The incumbent surface evaluation pins the manifest bytes it evaluated. + surface_path = candidate / "incumbent_surface_evaluation.json" + evaluation = json.loads(surface_path.read_text()) + evaluation["identity"]["candidate_manifest_sha256"] = _sha(manifest_path) + surface_path.write_text(json.dumps(evaluation)) + out = tmp_path / "releases" + argv = [ + "--candidate-dir", + str(candidate), + "--spine-h5", + str(spine), + "--incumbent-manifest", + str(incumbent), + "--out-dir", + str(out), + ] + with pytest.raises(SystemExit, match="--allow-missing-staging"): + assembler.main(argv) + assert not out.exists() + assert assembler.main([*argv, "--allow-missing-staging"]) == 0 + capsys.readouterr() + build_manifest = json.loads( + (out / UK_DENSE_RELEASE_ID / "build_manifest.json").read_text() + ) + assert build_manifest["staging"]["mode"] == "disabled" + assert build_manifest["staging"]["enabled"] is False + assert "--allow-missing-staging" in build_manifest["staging"]["opt_out_reason"] + dc.validate_release_dir(out / UK_DENSE_RELEASE_ID) + + @pytest.mark.parametrize( ("staging_delivery", "message"), [ @@ -461,19 +501,22 @@ def test_assembler_refuses_a_run_without_staging_evidence( else: manifest["staging_delivery"] = staging_delivery manifest_path.write_text(json.dumps(manifest)) + argv = [ + "--candidate-dir", + str(candidate), + "--spine-h5", + str(spine), + "--incumbent-manifest", + str(incumbent), + "--out-dir", + str(tmp_path / "releases"), + ] with pytest.raises(SystemExit, match=message): - assembler.main( - [ - "--candidate-dir", - str(candidate), - "--spine-h5", - str(spine), - "--incumbent-manifest", - str(incumbent), - "--out-dir", - str(tmp_path / "releases"), - ] - ) + assembler.main(argv) + if staging_delivery is not None: + # The override covers absent evidence only; invalid evidence stays refused. + with pytest.raises(SystemExit, match=message): + assembler.main([*argv, "--allow-missing-staging"]) assert not (tmp_path / "releases").exists() diff --git a/packages/microcosm-build/tests/test_uk_rowwise_candidate.py b/packages/microcosm-build/tests/test_uk_rowwise_candidate.py index ecb0ba4e9d..ccf0111088 100644 --- a/packages/microcosm-build/tests/test_uk_rowwise_candidate.py +++ b/packages/microcosm-build/tests/test_uk_rowwise_candidate.py @@ -2575,8 +2575,16 @@ def _load_tool(name: str): class _FakeHub: """One fake Hub serving the telemetry repo and the private dataset repo.""" - def __init__(self, *, fail_commit: bool = False, role: str = "write") -> None: + def __init__( + self, + *, + fail_commit: bool = False, + role: str = "write", + scopes: list[dict] | None = None, + ) -> None: self.role = role + self.scopes = scopes + self.commit_of: dict[tuple[str, str], str] = {} self.files: dict[tuple[str, str], bytes] = {} self.uploads: list[tuple[str, str]] = [] self.commits: list[dict[str, object]] = [] @@ -2611,7 +2619,21 @@ def repo_info(self, *, repo_id, repo_type): return SimpleNamespace(sha=self.sha) def whoami(self): - return {"name": "tester", "auth": {"accessToken": {"role": self.role}}} + token = {"role": self.role} + if self.role == "fineGrained": + token["fineGrained"] = {"global": [], "scoped": self.scopes or []} + return {"name": "tester", "auth": {"accessToken": token}} + + def get_paths_info(self, *, repo_id, paths, expand, repo_type): + assert expand and repo_type == "dataset" + return [ + SimpleNamespace( + path=path, + last_commit=SimpleNamespace(oid=self.commit_of[(repo_id, path)]), + ) + for path in paths + if (repo_id, path) in self.commit_of + ] def create_commit( self, *, repo_id, operations, commit_message, repo_type, parent_commit @@ -2631,7 +2653,9 @@ def create_commit( "paths": sorted(op.path_in_repo for op in operations), } ) - self.sha = "b" * 40 + self.sha = hashlib.sha256(commit_message.encode()).hexdigest()[:40] + for operation in operations: + self.commit_of[(repo_id, operation.path_in_repo)] = self.sha return SimpleNamespace(oid=self.sha) @@ -2847,6 +2871,82 @@ def test_size_candidate_stages_the_search_and_refit_phases(monkeypatch, tmp_path assert fit_summary["releasable"] is False +def test_staging_epoch_stride_keeps_the_forwarded_rows_bounded(): + builder = _load_builder_module() + + def stride(epochs, households): + return builder._staging_epoch_every( + SimpleNamespace(epochs=epochs, dataset_households=households) + ) + + assert stride(2, None) == 10 and stride(2, 300) == 10 + assert stride(2000, None) == 10 + # dense + ten probes + refit at 2,000 epochs: 24,000 epochs -> 2,400 rows + assert stride(2000, 55000) == 10 + assert stride(10000, 55000) == 50 + assert stride(100000, None) == 42 + for epochs, households in ((2000, 55000), (10000, 55000), (100000, None)): + solves = 1 if households is None else 2 + builder._BUDGET_ITERS + assert ( + epochs * solves / stride(epochs, households) + <= builder._STAGING_MAX_EPOCH_ROWS + ) + + +def test_telemetry_content_refusal_never_aborts_the_solve( + monkeypatch, tmp_path, capsys +): + from microcosm.build.staging_v2 import StagingContentError, validate_v2_bundle + + builder = _load_builder_module() + input_h5, ladder_path, flags = _staging_run_setup(builder, monkeypatch, tmp_path) + + class Refusing(builder.StagingTelemetryV2): + def calibration_progress(self, event): + raise StagingContentError("Staging file exceeds the 5242880-byte limit.") + + monkeypatch.setattr(builder, "StagingTelemetryV2", Refusing) + out = tmp_path / "refused-rows" + status = builder.main(_build_args(input_h5, ladder_path, flags, out)) + assert status == 0 + err = capsys.readouterr().err + assert err.count("no longer forwarded") == 1 + run_id = _single_run_id(out) + bundle = validate_v2_bundle(out / "staging", run_id) + assert bundle["run_manifest"]["status"] == "completed" + assert not ( + out / "staging" / "runs" / run_id / "calibration_progress.json" + ).exists() + manifest = json.loads((out / builder.MANIFEST_FILENAME).read_text()) + assert manifest["staging_delivery"]["mode"] == "local_only" + assert load_spool_rows(out / "logbook-spool")[0].disposition == "iterating" + + +def test_invalid_local_telemetry_bundle_is_a_warning_not_the_runs_failure( + monkeypatch, tmp_path, capsys +): + from microcosm.build.staging_v2 import StagingContractError + + builder = _load_builder_module() + input_h5, ladder_path, flags = _staging_run_setup(builder, monkeypatch, tmp_path) + + class Invalid(builder.StagingTelemetryV2): + def validate_local_bundle(self): + raise StagingContractError("synthetic bundle defect") + + monkeypatch.setattr(builder, "StagingTelemetryV2", Invalid) + out = tmp_path / "invalid-bundle" + status = builder.main(_build_args(input_h5, ladder_path, flags, out)) + assert status == 0 + err = capsys.readouterr().err + assert "does not validate" in err and "synthetic bundle defect" in err + manifest = json.loads((out / builder.MANIFEST_FILENAME).read_text()) + assert manifest["staging_delivery"]["mode"] == "local_only" + assert manifest["staged_dataset"]["status"] == "skipped" + rows = load_spool_rows(out / "logbook-spool") + assert rows and rows[0].disposition == "iterating" + + def test_no_staging_records_both_opt_outs(monkeypatch, tmp_path): builder = _load_builder_module() input_h5, ladder_path, flags = _staging_run_setup( @@ -2972,20 +3072,47 @@ def test_remote_staging_uploads_telemetry_and_the_bundle_in_one_commit( assert "dataset_staged" in rows[0].phases_reached assert rows[0].disposition == "iterating" - # Re-staging the same directory is a no-op on the output digests. + # Re-staging a directory whose record already says these outputs are + # uploaded touches nothing: the driver's record and revision stand, the + # sidecars keep their bytes, and no commit is made. stager = _load_tool("stage_uk_rowwise_candidate") monkeypatch.setattr(stager, "_hub_api", lambda: hub) + sidecar_bytes = (out / "staged_manifest.json").read_bytes() + capsys.readouterr() assert stager.main(["--run-dir", str(out)]) == 0 + assert "nothing to do" in capsys.readouterr().err restaged = json.loads((out / builder.MANIFEST_FILENAME).read_text())[ "staged_dataset" ] - assert restaged["status"] == "already_staged" + assert restaged == staged + assert (out / "staged_manifest.json").read_bytes() == sidecar_bytes for line in (out / "sha256sums.txt").read_text().splitlines(): digest, name = line.split(" ") assert hashlib.sha256((out / name).read_bytes()).hexdigest() == digest, name - assert restaged["run_id"] == run_id and restaged["revision"] == hub.sha assert len(hub.commits) == 1 + # A record that says the upload failed while the Hub already holds these + # outputs: the re-stage finds the bundle and records its own commit, not + # the repository head, which has moved on since. + manifest_path = out / builder.MANIFEST_FILENAME + manifest = json.loads(manifest_path.read_text()) + manifest["staged_dataset"] = { + **staged, + "status": "failed", + "revision": None, + "error_code": "UPLOAD_FAILED", + } + manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True)) + hub.sha = "e" * 40 + assert stager.main(["--run-dir", str(out)]) == 0 + recovered = json.loads(manifest_path.read_text())["staged_dataset"] + assert recovered["status"] == "already_staged" + assert recovered["revision"] == staged["revision"] != hub.sha + assert len(hub.commits) == 1 + for line in (out / "sha256sums.txt").read_text().splitlines(): + digest, name = line.split(" ") + assert hashlib.sha256((out / name).read_bytes()).hexdigest() == digest, name + # Consumers fetch by run id and get digest-verified local files. fetcher = _load_tool("fetch_uk_staged_dataset") monkeypatch.setattr(fetcher, "_hub_api", lambda: hub) @@ -3108,6 +3235,43 @@ def repo_info(self, **kwargs): builder.main(_build_args(input_h5, ladder_path, flags, out)) assert not out.exists() + # A fine-grained token scoped to another owner is refused the same way; + # one scoped to the repository's organisation passes the pre-flight. + user_scoped = [ + {"entity": {"type": "user", "name": "someone"}, "permissions": ["repo.write"]} + ] + monkeypatch.setattr( + builder, "_hub_api", lambda: _FakeHub(role="fineGrained", scopes=user_scoped) + ) + with pytest.raises(ValueError, match="repo.write"): + builder.main(_build_args(input_h5, ladder_path, flags, out)) + assert not out.exists() + org_scoped = [ + { + "entity": {"type": "org", "name": "policyengine"}, + "permissions": ["repo.write"], + } + ] + org_hub = _FakeHub(role="fineGrained", scopes=org_scoped) + monkeypatch.setattr(builder, "_hub_api", lambda: org_hub) + assert builder.main( + _build_args(input_h5, ladder_path, flags, tmp_path / "org-scoped") + ) in (0, 1) + assert org_hub.commits and org_hub.commits[0]["repo_id"] == ( + "policyengine/populace-uk-private" + ) + + # Argument refusals cost nothing and come first: a missing credential is + # never the reported reason when the arguments are wrong. + monkeypatch.setattr(builder, "_hub_token", lambda: None) + with pytest.raises(ValueError, match="only with --dry-run"): + builder.main( + _build_args( + input_h5, ladder_path, flags, out, "--candidate-clone-counts", "2,3" + ) + ) + assert not out.exists() + # The re-stage tool refuses the same credential the same way. stager = _load_tool("stage_uk_rowwise_candidate") monkeypatch.setattr(stager, "_hub_api", lambda: _FakeHub(role="read")) diff --git a/tools/assemble_uk_dense_release_dir.py b/tools/assemble_uk_dense_release_dir.py index 7d629644a4..8c73c44654 100644 --- a/tools/assemble_uk_dense_release_dir.py +++ b/tools/assemble_uk_dense_release_dir.py @@ -32,6 +32,7 @@ from microcosm.build.staging_v2 import ( StagingContractError, + disabled_staging_delivery, validate_staging_delivery, ) from microcosm.build.uk_runtime.release_identity import UK_DENSE_RELEASE_ID @@ -201,9 +202,22 @@ def _assemble(args: argparse.Namespace) -> dict[str, object]: # whose run intended to stage and delivered nothing. raw_staging_delivery = manifest.get("staging_delivery") if not isinstance(raw_staging_delivery, Mapping): - raise SystemExit( - "error: build record is missing valid staging-delivery evidence" - ) + if not args.allow_missing_staging: + raise SystemExit( + "error: build record is missing valid staging-delivery evidence " + "(a run built before the staging lane needs --allow-missing-staging)" + ) + # An explicit, recorded opt-out mirrors publication's override: the + # build manifest then says why no telemetry exists, instead of + # nothing at all. Invalid evidence is still refused. + manifest = { + **manifest, + "staging_delivery": disabled_staging_delivery( + "assembled with --allow-missing-staging: the run predates " + "staging telemetry" + ), + } + raw_staging_delivery = manifest["staging_delivery"] try: validate_staging_delivery(raw_staging_delivery) except StagingContractError as error: @@ -752,6 +766,14 @@ def _parse_args(argv: list[str] | None) -> argparse.Namespace: help="where the published H5 is cloned (default: beside the candidate)", ) parser.add_argument("--cut-tag") + parser.add_argument( + "--allow-missing-staging", + action="store_true", + help=( + "assemble a run built before staging telemetry existed, recording a " + "disabled-staging opt-out with this reason in build_manifest.json" + ), + ) return parser.parse_args(argv) diff --git a/tools/build_uk_rowwise_candidate.py b/tools/build_uk_rowwise_candidate.py index d3dcb571af..114b8b3ee5 100644 --- a/tools/build_uk_rowwise_candidate.py +++ b/tools/build_uk_rowwise_candidate.py @@ -184,6 +184,9 @@ # which would breach the cap mid-run. Forwarding every tenth epoch and the # last epoch of each phase keeps the loss curve and stays near 0.8 MB. _STAGING_EPOCH_EVERY = 10 +# ...and never more than this many forwarded epochs per run, whatever --epochs +# says: the stride grows with the run so the cap holds by construction. +_STAGING_MAX_EPOCH_ROWS = 2400 _STAGED_DATASET_PHASES = { "uploaded": "dataset_staged", "already_staged": "dataset_staged", @@ -1256,6 +1259,7 @@ def _run_candidate( pool_rows=int(problem.matrix.shape[1]), dataset_households=args.dataset_households, epochs=int(args.epochs), + epoch_every=_staging_epoch_every(args), resumed_from_checkpoint=resume_checkpoint is not None, ) solve = solve_uk_rowwise_weights_under_doctrine( @@ -1284,7 +1288,9 @@ def _run_candidate( progress_events=( None if telemetry is None - else _thinned_epochs(telemetry.calibration_progress) + else _thinned_epochs( + telemetry.calibration_progress, every=_staging_epoch_every(args) + ) ), ) _validate_solve_result(solve, problem=problem) @@ -1605,20 +1611,41 @@ def _preflight_staged_dataset(args: argparse.Namespace) -> None: try: storage.head_revision() except Exception as error: + # The transport's own message is not chained: it can carry request + # URLs and identifiers, and the type name is enough to act on. raise ValueError( f"remote dataset staging cannot reach {repo_id} " f"({type(error).__name__}); {hint}." - ) from error + ) from None + _require_write_credential(storage, hint=hint) + + +def _require_write_credential(storage: HuggingFaceDatasetStorage, *, hint: str) -> None: + """Refuse a credential that can see the repository but cannot write it. + + A read token, or a fine-grained token scoped to another owner, passes the + reachability check and is refused by the Hub with 403 only when the upload + starts, hours later. The scope is read from the Hub's own description of + the token; when it cannot be read the upload itself is the proof. + """ + try: - role = storage.credential_role() + can_write = storage.credential_can_write() except Exception: - role = None - if role == "read": - # A read token sees the private repository, so the reachability check - # passes; the upload at the end of the run would be refused (403). + can_write = None + if can_write is False: raise ValueError( - f"remote dataset staging to {repo_id} needs a write credential; the " - f"ambient Hugging Face token is read-only; {hint}." + f"remote dataset staging to {storage.repo_id} needs a write credential: " + "the ambient Hugging Face token is read-only or is not scoped to this " + "repository or its owner (a fine-grained token needs repo.write on " + f"{storage.repo_id} or on {storage.repo_id.split('/', 1)[0]}); {hint}." + ) + if can_write is None: + print( + "warning: the Hugging Face credential's write scope could not be read; " + "the upload at the end of the run will prove it.", + file=sys.stderr, + flush=True, ) @@ -1655,6 +1682,21 @@ def _stage( telemetry.stage(stage_id, event_status=event_status, **details) +def _staging_epoch_every(args: argparse.Namespace) -> int: + """The epoch stride that keeps the forwarded rows under the row budget. + + A dense run solves once; a size run solves the pool, up to ``budget_iters`` + full-length probes and the refit. The stride is at least + ``_STAGING_EPOCH_EVERY`` and grows so at most ``_STAGING_MAX_EPOCH_ROWS`` + epochs are forwarded, keeping ``calibration_progress.json`` and + ``events.ndjson`` under the contract's 5 MiB cap for any ``--epochs``. + """ + + solves = 1 if args.dataset_households is None else 2 + _BUDGET_ITERS + total = int(args.epochs) * solves + return max(_STAGING_EPOCH_EVERY, -(-total // _STAGING_MAX_EPOCH_ROWS)) + + def _thinned_epochs( sink: Callable[[Mapping[str, Any]], None], *, every: int = _STAGING_EPOCH_EVERY ) -> Callable[[dict[str, object]], None]: @@ -1662,11 +1704,16 @@ def _thinned_epochs( The kernel flags probe epochs with ``budget_search: True``; the staging contract records that field as an integer or null, so the flag becomes 1 - (the national command never runs a budget search and never met this). + (the national command never runs a budget search and never met this). A + contract or content refusal from the telemetry is reported once and stops + the forwarding: the solve must never abort on its own progress report. """ + disabled = False + def callback(event: dict[str, object]) -> None: - if event.get("kind") != "calibration_epoch": + nonlocal disabled + if disabled or event.get("kind") != "calibration_epoch": return epoch = int(event["epoch"]) epochs = int(event["epochs"]) @@ -1676,7 +1723,17 @@ def callback(event: dict[str, object]) -> None: budget_search = forwarded.get("budget_search") if isinstance(budget_search, bool): forwarded["budget_search"] = 1 if budget_search else None - sink(forwarded) + try: + sink(forwarded) + except StagingContractError as error: + disabled = True + print( + "warning: staging telemetry refused a calibration progress row " + f"({type(error).__name__}); epoch progress is no longer forwarded, " + "the solve continues.", + file=sys.stderr, + flush=True, + ) return callback @@ -1717,12 +1774,30 @@ def _finalize_staging_telemetry( ) -> None: if telemetry is None: return - telemetry.complete(message="UK rowwise candidate staging run completed.") + try: + telemetry.complete(message="UK rowwise candidate staging run completed.") + except StagingContractError as error: + _warn_telemetry("could not close the staging run", error) + return try: if args.staging_read_back: + # Requested explicitly, so a failed read-back is the run's failure, + # as on the national command. telemetry.verify_remote() finally: - telemetry.validate_local_bundle() + try: + telemetry.validate_local_bundle() + except StagingContractError as error: + _warn_telemetry("the local staging bundle does not validate", error) + + +def _warn_telemetry(what: str, error: BaseException) -> None: + print( + f"warning: {what} ({type(error).__name__}: {error}); the build's own " + "evidence is unaffected.", + file=sys.stderr, + flush=True, + ) def _staging_delivery(telemetry: StagingTelemetryV2 | None) -> dict[str, Any]: diff --git a/tools/stage_uk_rowwise_candidate.py b/tools/stage_uk_rowwise_candidate.py index f7de7612ac..0dcfa3b809 100644 --- a/tools/stage_uk_rowwise_candidate.py +++ b/tools/stage_uk_rowwise_candidate.py @@ -23,6 +23,7 @@ from microcosm.build.staging_dataset import ( StagedDatasetBundle, + delivery_covers_bundle, local_only_staged_dataset, refresh_sha256sums_entry, stage_bundle, @@ -119,6 +120,18 @@ def main(argv: list[str] | None = None) -> int: "release_posture": manifest.get("release_posture"), }, ) + existing = manifest.get("staged_dataset") + if delivery_covers_bundle(existing, bundle, repository=repository): + # Already on the Hub with these very outputs: keep the driver's record + # (its revision is the bundle's commit) and touch nothing on disk. + print( + f"already staged at {existing['repository']}/{existing['prefix']} " + f"(revision {existing['revision']}); nothing to do.", + file=sys.stderr, + flush=True, + ) + print(json.dumps(existing, indent=2, sort_keys=True)) + return 0 telemetry = manifest.get("staging_delivery") write_sidecars( bundle, @@ -149,13 +162,15 @@ def main(argv: list[str] | None = None) -> int: ) storage = HuggingFaceDatasetStorage(repository, api=_hub_api()) try: - role = storage.credential_role() + can_write = storage.credential_can_write() except Exception: - role = None - if role == "read": + can_write = None + if can_write is False: raise SystemExit( - f"error: the ambient Hugging Face token is read-only; staging to " - f"{repository} needs a write credential (HF_TOKEN or `hf auth login`)." + f"error: the ambient Hugging Face token cannot write {repository}: " + "it is read-only or not scoped to this repository or its owner " + "(a fine-grained token needs repo.write on the repository or on " + f"{repository.split('/', 1)[0]}); set HF_TOKEN or `hf auth login`." ) delivery = stage_bundle( bundle, storage=storage, prefix=UK_STAGED_DATASET_PREFIX