From cb5abcc4026b0586d2b0f5703339df6589862a6d Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Thu, 1 Oct 2026 13:19:24 -0400 Subject: [PATCH 1/3] Keep US staging telemetry on for every build - Exact-count builds stage under their own run prefix and never move latest_staging.json or runs.json, so they no longer need --no-staging (#578 forbids automatic pointer moves, not visibility). The ladder launcher stops passing --no-staging. - The release uploads from a background thread, so a slow or unreachable Hub never stalls the build; the terminal upload is awaited for at most 120 s. Uploads snapshot files under the write lock. - The default upload interval is 300 s, as for UK builds: each cycle is up to six single-file Hub commits against roughly 128 per hour per repository. - SIGTERM raises BuildTerminatedError, so a supervisor or budget stop marks the staging run failed instead of leaving it running. Co-Authored-By: Claude Opus 5.5 --- .../src/microcosm/build/staging.py | 172 ++++++++++++++---- .../tests/engine_free/shared/test_staging.py | 87 +++++++++ .../us/test_us_exact_k_ladder_launcher.py | 37 ++-- .../us/test_us_fiscal_refresh_builder.py | 46 ++++- tools/build_us_exact_k_ladder_release.py | 1 - tools/build_us_fiscal_refresh_release.py | 55 +++++- 6 files changed, 338 insertions(+), 60 deletions(-) diff --git a/packages/microcosm-build/src/microcosm/build/staging.py b/packages/microcosm-build/src/microcosm/build/staging.py index d7ebd59a3c..4aed65de33 100644 --- a/packages/microcosm-build/src/microcosm/build/staging.py +++ b/packages/microcosm-build/src/microcosm/build/staging.py @@ -13,6 +13,7 @@ import math import shutil import sys +import threading import time import traceback from dataclasses import dataclass, field @@ -71,6 +72,14 @@ class StagingTelemetry: api: Optional ``huggingface_hub.HfApi``-shaped object for tests. upload_interval_seconds: Minimum interval between best-effort progress uploads. Final states always upload. + update_pointers: Whether uploads also refresh the repository-level + ``latest_staging.json`` pointer and ``runs.json`` index. A run + that must never move a shared pointer (an exact-count release) + sets this to ``False`` and writes only under its own run prefix. + background_uploads: Upload from a worker thread instead of the build + thread, so a slow or unreachable Hub never stalls the build. The + terminal ``complete``/``fail`` upload is still awaited, for at most + ``final_upload_timeout_seconds``. """ run_id: str @@ -80,6 +89,9 @@ class StagingTelemetry: path_prefix: str = DEFAULT_STAGING_PREFIX api: Any = None upload_interval_seconds: float = 30.0 + update_pointers: bool = True + background_uploads: bool = False + final_upload_timeout_seconds: float = 120.0 started_at: str = field(default_factory=_now) def __post_init__(self) -> None: @@ -105,6 +117,20 @@ def __post_init__(self) -> None: ) self._calibration_events: list[dict[str, Any]] = [] self._artifacts: dict[str, dict[str, Any]] = {} + # Local writes and the uploader's snapshots take this lock, so a + # background upload never ships a half-written file. + self._io_lock = threading.RLock() + self._pending_artifacts: list[tuple[Path, str]] = [] + self._upload_requested = threading.Event() + self._uploads_closing = False + self._upload_thread: threading.Thread | None = None + if self._upload_session is not None and self.background_uploads: + self._upload_thread = threading.Thread( + target=self._upload_worker, + name=f"staging-upload-{self.run_id}", + daemon=True, + ) + self._upload_thread.start() self._progress: dict[str, Any] = { "schema_version": STAGING_SCHEMA_VERSION, "run_id": self.run_id, @@ -166,17 +192,81 @@ def _maybe_upload(self, *, force: bool = False) -> None: if not force and now - self._last_upload_at < self.upload_interval_seconds: return self._last_upload_at = now - for filename in ( - "run_manifest.json", - "progress.json", - "calibration_progress.json", - "events.ndjson", - ): - local = self.run_dir / filename - if local.exists(): - self._upload_file(local, f"{self.repo_run_prefix}/{filename}") - self._upload_latest_pointer() - self._upload_runs_index() + if self._upload_thread is None: + self._upload_cycle() + else: + self._upload_requested.set() + + def _upload_cycle(self) -> None: + with self._io_lock: + pending = [ + (self._upload_copy(local), path_in_repo) + for local, path_in_repo in self._pending_artifacts + if local.exists() + ] + self._pending_artifacts.clear() + for filename in ( + "run_manifest.json", + "progress.json", + "calibration_progress.json", + "events.ndjson", + ): + local = self.run_dir / filename + if local.exists(): + pending.append( + (self._upload_copy(local), f"{self.repo_run_prefix}/{filename}") + ) + for local, path_in_repo in pending: + self._upload_file(local, path_in_repo) + if self.update_pointers: + self._upload_latest_pointer() + self._upload_runs_index() + + def _upload_copy(self, local: Path) -> Path: + """The file to upload: a snapshot when a worker thread uploads it. + + The build thread keeps appending to these files while the worker + uploads, so the worker ships a copy taken under the write lock. + """ + + if self._upload_thread is None: + return local + snapshots = self.run_dir / ".upload" + snapshots.mkdir(exist_ok=True) + copy = snapshots / local.name + shutil.copyfile(local, copy) + return copy + + def _upload_worker(self) -> None: + while True: + self._upload_requested.wait() + self._upload_requested.clear() + try: + self._upload_cycle() + except Exception as error: # pragma: no cover - defensive + print( + f"warning: staging upload cycle failed: {error}", + file=sys.stderr, + ) + if self._uploads_closing and not self._upload_requested.is_set(): + return + + def _finish_uploads(self) -> None: + """Wait for the terminal upload when uploads run in the background.""" + + thread = self._upload_thread + if thread is None: + return + self._uploads_closing = True + self._upload_requested.set() + thread.join(timeout=self.final_upload_timeout_seconds) + if thread.is_alive(): + print( + "warning: the final staging upload did not finish within " + f"{self.final_upload_timeout_seconds:g} s; local staging " + "artifacts are complete.", + file=sys.stderr, + ) def _upload_latest_pointer(self) -> None: payload = { @@ -249,38 +339,41 @@ def run_summary(self) -> dict[str, Any]: } def _write_run_manifest(self) -> None: - _write_json( - self.run_dir / "run_manifest.json", - { - "schema_version": STAGING_SCHEMA_VERSION, - "run_id": self.run_id, - "candidate_release_id": self.candidate_release_id, - "started_at": self.started_at, - "repo_id": self.repo_id, - "path_prefix": self.path_prefix, - "artifacts": self._artifacts, - }, - ) + with self._io_lock: + _write_json( + self.run_dir / "run_manifest.json", + { + "schema_version": STAGING_SCHEMA_VERSION, + "run_id": self.run_id, + "candidate_release_id": self.candidate_release_id, + "started_at": self.started_at, + "repo_id": self.repo_id, + "path_prefix": self.path_prefix, + "artifacts": self._artifacts, + }, + ) def _write_progress(self) -> None: - _write_json(self.run_dir / "progress.json", self._progress) + with self._io_lock: + _write_json(self.run_dir / "progress.json", self._progress) def _write_calibration_progress(self) -> None: - _write_json( - self.run_dir / "calibration_progress.json", - { - "schema_version": STAGING_SCHEMA_VERSION, - "run_id": self.run_id, - "candidate_release_id": self.candidate_release_id, - "updated_at": _now(), - "events": self._calibration_events, - }, - ) + with self._io_lock: + _write_json( + self.run_dir / "calibration_progress.json", + { + "schema_version": STAGING_SCHEMA_VERSION, + "run_id": self.run_id, + "candidate_release_id": self.candidate_release_id, + "updated_at": _now(), + "events": self._calibration_events, + }, + ) def _append_event(self, event: dict[str, Any]) -> None: path = self.run_dir / "events.ndjson" path.parent.mkdir(parents=True, exist_ok=True) - with path.open("a") as stream: + with self._io_lock, path.open("a") as stream: stream.write(json.dumps(_jsonable(event), allow_nan=False) + "\n") def stage( @@ -359,7 +452,12 @@ def attach_artifact( } self._write_run_manifest() if local.exists(): - self._upload_file(local, f"{self.repo_run_prefix}/{source.name}") + path_in_repo = f"{self.repo_run_prefix}/{source.name}" + if self._upload_thread is None: + self._upload_file(local, path_in_repo) + else: + with self._io_lock: + self._pending_artifacts.append((local, path_in_repo)) self._maybe_upload(force=force_upload) def fail(self, error: BaseException) -> None: @@ -371,6 +469,7 @@ def fail(self, error: BaseException) -> None: error_type=type(error).__name__, traceback=traceback.format_exc(), ) + self._finish_uploads() def complete(self) -> None: self.stage( @@ -379,3 +478,4 @@ def complete(self) -> None: message="Staging run completed.", force_upload=True, ) + self._finish_uploads() diff --git a/packages/microcosm-build/tests/engine_free/shared/test_staging.py b/packages/microcosm-build/tests/engine_free/shared/test_staging.py index be9490272d..bc9c59aeca 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_staging.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_staging.py @@ -1,4 +1,5 @@ import json +import threading import microcosm.build.staging as staging_module from microcosm.build.staging import StagingTelemetry @@ -243,3 +244,89 @@ def hf_hub_download(self, **kwargs): assert (run_dir / "events.ndjson").read_text().splitlines() == ( V1_FIXTURE / "events.ndjson" ).read_text().splitlines() + + +def test_pointer_free_runs_upload_only_under_their_prefix(tmp_path) -> None: + api = FakeApi() + telemetry = StagingTelemetry( + run_id="run-k", + candidate_release_id="populace-us-2024-k20000-fixture", + run_dir=tmp_path / "run-k", + repo_id="policyengine/populace-us-staging", + api=api, + upload_interval_seconds=0, + update_pointers=False, + ) + telemetry.stage("target_compilation", force_upload=True) + telemetry.complete() + + uploaded_paths = {upload[1] for upload in api.uploads} + assert "runs/run-k/progress.json" in uploaded_paths + assert all(path.startswith("runs/run-k/") for path in uploaded_paths) + + +class BlockingApi(FakeApi): + """Holds every upload until released, like a slow or unreachable Hub.""" + + def __init__(self) -> None: + super().__init__() + self.release = threading.Event() + + def upload_file(self, **kwargs): + self.release.wait(timeout=10) + return super().upload_file(**kwargs) + + def hf_hub_download(self, **kwargs): + raise FileNotFoundError("no runs index yet") + + +def test_background_uploads_never_block_the_build(tmp_path) -> None: + api = BlockingApi() + telemetry = StagingTelemetry( + run_id="run-bg", + candidate_release_id="populace-us-2024-bg-20260618T000000Z", + run_dir=tmp_path / "run-bg", + repo_id="policyengine/populace-us-staging", + api=api, + upload_interval_seconds=0, + background_uploads=True, + ) + + # The Hub is stuck, yet the build thread keeps going and writes locally. + for stage in ("target_registry", "load_base_frame", "target_compilation"): + telemetry.stage(stage, force_upload=True) + progress = json.loads((tmp_path / "run-bg" / "progress.json").read_text()) + assert progress["stage"] == "target_compilation" + assert api.uploads == [] + + api.release.set() + telemetry.complete() + + # The terminal state is awaited and uploaded, from a consistent snapshot. + uploaded = {path: local for local, path, _, _ in api.uploads} + final = json.loads(open(uploaded["runs/run-bg/progress.json"]).read()) + assert final["status"] == "passed" + events = open(uploaded["runs/run-bg/events.ndjson"]).read().splitlines() + assert json.loads(events[-1])["stage"] == "complete" + assert telemetry.uploads_succeeded > 0 + + +def test_background_final_upload_wait_is_bounded(tmp_path, capsys) -> None: + api = BlockingApi() + telemetry = StagingTelemetry( + run_id="run-stuck", + candidate_release_id="populace-us-2024-stuck-20260618T000000Z", + run_dir=tmp_path / "run-stuck", + repo_id="policyengine/populace-us-staging", + api=api, + upload_interval_seconds=0, + background_uploads=True, + final_upload_timeout_seconds=0.2, + ) + + telemetry.complete() + + assert "did not finish within 0.2 s" in capsys.readouterr().err + progress = json.loads((tmp_path / "run-stuck" / "progress.json").read_text()) + assert progress["status"] == "passed" + api.release.set() diff --git a/packages/microcosm-build/tests/engine_free/us/test_us_exact_k_ladder_launcher.py b/packages/microcosm-build/tests/engine_free/us/test_us_exact_k_ladder_launcher.py index 484e548f53..2f40f43f44 100644 --- a/packages/microcosm-build/tests/engine_free/us/test_us_exact_k_ladder_launcher.py +++ b/packages/microcosm-build/tests/engine_free/us/test_us_exact_k_ladder_launcher.py @@ -250,7 +250,7 @@ def test_launcher_arguments_are_accepted_by_the_house_builder_parser( assert parsed.seed == 17 assert parsed.pool_release_id == "fixture-publication" assert parsed.release_id == "populace-us-2024-k20000-fixture" - assert parsed.no_staging is True + assert parsed.no_staging is False assert parsed.ssi_take_up_prior_weight_basis == tmp_path / "ssi.json" assert parsed.ssi_take_up_prior_weight_basis_sha256 == "1" * 64 @@ -285,7 +285,7 @@ def fake_builder(argv): assert captured[captured.index("--exact-k") + 1] == "N" assert captured[captured.index("--seed") + 1] == "17" - assert "--no-staging" in captured + assert "--no-staging" not in captured assert captured[captured.index("--pool-manifest-sha256") + 1] == _sha256(manifest) assert result["automatic_publish"] is False assert result["release_dir"] == str( @@ -356,17 +356,24 @@ def test_staging_credentials_cannot_enable_pointer_writes( config = launcher._read_config(_write_config(tmp_path, payload)) monkeypatch.setenv("HF_TOKEN", "credentialed-fixture") monkeypatch.setenv("POPULACE_STAGING_REPO_ID", "fixture/staging") - constructed = False + uploads: list[str] = [] - class UnexpectedTelemetry: - def __init__(self, **kwargs): - nonlocal constructed - constructed = True + class RecordingApi: + def upload_file(self, *, path_or_fileobj, path_in_repo, repo_id, repo_type): + uploads.append(path_in_repo) + + def hf_hub_download(self, **kwargs): # pragma: no cover - never reached + raise AssertionError("an exact-count run must not read the runs index") + + real_telemetry = launcher.fiscal_release.StagingTelemetry + + def telemetry_with_recording_api(**kwargs): + return real_telemetry(api=RecordingApi(), **kwargs) monkeypatch.setattr( launcher.fiscal_release, "StagingTelemetry", - UnexpectedTelemetry, + telemetry_with_recording_api, ) argv = launcher._builder_argv( config=config, @@ -381,8 +388,12 @@ def __init__(self, **kwargs): release_root=tmp_path / "out", release_id=config.release_id, ) - - assert parsed.no_staging is True - assert telemetry is None - assert constructed is False - assert not (tmp_path / "out" / "staging").exists() + telemetry.stage("target_registry", force_upload=True) + telemetry.complete() + + # The run is visible, but only under its own prefix: credentials never let + # an exact-count build move latest_staging.json or the runs.json index. + assert uploads + assert all(path.startswith(f"runs/{config.release_id}/") for path in uploads) + assert "latest_staging.json" not in uploads + assert "runs.json" not in uploads diff --git a/packages/microcosm-build/tests/engine_free/us/test_us_fiscal_refresh_builder.py b/packages/microcosm-build/tests/engine_free/us/test_us_fiscal_refresh_builder.py index fe38277f82..59f8997bc5 100644 --- a/packages/microcosm-build/tests/engine_free/us/test_us_fiscal_refresh_builder.py +++ b/packages/microcosm-build/tests/engine_free/us/test_us_fiscal_refresh_builder.py @@ -4,6 +4,7 @@ import inspect import json import os +import signal import sys from dataclasses import replace from pathlib import Path @@ -2035,15 +2036,27 @@ def test_builder_gate_failed_base_pool_override_is_explicit(capsys) -> None: assert "requires --base-h5" in capsys.readouterr().err -def test_builder_exact_k_requires_pointer_suppression(capsys) -> None: +def test_builder_exact_k_stages_without_moving_pointers(tmp_path, monkeypatch) -> None: builder = _load_builder_module() argv = _exact_k_builder_argv("20000") argv.remove("--no-staging") + args = builder._parse_args(argv) + constructed: dict[str, object] = {} - with pytest.raises(SystemExit): - builder._parse_args(argv) + class RecordingTelemetry: + def __init__(self, **kwargs): + constructed.update(kwargs) + + monkeypatch.setattr(builder, "StagingTelemetry", RecordingTelemetry) + telemetry = builder._staging_telemetry( + args, release_root=tmp_path, release_id="populace-us-2024-k20000-fixture" + ) - assert "ExactKPointerSuppressionError" in capsys.readouterr().err + assert isinstance(telemetry, RecordingTelemetry) + # #578: an exact-count release never moves a shared pointer, so its + # telemetry stays under its own run prefix. + assert constructed["update_pointers"] is False + assert constructed["background_uploads"] is True def test_builder_pool_release_identity_is_manifest_authenticated() -> None: @@ -13429,6 +13442,31 @@ def fail(self, error): assert [str(error) for error in recorded] == ["build exploded"] +def test_sigterm_marks_the_staging_run_failed(monkeypatch) -> None: + module = _load_builder_module() + recorded: list[BaseException] = [] + + class Telemetry: + def fail(self, error): + recorded.append(error) + + def terminated(argv=None): + os.kill(os.getpid(), signal.SIGTERM) + raise AssertionError("SIGTERM should have ended the build") + + before = signal.getsignal(signal.SIGTERM) + monkeypatch.setattr(module, "_ACTIVE_TELEMETRY", Telemetry()) + monkeypatch.setattr(module, "_main", terminated) + + with pytest.raises(SystemExit) as raised: + module.main() + + assert raised.value.code == 128 + signal.SIGTERM + assert [type(error).__name__ for error in recorded] == ["BuildTerminatedError"] + assert str(recorded[0]) == "The build was terminated by SIGTERM." + assert signal.getsignal(signal.SIGTERM) is before + + def test_crash_reporting_never_replaces_the_real_traceback(monkeypatch, capsys) -> None: module = _load_builder_module() diff --git a/tools/build_us_exact_k_ladder_release.py b/tools/build_us_exact_k_ladder_release.py index 17668c5e65..87af599237 100644 --- a/tools/build_us_exact_k_ladder_release.py +++ b/tools/build_us_exact_k_ladder_release.py @@ -441,7 +441,6 @@ def _builder_argv( str(config.l2_lambda), "--refit-l2-lambda", str(config.refit_l2_lambda), - "--no-staging", ] if config.ssi_take_up_prior_weight_basis is not None: argv.extend( diff --git a/tools/build_us_fiscal_refresh_release.py b/tools/build_us_fiscal_refresh_release.py index 9c16a01e1c..ca0fc292ca 100644 --- a/tools/build_us_fiscal_refresh_release.py +++ b/tools/build_us_fiscal_refresh_release.py @@ -39,6 +39,7 @@ import platform import re import shutil +import signal import subprocess import sys import time @@ -1827,8 +1828,14 @@ def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser.add_argument( "--staging-upload-interval-seconds", type=float, - default=30.0, - help="Minimum seconds between progress uploads to the staging repo.", + default=300.0, + help=( + "Minimum seconds between progress uploads to the staging repo. " + "Each upload is up to six single-file Hub commits, and the Hub " + "allows roughly 128 commits per hour per repository, so the " + "default matches the UK builds' 300 s. Stage failures and the " + "final state always upload." + ), ) parser.add_argument( _DRY_RUN_REPORT_FLAG, @@ -1963,10 +1970,6 @@ def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: "ExactKExplicitSeedError: --exact-k requires an explicit " "non-negative --seed." ) - if not args.no_staging: - parser.error( - "ExactKPointerSuppressionError: --exact-k requires --no-staging." - ) if not math.isfinite(args.exact_k_pi_hi) or not ( 0.0 <= args.exact_k_pi_hi <= 1.0 ): @@ -12053,6 +12056,12 @@ def _staging_telemetry( repo_id=args.staging_repo_id, path_prefix=args.staging_prefix, upload_interval_seconds=args.staging_upload_interval_seconds, + # An exact-count release never moves a shared pointer (#578): it + # stages under its own run prefix and leaves latest_staging.json and + # runs.json alone. + update_pointers=getattr(args, "exact_k", None) is None, + # A slow or unreachable Hub must not stall the build thread. + background_uploads=True, ) return _ACTIVE_TELEMETRY @@ -12153,6 +12162,7 @@ def main(argv: Sequence[str] | None = None) -> None: global _ACTIVE_DRY_RUN _ACTIVE_DRY_RUN = None + previous_sigterm = _raise_build_terminated_on_sigterm() try: dry_run_exit = _main(argv) except BaseException as error: @@ -12172,10 +12182,43 @@ def main(argv: Sequence[str] | None = None) -> None: file=sys.stderr, ) raise + finally: + if previous_sigterm is not None: + signal.signal(signal.SIGTERM, previous_sigterm) if dry_run_exit is not None: raise SystemExit(dry_run_exit) +class BuildTerminatedError(SystemExit): + """SIGTERM ended the build (a supervisor, a Modal budget stop, ``kill``). + + Raised from the signal handler so ``main`` records the staging run as + failed instead of leaving it ``running``; exits with 128 + SIGTERM. + """ + + def __init__(self) -> None: + super().__init__(128 + signal.SIGTERM) + + def __str__(self) -> str: + return "The build was terminated by SIGTERM." + + +def _raise_build_terminated_on_sigterm(): + """Turn SIGTERM into ``BuildTerminatedError``; return the prior handler. + + Returns ``None`` (and changes nothing) off the main thread, where Python + does not allow installing signal handlers. + """ + + def handle(signum, frame): # noqa: ARG001 - signal handler signature + raise BuildTerminatedError() + + try: + return signal.signal(signal.SIGTERM, handle) + except ValueError: + return None + + def _check_committed_us_ledger_feed_pin( facts_sha256: str, *, From 21172609f7de942a76d4e91e3810859d4ce06a20 Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Thu, 1 Oct 2026 13:26:37 -0400 Subject: [PATCH 2/3] Document staging on every US build Co-Authored-By: Claude Opus 5.5 --- README.md | 6 +++++- changelog.d/us-staging-on-every-build.changed.md | 1 + 2 files changed, 6 insertions(+), 1 deletion(-) create mode 100644 changelog.d/us-staging-on-every-build.changed.md diff --git a/README.md b/README.md index e584a6a021..de110f4cfd 100644 --- a/README.md +++ b/README.md @@ -51,7 +51,11 @@ progress JSON is uploaded to `policyengine/populace-us-staging` while the build runs (best-effort — a missing token or failed upload never fails the build), so every candidate shows up on the staging dashboard before it is published. Disable with `--no-staging`, or point elsewhere with `--staging-repo-id` / -`POPULACE_STAGING_REPO_ID`. An *empty* `POPULACE_STAGING_REPO_ID` is ignored +`POPULACE_STAGING_REPO_ID`. Uploads run on a background thread every 300 s +(`--staging-upload-interval-seconds`), so a slow or unreachable Hub never +stalls the build, and SIGTERM marks the run failed. Exact-count (`--exact-k`) +builds stage too, under their own `runs//` prefix only: they never +move `latest_staging.json` or `runs.json`. An *empty* `POPULACE_STAGING_REPO_ID` is ignored rather than read as off, and staging with no destination at all is an argparse error — `--no-staging` is the only way a build produces no telemetry. diff --git a/changelog.d/us-staging-on-every-build.changed.md b/changelog.d/us-staging-on-every-build.changed.md new file mode 100644 index 0000000000..2a6649b15a --- /dev/null +++ b/changelog.d/us-staging-on-every-build.changed.md @@ -0,0 +1 @@ +US fiscal-refresh staging telemetry is now safe to leave on for every build. Exact-count builds stage under their own `runs//` prefix and never write `latest_staging.json` or `runs.json`, so `--exact-k` no longer requires `--no-staging` and the exact-k ladder launcher stops passing it. The release uploads from a background thread, so a slow or unreachable Hub never stalls the build; the terminal upload is awaited for at most 120 s. `--staging-upload-interval-seconds` defaults to 300 s (was 30 s), as for UK builds, because each cycle is up to six single-file Hub commits. SIGTERM now raises `BuildTerminatedError`, so a supervisor or budget stop marks the staging run failed instead of leaving it `running`. From 09b8170074945ecbaeab3bfc2ce3b49367faf354 Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Thu, 1 Oct 2026 15:47:44 -0400 Subject: [PATCH 3/3] Warn and stay local without a staging write token; add a re-stage tool A US build whose Hugging Face credential cannot write the staging repository is never refused. It checks at the start, prints how to upload the run later, records delivery_check in the run manifest, keeps its telemetry local and builds as usual. tools/restage_us_staging_run.py uploads a finished run folder from any machine with a token; it adds the run to runs.json and never moves latest_staging.json. Co-Authored-By: Claude Opus 5.5 --- README.md | 9 +- .../us-staging-on-every-build.changed.md | 2 +- .../src/microcosm/build/staging.py | 155 ++++++++++++++++++ .../tests/engine_free/shared/test_staging.py | 99 +++++++++++ .../us/test_us_fiscal_refresh_builder.py | 2 + tools/build_us_fiscal_refresh_release.py | 2 + tools/restage_us_staging_run.py | 53 ++++++ 7 files changed, 320 insertions(+), 2 deletions(-) create mode 100644 tools/restage_us_staging_run.py diff --git a/README.md b/README.md index de110f4cfd..3506efbaf0 100644 --- a/README.md +++ b/README.md @@ -55,7 +55,14 @@ Disable with `--no-staging`, or point elsewhere with `--staging-repo-id` / (`--staging-upload-interval-seconds`), so a slow or unreachable Hub never stalls the build, and SIGTERM marks the run failed. Exact-count (`--exact-k`) builds stage too, under their own `runs//` prefix only: they never -move `latest_staging.json` or `runs.json`. An *empty* `POPULACE_STAGING_REPO_ID` is ignored +move `latest_staging.json` or `runs.json`. A build whose Hugging Face token +cannot write the staging repository is never refused: it warns, keeps its +telemetry under `runs//` locally, and builds as usual. Upload that +folder later from any machine with a token: + +```bash +uv run python tools/restage_us_staging_run.py --run-dir /staging/runs/ +``` An *empty* `POPULACE_STAGING_REPO_ID` is ignored rather than read as off, and staging with no destination at all is an argparse error — `--no-staging` is the only way a build produces no telemetry. diff --git a/changelog.d/us-staging-on-every-build.changed.md b/changelog.d/us-staging-on-every-build.changed.md index 2a6649b15a..8ce5c0b1ef 100644 --- a/changelog.d/us-staging-on-every-build.changed.md +++ b/changelog.d/us-staging-on-every-build.changed.md @@ -1 +1 @@ -US fiscal-refresh staging telemetry is now safe to leave on for every build. Exact-count builds stage under their own `runs//` prefix and never write `latest_staging.json` or `runs.json`, so `--exact-k` no longer requires `--no-staging` and the exact-k ladder launcher stops passing it. The release uploads from a background thread, so a slow or unreachable Hub never stalls the build; the terminal upload is awaited for at most 120 s. `--staging-upload-interval-seconds` defaults to 300 s (was 30 s), as for UK builds, because each cycle is up to six single-file Hub commits. SIGTERM now raises `BuildTerminatedError`, so a supervisor or budget stop marks the staging run failed instead of leaving it `running`. +US fiscal-refresh staging telemetry is now safe to leave on for every build. Exact-count builds stage under their own `runs//` prefix and never write `latest_staging.json` or `runs.json`, so `--exact-k` no longer requires `--no-staging` and the exact-k ladder launcher stops passing it. The release uploads from a background thread, so a slow or unreachable Hub never stalls the build; the terminal upload is awaited for at most 120 s. `--staging-upload-interval-seconds` defaults to 300 s (was 30 s), as for UK builds, because each cycle is up to six single-file Hub commits. SIGTERM now raises `BuildTerminatedError`, so a supervisor or budget stop marks the staging run failed instead of leaving it `running`. A build without a token that can write the staging repository is never refused: it warns at the start, keeps its telemetry local, records `delivery_check` in the run manifest, and builds as usual; `tools/restage_us_staging_run.py --run-dir ` uploads the run later from any machine with a token. diff --git a/packages/microcosm-build/src/microcosm/build/staging.py b/packages/microcosm-build/src/microcosm/build/staging.py index 4aed65de33..96125ea280 100644 --- a/packages/microcosm-build/src/microcosm/build/staging.py +++ b/packages/microcosm-build/src/microcosm/build/staging.py @@ -53,6 +53,105 @@ def _jsonable(value: Any) -> Any: return str(value) +def _ambient_hub_token() -> str | None: + """The token huggingface_hub would use (env or ``hf auth login``), if any.""" + + try: + from huggingface_hub.utils import get_token + except ImportError: # pragma: no cover - very old huggingface_hub + return None + try: + return get_token() + except Exception: # pragma: no cover - defensive + return None + + +RUN_FILES = ( + "run_manifest.json", + "progress.json", + "calibration_progress.json", + "events.ndjson", +) + + +def restage_run( + run_dir: Path | str, + *, + repo_id: str, + path_prefix: str = DEFAULT_STAGING_PREFIX, + update_index: bool = True, + api: Any = None, +) -> list[str]: + """Upload a finished local run folder to the staging repository. + + For a run that kept its telemetry local (no write token, a failed + upload, ``--staging-dir`` only). Uploads the run files and the artifacts + its manifest lists under ``//``, and upserts the run + into ``runs.json`` unless ``update_index`` is false. Never moves + ``latest_staging.json``. Returns the repository paths written. + """ + + run_dir = Path(run_dir) + manifest = json.loads((run_dir / "run_manifest.json").read_text()) + run_id = str(manifest["run_id"]) + prefix = path_prefix.strip().strip("/") or DEFAULT_STAGING_PREFIX + storage = HuggingFaceDatasetStorage(repo_id, api=api) + artifacts = manifest.get("artifacts") or {} + names = list(RUN_FILES) + [ + str(artifact["path"]) + for artifact in artifacts.values() + if isinstance(artifact, dict) and artifact.get("path") + ] + written: list[str] = [] + for name in dict.fromkeys(names): + local = run_dir / name + if not local.is_file(): + continue + path_in_repo = f"{prefix}/{run_id}/{name}" + storage.upload(local, path_in_repo) + written.append(path_in_repo) + if update_index and (run_dir / "progress.json").is_file(): + progress = json.loads((run_dir / "progress.json").read_text()) + try: + index = json.loads(storage.download(RUNS_INDEX)) + runs = [ + run + for run in (index.get("runs") or []) + if isinstance(run, dict) and run.get("run_id") != run_id + ] + except Exception: + runs = [] + runs.append( + { + "run_id": run_id, + "candidate_release_id": manifest.get("candidate_release_id"), + "status": progress.get("status"), + "stage": progress.get("stage"), + "started_at": manifest.get("started_at"), + "updated_at": progress.get("updated_at"), + "progress_path": f"{prefix}/{run_id}/progress.json", + "run_manifest_path": f"{prefix}/{run_id}/run_manifest.json", + } + ) + runs.sort( + key=lambda run: str(run.get("updated_at") or run.get("started_at") or ""), + reverse=True, + ) + local_index = run_dir / ".upload" / RUNS_INDEX + local_index.parent.mkdir(exist_ok=True) + _write_json( + local_index, + { + "schema_version": STAGING_SCHEMA_VERSION, + "updated_at": _now(), + "runs": runs, + }, + ) + storage.upload(local_index, RUNS_INDEX) + written.append(RUNS_INDEX) + return written + + def _write_json(path: Path, payload: dict[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(_jsonable(payload), indent=1, allow_nan=False)) @@ -80,6 +179,11 @@ class StagingTelemetry: thread, so a slow or unreachable Hub never stalls the build. The terminal ``complete``/``fail`` upload is still awaited, for at most ``final_upload_timeout_seconds``. + check_write_access: Check at the start that the ambient Hugging Face + credential can write ``repo_id``. Without one, the run warns, keeps + its telemetry local, records why in the run manifest, and builds + as usual; ``restage_run`` uploads the folder later. A check that + cannot reach the Hub changes nothing: uploads stay best-effort. """ run_id: str @@ -92,6 +196,7 @@ class StagingTelemetry: update_pointers: bool = True background_uploads: bool = False final_upload_timeout_seconds: float = 120.0 + check_write_access: bool = False started_at: str = field(default_factory=_now) def __post_init__(self) -> None: @@ -115,6 +220,9 @@ def __post_init__(self) -> None: self._upload_session = ( BestEffortUploadSession(self._storage) if self._storage else None ) + self._delivery_check: dict[str, Any] | None = None + if self._storage is not None and self.check_write_access: + self._delivery_check = self._check_write_access() self._calibration_events: list[dict[str, Any]] = [] self._artifacts: dict[str, dict[str, Any]] = {} # Local writes and the uploader's snapshots take this lock, so a @@ -143,6 +251,48 @@ def __post_init__(self) -> None: self._write_run_manifest() self.stage("created", message="Staging run created.") + def _check_write_access(self) -> dict[str, Any] | None: + """Keep the run local, loudly, when no credential can write the repo. + + Returns the reason recorded in the run manifest, or ``None`` when + uploads go ahead (write access confirmed, or not determinable). + """ + + repo_id = self.repo_id + reason: str | None = None + if self.api is None and _ambient_hub_token() is None: + reason = "no Hugging Face token is configured" + else: + try: + can_write = self._storage.credential_can_write() + except Exception as error: + print( + "warning: could not check write access to the staging " + f"repository {repo_id} ({type(error).__name__}); uploads " + "stay best-effort.", + file=sys.stderr, + ) + return None + if can_write is False: + reason = "the Hugging Face token cannot write this repository" + if reason is None: + return None + self.repo_id = None + self._storage = None + self._upload_session = None + print( + f"warning: staging uploads are off for this run: {reason} " + f"({repo_id}). The build continues and its telemetry stays in " + f"{self.run_dir}. Upload it afterwards with:\n" + f" uv run python tools/restage_us_staging_run.py --run-dir {self.run_dir}", + file=sys.stderr, + ) + return { + "uploads": "local_only", + "repository": repo_id, + "reason": reason, + } + @property def repo_run_prefix(self) -> str: # path_prefix is normalized non-empty at construction, so there is no @@ -350,6 +500,11 @@ def _write_run_manifest(self) -> None: "repo_id": self.repo_id, "path_prefix": self.path_prefix, "artifacts": self._artifacts, + **( + {"delivery_check": self._delivery_check} + if self._delivery_check + else {} + ), }, ) diff --git a/packages/microcosm-build/tests/engine_free/shared/test_staging.py b/packages/microcosm-build/tests/engine_free/shared/test_staging.py index bc9c59aeca..c0c630f1ec 100644 --- a/packages/microcosm-build/tests/engine_free/shared/test_staging.py +++ b/packages/microcosm-build/tests/engine_free/shared/test_staging.py @@ -330,3 +330,102 @@ def test_background_final_upload_wait_is_bounded(tmp_path, capsys) -> None: progress = json.loads((tmp_path / "run-stuck" / "progress.json").read_text()) assert progress["status"] == "passed" api.release.set() + + +class RoleApi(FakeApi): + """A Hub client whose credential has the given role.""" + + def __init__(self, role: str) -> None: + super().__init__() + self.role = role + + def whoami(self): + return {"auth": {"accessToken": {"role": self.role}}} + + def hf_hub_download(self, **kwargs): + raise FileNotFoundError("no runs index yet") + + +def test_a_run_without_write_access_builds_on_and_stays_local(tmp_path, capsys): + api = RoleApi("read") + telemetry = StagingTelemetry( + run_id="run-ro", + candidate_release_id="populace-us-2024-ro", + run_dir=tmp_path / "run-ro", + repo_id="policyengine/populace-us-staging", + api=api, + upload_interval_seconds=0, + check_write_access=True, + ) + telemetry.stage("target_compilation", force_upload=True) + telemetry.complete() + + # Nothing is refused: the run completes with full local telemetry. + progress = json.loads((tmp_path / "run-ro" / "progress.json").read_text()) + assert progress["status"] == "passed" + assert api.uploads == [] + manifest = json.loads((tmp_path / "run-ro" / "run_manifest.json").read_text()) + assert manifest["delivery_check"] == { + "uploads": "local_only", + "repository": "policyengine/populace-us-staging", + "reason": "the Hugging Face token cannot write this repository", + } + assert "restage_us_staging_run.py --run-dir" in capsys.readouterr().err + + +def test_a_run_with_write_access_uploads(tmp_path): + api = RoleApi("write") + telemetry = StagingTelemetry( + run_id="run-rw", + candidate_release_id="populace-us-2024-rw", + run_dir=tmp_path / "run-rw", + repo_id="policyengine/populace-us-staging", + api=api, + upload_interval_seconds=0, + check_write_access=True, + ) + telemetry.complete() + + assert "runs/run-rw/progress.json" in {upload[1] for upload in api.uploads} + manifest = json.loads((tmp_path / "run-rw" / "run_manifest.json").read_text()) + assert "delivery_check" not in manifest + + +def test_restage_uploads_a_local_run_without_moving_the_pointer(tmp_path): + run_dir = tmp_path / "run-local" + telemetry = StagingTelemetry( + run_id="run-local", + candidate_release_id="populace-us-2024-local", + run_dir=run_dir, + ) + (run_dir / "calibration_diagnostics.json").write_text("{}") + telemetry.attach_artifact( + "calibration_diagnostics", run_dir / "calibration_diagnostics.json" + ) + telemetry.complete() + + api = RoleApi("write") + written = staging_module.restage_run( + run_dir, repo_id="policyengine/populace-us-staging", api=api + ) + + assert set(written) == { + "runs/run-local/run_manifest.json", + "runs/run-local/progress.json", + "runs/run-local/events.ndjson", + "runs/run-local/calibration_diagnostics.json", + "runs.json", + } + assert "latest_staging.json" not in {upload[1] for upload in api.uploads} + uploaded = {path_in_repo: local for local, path_in_repo, _, _ in api.uploads} + index = json.loads(open(uploaded["runs.json"]).read()) + assert [run["run_id"] for run in index["runs"]] == ["run-local"] + + api = RoleApi("write") + written = staging_module.restage_run( + run_dir, + repo_id="policyengine/populace-us-staging", + api=api, + update_index=False, + ) + assert "runs.json" not in written diff --git a/packages/microcosm-build/tests/engine_free/us/test_us_fiscal_refresh_builder.py b/packages/microcosm-build/tests/engine_free/us/test_us_fiscal_refresh_builder.py index 59f8997bc5..58d174ac89 100644 --- a/packages/microcosm-build/tests/engine_free/us/test_us_fiscal_refresh_builder.py +++ b/packages/microcosm-build/tests/engine_free/us/test_us_fiscal_refresh_builder.py @@ -2057,6 +2057,8 @@ def __init__(self, **kwargs): # telemetry stays under its own run prefix. assert constructed["update_pointers"] is False assert constructed["background_uploads"] is True + # Without a write token the run warns and stays local; it is never refused. + assert constructed["check_write_access"] is True def test_builder_pool_release_identity_is_manifest_authenticated() -> None: diff --git a/tools/build_us_fiscal_refresh_release.py b/tools/build_us_fiscal_refresh_release.py index ca0fc292ca..6c3b953487 100644 --- a/tools/build_us_fiscal_refresh_release.py +++ b/tools/build_us_fiscal_refresh_release.py @@ -12062,6 +12062,8 @@ def _staging_telemetry( update_pointers=getattr(args, "exact_k", None) is None, # A slow or unreachable Hub must not stall the build thread. background_uploads=True, + # No write token: warn, keep telemetry local, build as usual. + check_write_access=True, ) return _ACTIVE_TELEMETRY diff --git a/tools/restage_us_staging_run.py b/tools/restage_us_staging_run.py new file mode 100644 index 0000000000..877dcd06a8 --- /dev/null +++ b/tools/restage_us_staging_run.py @@ -0,0 +1,53 @@ +"""Upload a finished US staging run folder to the staging repository. + +A build that could not upload (no write token, uploads that failed, or a +``--staging-dir``-only run) keeps its telemetry in +``/staging/runs//``. Run this from any machine with a +token that can write the staging repository to put the run on the staging +dashboard: + + uv run python tools/restage_us_staging_run.py --run-dir + +It uploads the run files and the artifacts the run manifest lists, and adds +the run to ``runs.json``. It never moves ``latest_staging.json``. +""" + +from __future__ import annotations + +import argparse +import os +from pathlib import Path + +from microcosm.build.staging import DEFAULT_STAGING_PREFIX, restage_run + +DEFAULT_REPO_ID = "policyengine/populace-us-staging" + + +def main(argv: list[str] | None = None) -> None: + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument("--run-dir", type=Path, required=True) + parser.add_argument( + "--repo-id", + default=os.environ.get("POPULACE_STAGING_REPO_ID") or DEFAULT_REPO_ID, + help="Staging repository. Defaults to POPULACE_STAGING_REPO_ID or %(default)s.", + ) + parser.add_argument("--path-prefix", default=DEFAULT_STAGING_PREFIX) + parser.add_argument( + "--no-index", + action="store_true", + help="Do not add the run to runs.json (an exact-count run, #578).", + ) + args = parser.parse_args(argv) + if not (args.run_dir / "run_manifest.json").is_file(): + parser.error(f"{args.run_dir} has no run_manifest.json.") + written = restage_run( + args.run_dir, + repo_id=args.repo_id, + path_prefix=args.path_prefix, + update_index=not args.no_index, + ) + print(f"Uploaded {len(written)} file(s) to {args.repo_id}.") + + +if __name__ == "__main__": + main()