From 57d33105231e63b0554ed0b2e8a63fde8bee658f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 22 Sep 2026 23:18:14 -0400 Subject: [PATCH 01/19] Run a US build stage on Modal: plan, image, receipts tools/modal_us_stage.py is a Modal app that runs one stage of the ACS local release tool (materialize, calibrate, qa, finalize, package, or all) off this machine. The image is a shallow clone of the plan's pushed commit synced from that tree's own uv.lock (--frozen), asserted clean so the tool's git identity names the real code vintage. Inputs come from a content-addressed Modal volume or the Hugging Face Hub at an explicit revision and are verified against the plan's sha256 before the tool starts. The tool's state is mirrored to a runs volume and every file in it is listed with its sha256 in a receipt. The default invocation is a cheap check (image, clone, environment, the pinned tool's own argument parser on the built argv, input digests, the run's prior state); --run executes. tools/modal_us_stage_plan.py is the pure, stdlib-only half (plan validation, argv, resource classes sized from the #974 measured peaks, hashing, mirroring, receipts, receipt verification) with unit tests. Co-Authored-By: Claude Opus 5.5 --- docs/us-modal-stage-example-plan.json | 32 + .../tests/test_us_modal_stage_plan_tool.py | 463 +++++++++ tools/modal_us_stage.py | 493 ++++++++++ tools/modal_us_stage_plan.py | 931 ++++++++++++++++++ 4 files changed, 1919 insertions(+) create mode 100644 docs/us-modal-stage-example-plan.json create mode 100644 packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py create mode 100644 tools/modal_us_stage.py create mode 100644 tools/modal_us_stage_plan.py diff --git a/docs/us-modal-stage-example-plan.json b/docs/us-modal-stage-example-plan.json new file mode 100644 index 000000000..a441364df --- /dev/null +++ b/docs/us-modal-stage-example-plan.json @@ -0,0 +1,32 @@ +{ + "schema": "microcosm-modal-us-stage-plan/1", + "tool": "us-acs-local-release", + "stage": "materialize", + "run_id": "replay-974-materialize", + "source": { + "commit": "cadaf418097964038ae539a3d5ce5acedaf49912", + "branch": "local-acs-hours-rebuild-20260921" + }, + "inputs": { + "staging_h5": { + "uri": "volume://cas/sha256/f335f5737537b7efbf05840bfadcf84c74caeef15a3b27a0539e3df42cca1637/acs_multispine_staging.h5", + "sha256": "f335f5737537b7efbf05840bfadcf84c74caeef15a3b27a0539e3df42cca1637" + }, + "staging_summary": { + "uri": "volume://cas/sha256/ebee1ac61811810082936b827a75e6f95ab0db1b139715f1d663bcb812c9e541/acs_multispine_staging.summary.json", + "sha256": "ebee1ac61811810082936b827a75e6f95ab0db1b139715f1d663bcb812c9e541" + }, + "feed": { + "uri": "volume://cas/sha256/4d1dba8c1b6274877bf184fa6de5d99b13fc61f34709ccab1487db2b5c64a79f/consumer_facts.jsonl", + "sha256": "4d1dba8c1b6274877bf184fa6de5d99b13fc61f34709ccab1487db2b5c64a79f" + }, + "ladder": { + "uri": "volume://cas/sha256/39a2ab2abeab07a88362af7ab2940e0e1d50a297c919e4bbc6fb65bab51147d8/us_puma_ladder_2020.npz", + "sha256": "39a2ab2abeab07a88362af7ab2940e0e1d50a297c919e4bbc6fb65bab51147d8" + } + }, + "options": { + "soi_mode": "totals", + "hh_chunk": 20000 + } +} diff --git a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py new file mode 100644 index 000000000..ded601186 --- /dev/null +++ b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py @@ -0,0 +1,463 @@ +"""Unit tests for the pure half of the US Modal stage runner. + +``tools/modal_us_stage_plan.py`` builds the plan, argv and sha256 receipts +that ``tools/modal_us_stage.py`` executes on Modal. Nothing here imports +``modal`` or touches the network. +""" + +from __future__ import annotations + +import copy +import hashlib +import importlib.util +import json +import os +import sys +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[3] + + +def _load(): + path = ROOT / "tools" / "modal_us_stage_plan.py" + spec = importlib.util.spec_from_file_location("modal_us_stage_plan", path) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + # dataclasses resolve string annotations through sys.modules. + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +plan_lib = _load() + +COMMIT = "4d773a4785a1e2c7f0b9d3e6a8c5b1f2e3d4c5b6" +STAGING_SHA = "a" * 64 +SUMMARY_SHA = "b" * 64 +FEED_SHA = "4d1dba8c1b6274877bf184fa6de5d99b13fc61f34709ccab1487db2b5c64a79f" +LADDER_SHA = "c" * 64 + + +def _plan_data(stage: str = "materialize", **overrides) -> dict: + data = { + "schema": plan_lib.PLAN_SCHEMA, + "tool": "us-acs-local-release", + "stage": stage, + "run_id": "acs-local-20260923", + "source": {"commit": COMMIT, "branch": "us-modal-stage-runner"}, + "inputs": { + "staging_h5": { + "uri": f"volume://cas/sha256/{STAGING_SHA}/acs_multispine_staging.h5", + "sha256": STAGING_SHA, + }, + "staging_summary": { + "uri": ( + f"volume://cas/sha256/{SUMMARY_SHA}/" + "acs_multispine_staging.summary.json" + ), + "sha256": SUMMARY_SHA, + }, + "feed": { + "uri": f"volume://cas/sha256/{FEED_SHA}/consumer_facts.jsonl", + "sha256": FEED_SHA, + }, + "ladder": { + "uri": ( + "hf://datasets/policyengine/populace-us@" + "populace-us-2024-spm-receipts-20260923/" + "inputs/us_puma_ladder_2020.npz" + ), + "sha256": LADDER_SHA, + }, + }, + "options": {"soi_mode": "totals", "hh_chunk": 20000}, + } + data.update(overrides) + return data + + +# --------------------------------------------------------------------------- # +# Plan validation # +# --------------------------------------------------------------------------- # + + +def test_valid_plan_parses_with_stage_resources() -> None: + plan = plan_lib.parse_plan(_plan_data()) + assert plan.stage == "materialize" + assert plan.resources is plan_lib.HEAVY + assert plan.repo_url == plan_lib.DEFAULT_REPO_URL + assert set(plan.inputs) == {"staging_h5", "staging_summary", "feed", "ladder"} + ladder = plan.inputs["ladder"] + assert (ladder.kind, ladder.repo_type, ladder.repo_id) == ( + "hf", + "dataset", + "policyengine/populace-us", + ) + assert ladder.revision == "populace-us-2024-spm-receipts-20260923" + assert ladder.path_in_repo == "inputs/us_puma_ladder_2020.npz" + assert plan.inputs["feed"].volume_path.endswith("/consumer_facts.jsonl") + + +@pytest.mark.parametrize( + ("stage", "resources"), + [ + ("materialize", "heavy"), + ("calibrate", "heavy"), + ("all", "heavy"), + ("qa", "light"), + ("finalize", "light"), + ("package", "light"), + ], +) +def test_every_stage_maps_to_a_sized_class(stage: str, resources: str) -> None: + assert plan_lib.parse_plan(_plan_data(stage)).resources.name == resources + + +def test_heavy_class_covers_measured_peaks_with_headroom() -> None: + for (tool, stage), measured in plan_lib.MEASURED.items(): + spec = plan_lib.TOOLS[tool].stages[stage] + request_bytes = spec.resources.memory_mib * 1024 * 1024 + assert request_bytes >= 1.5 * measured.peak_rss_bytes, (stage, measured) + # The state SOI surface reported ~94 GB for materialize (2026-09-22). + assert plan_lib.HEAVY.memory_mib * 1024 * 1024 > 94e9 * 1.3 + + +@pytest.mark.parametrize( + ("mutate", "message"), + [ + (lambda d: d.update(schema="other/1"), "schema"), + (lambda d: d.update(tool="us-nope"), "unknown tool"), + (lambda d: d.update(stage="publish"), "no stage"), + (lambda d: d.update(run_id="Bad Run"), "run_id"), + (lambda d: d["source"].update(commit="4d773a4"), "40-hex"), + (lambda d: d["source"].pop("branch"), "branch"), + (lambda d: d["source"].update(branch="../x"), "branch"), + ( + lambda d: d["source"].update(repo_url="https://evil.example/x.git"), + "repo_url", + ), + (lambda d: d.update(extra=1), "unknown plan keys"), + (lambda d: d["inputs"].pop("feed"), "requires inputs"), + ( + lambda d: d["inputs"].update(puf={"uri": "volume://x", "sha256": "d" * 64}), + "takes no inputs", + ), + (lambda d: d["inputs"]["feed"].update(sha256="ABC"), "64 lowercase hex"), + (lambda d: d["inputs"]["feed"].pop("sha256"), "exactly"), + ( + lambda d: d["inputs"]["feed"].update(uri="s3://bucket/feed.jsonl"), + "unsupported uri", + ), + ( + lambda d: d["inputs"]["feed"].update(uri="volume://../etc/passwd"), + "clean relative", + ), + ( + lambda d: d["inputs"]["feed"].update(uri="volume:///abs/feed.jsonl"), + "clean relative", + ), + ( + lambda d: d["inputs"]["feed"].update( + uri=f"volume://cas/sha256/{'e' * 64}/consumer_facts.jsonl" + ), + "own digest", + ), + ( + lambda d: d["inputs"]["ladder"].update( + uri="hf://datasets/policyengine/populace-us/inputs/x.npz" + ), + "@revision", + ), + ( + lambda d: d["inputs"]["ladder"].update( + uri="hf://spaces/policyengine/populace-us@main/x.npz" + ), + "hf://datasets/", + ), + (lambda d: d["options"].update(allow_dirty=True), "not allowlisted"), + (lambda d: d["options"].update(checkpoint_dir="/tmp"), "not allowlisted"), + (lambda d: d["options"].update(hh_chunk="20000"), "must be int"), + (lambda d: d["options"].update(hh_chunk=True), "must be int"), + (lambda d: d["options"].update(resume="yes"), "must be bool"), + (lambda d: d["options"].update(soi_mode="--out=/x"), "must be str"), + (lambda d: d.update(env={"HF_TOKEN": "x"}), "not allowlisted"), + (lambda d: d.update(env={"MICROCOSM_X": 1}), "must be a string"), + ], +) +def test_loose_or_unsafe_plans_are_refused(mutate, message: str) -> None: + data = copy.deepcopy(_plan_data()) + mutate(data) + with pytest.raises(plan_lib.PlanError, match=message): + plan_lib.parse_plan(data) + + +def test_later_stages_do_not_need_the_feed() -> None: + data = _plan_data("calibrate") + data["inputs"].pop("feed") + plan = plan_lib.parse_plan(data) + assert "--feed" not in plan_lib.planned_argv(plan) + + +def test_plan_digest_is_key_order_independent() -> None: + data = _plan_data() + reordered = json.loads(json.dumps(data, sort_keys=True)) + assert plan_lib.plan_digest(data) == plan_lib.plan_digest(reordered) + changed = _plan_data(run_id="acs-local-20260924") + assert plan_lib.plan_digest(data) != plan_lib.plan_digest(changed) + + +# --------------------------------------------------------------------------- # +# Argv # +# --------------------------------------------------------------------------- # + + +def test_materialize_argv_is_exact() -> None: + plan = plan_lib.parse_plan(_plan_data()) + assert plan_lib.planned_argv(plan) == [ + "/opt/venv/bin/python", + "-B", + "tools/build_us_acs_local_release.py", + "--stage", + "materialize", + "--staging-h5", + "/work/inputs/staging_h5/acs_multispine_staging.h5", + "--staging-summary", + "/work/inputs/staging_summary/acs_multispine_staging.summary.json", + "--ladder", + "/work/inputs/ladder/us_puma_ladder_2020.npz", + "--feed", + "/work/inputs/feed/consumer_facts.jsonl", + "--feed-sha256", + FEED_SHA, + "--checkpoint-dir", + "/work/state/checkpoints", + "--out-h5", + "/work/state/populace_us_2024_acs_local.h5", + "--hh-chunk", + "20000", + "--soi-mode", + "totals", + ] + + +def test_package_argv_adds_release_root_and_bool_flags() -> None: + data = _plan_data("package") + data["options"] = {"resume": False, "allow_partial_geography": True} + argv = plan_lib.planned_argv(plan_lib.parse_plan(data)) + assert argv[argv.index("--out") + 1] == "/work/state/out" + assert "--allow-partial-geography" in argv + assert "--resume" not in argv + assert "--allow-dirty" not in argv + + +def test_argv_paths_are_stable_across_stages() -> None: + # run_identity.json records the staging path and every later stage + # re-verifies the staging digest, so staged paths must not move. + argvs = [ + plan_lib.planned_argv(plan_lib.parse_plan(_plan_data(stage))) + for stage in ("materialize", "calibrate", "qa", "finalize", "package") + ] + for argv in argvs: + for flag in ("--staging-h5", "--checkpoint-dir", "--out-h5", "--ladder"): + assert argv[argv.index(flag) + 1] == argvs[0][argvs[0].index(flag) + 1] + + +def test_no_option_maps_to_a_runner_owned_flag() -> None: + for tool in plan_lib.TOOLS.values(): + flags = {option.flag for option in tool.options.values()} + assert not flags & tool.owned_flags + + +def test_build_stage_argv_requires_every_staged_path() -> None: + plan = plan_lib.parse_plan(_plan_data()) + with pytest.raises(plan_lib.PlanError, match="no staged path"): + plan_lib.build_stage_argv( + plan, python="python", input_paths={"staging_h5": "x"}, state_dir="/s" + ) + + +def test_image_pins_the_commit_and_the_lock() -> None: + plan = plan_lib.parse_plan(_plan_data()) + commands = plan_lib.image_build_commands(plan) + joined = "\n".join(commands) + assert f"fetch -q --depth 1 origin {COMMIT}" in joined + assert f'test "$(git -C /root/microcosm rev-parse HEAD)" = {COMMIT}' in joined + assert "checkout -q -B us-modal-stage-runner FETCH_HEAD" in joined + assert "uv sync --all-packages --extra us --frozen" in joined + # The clean-tree assertion runs after the sync, so the release tool's + # own dirty-tree refusal and recorded sha match the pushed commit. + assert commands[-1].startswith('test -z "$(git') + + +# --------------------------------------------------------------------------- # +# Hashing, mirroring, receipts # +# --------------------------------------------------------------------------- # + + +def test_copy_with_sha256_hashes_in_the_same_pass(tmp_path: Path) -> None: + payload = os.urandom(3 * 1024 * 1024 + 17) + src = tmp_path / "src.bin" + src.write_bytes(payload) + sha, size = plan_lib.copy_with_sha256(src, tmp_path / "a" / "b" / "dst.bin") + assert (sha, size) == (hashlib.sha256(payload).hexdigest(), len(payload)) + assert (tmp_path / "a" / "b" / "dst.bin").read_bytes() == payload + assert not list(tmp_path.rglob("*.partial")) + + +def test_verify_digest_refuses_a_mismatch() -> None: + ref = plan_lib.parse_plan(_plan_data()).inputs["feed"] + plan_lib.verify_digest(ref, FEED_SHA) + with pytest.raises(plan_lib.PlanError, match="plan pins"): + plan_lib.verify_digest(ref, "0" * 64) + + +def test_mirror_actions_copy_changed_and_delete_removed() -> None: + source = {"a": (1, 10), "b": (2, 20), "c": (3, 30)} + destination = {"a": (1, 10), "b": (2, 99), "stale.mmap": (5, 50)} + copy_, delete = plan_lib.mirror_actions(source, destination) + assert copy_ == ["b", "c"] + assert delete == ["stale.mmap"] + + +def test_mirror_tree_round_trip_is_a_no_op(tmp_path: Path) -> None: + src, dst = tmp_path / "src", tmp_path / "dst" + (src / "checkpoints").mkdir(parents=True) + (src / "checkpoints" / "targets.json").write_text("[]") + (src / "artifact.h5").write_bytes(b"h5") + assert plan_lib.mirror_tree(src, dst) == {"copied": 2, "deleted": 0} + assert plan_lib.mirror_tree(src, dst) == {"copied": 0, "deleted": 0} + (src / "artifact.h5").unlink() + assert plan_lib.mirror_tree(src, dst) == {"copied": 0, "deleted": 1} + assert plan_lib.tree_listing(dst).keys() == {"checkpoints/targets.json"} + + +def _receipt(tmp_path: Path, returncode: int = 0) -> tuple[dict, Path]: + state = tmp_path / "state" + (state / "checkpoints").mkdir(parents=True) + (state / "checkpoints" / "run_identity.json").write_text('{"staging_sha256": "x"}') + (state / "populace_us_2024_acs_local.h5").write_bytes(b"\x89HDF" * 100) + data = _plan_data() + plan = plan_lib.parse_plan(data) + receipt = plan_lib.build_receipt( + plan, + data, + argv=plan_lib.planned_argv(plan), + returncode=returncode, + started_at="2026-09-23T03:00:00Z", + finished_at="2026-09-23T04:24:27Z", + wall_seconds=5067.1, + peak_rss_bytes=74_760_110_080, + inputs_verified=[{"name": "feed", "sha256": FEED_SHA, "bytes": 1}], + outputs=plan_lib.hash_tree(state), + git={"head": COMMIT, "tree_clean": True}, + runner={"python": "3.13"}, + ) + return json.loads(json.dumps(receipt)), state + + +def test_receipt_records_plan_source_outputs_and_cost(tmp_path: Path) -> None: + receipt, state = _receipt(tmp_path) + assert receipt["schema"] == plan_lib.RECEIPT_SCHEMA + assert receipt["status"] == "COMPLETED" + assert receipt["plan_sha256"] == plan_lib.plan_digest(_plan_data()) + assert receipt["source"]["commit"] == COMMIT + assert receipt["resources"] == { + "class": "heavy", + "cpu": 4.0, + "memory_mib": 131072, + "timeout_s": 28800, + } + paths = [item["path"] for item in receipt["outputs"]] + assert paths == ["checkpoints/run_identity.json", "populace_us_2024_acs_local.h5"] + h5 = receipt["outputs"][1] + assert h5["sha256"] == hashlib.sha256(b"\x89HDF" * 100).hexdigest() + # 4 cores + 128 GiB at list price for the measured materialize wall. + assert receipt["estimated_usd_at_list_price"] == pytest.approx(1.71, abs=0.01) + assert plan_lib.verify_receipt(receipt, state) == [] + + +def test_failed_stage_receipt_is_marked_failed(tmp_path: Path) -> None: + receipt, _ = _receipt(tmp_path, returncode=1) + assert receipt["status"] == "FAILED" + + +def test_verify_receipt_reports_tampering(tmp_path: Path) -> None: + receipt, state = _receipt(tmp_path) + (state / "populace_us_2024_acs_local.h5").write_bytes(b"\x89HDF" * 99 + b"XXXX") + (state / "checkpoints" / "run_identity.json").unlink() + (state / "extra.json").write_text("{}") + problems = plan_lib.verify_receipt(receipt, state) + assert problems == [ + "missing: checkpoints/run_identity.json", + "sha256 mismatch: populace_us_2024_acs_local.h5", + ] + strict = plan_lib.verify_receipt(receipt, state, strict=True) + assert strict[-1] == "not in receipt: extra.json" + assert plan_lib.verify_receipt({"schema": "x"}, state) == [ + f"not a {plan_lib.RECEIPT_SCHEMA} receipt" + ] + + +def test_prior_state_check_refuses_missing_or_foreign_runs() -> None: + materialize = plan_lib.parse_plan(_plan_data("materialize")) + calibrate = plan_lib.parse_plan(_plan_data("calibrate")) + assert plan_lib.prior_state_problems(materialize, None) == [] + assert "run materialize" in plan_lib.prior_state_problems(calibrate, None)[0] + good = {"staging_sha256": STAGING_SHA, "ladder_sha256": LADDER_SHA} + assert plan_lib.prior_state_problems(calibrate, good) == [] + foreign = {"staging_sha256": "f" * 64, "ladder_sha256": "e" * 64} + problems = plan_lib.prior_state_problems(calibrate, foreign) + assert [p.split(" sha256")[0] for p in problems] == ["staging_h5", "ladder"] + + +def test_cas_path_and_digest_cli(tmp_path: Path, capsys) -> None: + path = tmp_path / "consumer_facts.jsonl" + path.write_text('{"fact": 1}\n') + sha = hashlib.sha256(path.read_bytes()).hexdigest() + assert plan_lib.cas_volume_path(sha, path.name) == ( + f"cas/sha256/{sha}/consumer_facts.jsonl" + ) + assert plan_lib.main(["digest", str(path)]) == 0 + line = json.loads(capsys.readouterr().out) + assert line["input"] == { + "uri": f"volume://cas/sha256/{sha}/consumer_facts.jsonl", + "sha256": sha, + } + assert line["upload"].startswith(f"modal volume put {plan_lib.INPUTS_VOLUME} ") + # The emitted input parses as a plan input. + ref = plan_lib.parse_input("feed", line["input"]) + assert ref.volume_path == f"cas/sha256/{sha}/consumer_facts.jsonl" + + +def test_validate_cli_prints_argv_and_estimate(tmp_path: Path, capsys) -> None: + plan_path = tmp_path / "plan.json" + plan_path.write_text(json.dumps(_plan_data())) + assert plan_lib.main(["validate", str(plan_path)]) == 0 + summary = json.loads(capsys.readouterr().out) + assert summary["resources"]["class"] == "heavy" + assert summary["measured_locally"]["peak_rss_gb"] == 74.8 + assert summary["estimated_usd_at_measured_wall"] == pytest.approx(1.71, abs=0.01) + bad = _plan_data() + bad["source"]["commit"] = "abc" + plan_path.write_text(json.dumps(bad)) + assert plan_lib.main(["validate", str(plan_path)]) == 2 + assert "REFUSED" in capsys.readouterr().err + + +def test_verify_receipt_cli_exit_codes(tmp_path: Path, capsys) -> None: + receipt, state = _receipt(tmp_path) + receipt_path = tmp_path / "receipt.json" + receipt_path.write_text(json.dumps(receipt)) + args = ["verify-receipt", str(receipt_path), "--state-root", str(state)] + assert plan_lib.main(args) == 0 + assert json.loads(capsys.readouterr().out)["verified"] is True + (state / "populace_us_2024_acs_local.h5").write_bytes(b"changed") + assert plan_lib.main(args) == 1 + + +def test_example_plan_in_the_runbook_parses() -> None: + example = ROOT / "docs" / "us-modal-stage-example-plan.json" + plan = plan_lib.parse_plan(json.loads(example.read_text())) + assert plan.tool is plan_lib.US_ACS_LOCAL_RELEASE diff --git a/tools/modal_us_stage.py b/tools/modal_us_stage.py new file mode 100644 index 000000000..3152f965a --- /dev/null +++ b/tools/modal_us_stage.py @@ -0,0 +1,493 @@ +"""Modal app: run one US Microcosm build stage off this machine. + +The plan (``MICROCOSM_MODAL_PLAN``, see ``tools/modal_us_stage_plan.py``) +names one stage of a registered tool, the pushed commit that runs it, the +run it belongs to, and every input by URI and sha256. The image is a +shallow clone of that commit synced from its own ``uv.lock``; inputs are +copied from the content-addressed inputs volume or downloaded from the +Hugging Face Hub at an explicit revision, and verified against the plan's +digests before the tool starts; the tool's state (checkpoints, calibrated +H5, release directory) is mirrored to the runs volume, and every file in it +is listed with its sha256 in a receipt written next to it. + +Check (cheap, the default; builds the image, verifies the clone, the +environment, the tool's own argument parser on the built argv, every input +digest and the run's prior state, without running the stage):: + + MICROCOSM_MODAL_PLAN=plan.json modal run tools/modal_us_stage.py + +Run the stage (paid; sized per stage in the plan module):: + + MICROCOSM_MODAL_PLAN=plan.json modal run --detach tools/modal_us_stage.py --run + +Set ``MICROCOSM_MODAL_HF_SECRET=`` only when an input +lives in a gated or private Hugging Face repo; public releases need no +token, and no secret is attached otherwise. Nothing here uploads to the Hub +or publishes: packaging a release directory is as far as a stage goes. + +Runbook: ``docs/us-modal-stage-runbook.md``. +""" + +from __future__ import annotations + +import json +import os +import resource +import shutil +import subprocess +import sys +import time +from datetime import UTC, datetime +from pathlib import Path + +import modal + +_HERE = Path(__file__).resolve().parent +if str(_HERE) not in sys.path: + sys.path.insert(0, str(_HERE)) + +import modal_us_stage_plan as plan_lib # noqa: E402 + +PLAN_ENV = "MICROCOSM_MODAL_PLAN" +HF_SECRET_ENV = "MICROCOSM_MODAL_HF_SECRET" +APP_NAME = "microcosm-us-stage" +PLAN_MODULE_REMOTE = "/root/modal_us_stage_plan.py" + + +def _local_plan() -> tuple[dict, plan_lib.Plan] | None: + """The plan, read locally when the app is defined; None in the container.""" + + if not modal.is_local(): + return None + path = os.environ.get(PLAN_ENV) + if not path: + raise SystemExit( + f"Set {PLAN_ENV}=; the image is pinned to the plan's commit." + ) + try: + return plan_lib.load_plan(Path(path)) + except plan_lib.PlanError as error: + raise SystemExit(f"REFUSED plan {path}: {error}") from None + + +_LOADED = _local_plan() + + +def _image() -> modal.Image: + image = ( + modal.Image.debian_slim(python_version=plan_lib.IMAGE_PYTHON_VERSION) + .apt_install("git", "build-essential") + # The runner itself (not the build) runs on the image's interpreter; + # it needs the Hub client to stage hf:// inputs. The build runs in + # IMAGE_VENV, synced from the pinned tree's own lock. + .pip_install( + f"uv=={plan_lib.IMAGE_UV_VERSION}", + f"huggingface_hub=={plan_lib.RUNNER_HF_HUB_VERSION}", + ) + .env( + { + "UV_PROJECT_ENVIRONMENT": plan_lib.IMAGE_VENV, + "UV_PYTHON_DOWNLOADS": "never", + "UV_LINK_MODE": "copy", + "PYTHONDONTWRITEBYTECODE": "1", + } + ) + ) + if _LOADED is not None: + _, plan = _LOADED + image = image.run_commands(*plan_lib.image_build_commands(plan)) + image = image.add_local_file( + str(_HERE / "modal_us_stage_plan.py"), PLAN_MODULE_REMOTE + ) + return image + + +def _secrets() -> list[modal.Secret]: + name = os.environ.get(HF_SECRET_ENV) if modal.is_local() else None + return [modal.Secret.from_name(name)] if name else [] + + +app = modal.App(APP_NAME) +image = _image() +inputs_volume = modal.Volume.from_name(plan_lib.INPUTS_VOLUME, create_if_missing=True) +runs_volume = modal.Volume.from_name(plan_lib.RUNS_VOLUME, create_if_missing=True) +VOLUMES = {plan_lib.INPUTS_MOUNT: inputs_volume, plan_lib.RUNS_MOUNT: runs_volume} + + +# --------------------------------------------------------------------------- # +# In-container helpers # +# --------------------------------------------------------------------------- # + + +def _now() -> str: + return datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%SZ") + + +def _git(*parts: str) -> str: + return subprocess.check_output( + ["git", "-C", plan_lib.IMAGE_REPO_ROOT, *parts], text=True + ).strip() + + +def _git_state(plan: plan_lib.Plan) -> dict[str, object]: + head = _git("rev-parse", "HEAD") + clean = _git("status", "--porcelain") == "" + return { + "head": head, + "head_matches_plan": head == plan.commit, + "tree_clean": clean, + "branch_checked_out": _git("branch", "--show-current"), + "tool_present": (Path(plan_lib.IMAGE_REPO_ROOT) / plan.tool.script).is_file(), + } + + +def _runner_identity() -> dict[str, object]: + identity: dict[str, object] = {"python": sys.version.split()[0]} + for name, path in { + "modal_us_stage_plan.py": PLAN_MODULE_REMOTE, + "modal_us_stage.py": __file__, + }.items(): + try: + identity[f"{name}_sha256"] = plan_lib.sha256_file(path)[0] + except OSError: + identity[f"{name}_sha256"] = None + lock = Path(plan_lib.IMAGE_REPO_ROOT) / "uv.lock" + identity["uv_lock_sha256"] = ( + plan_lib.sha256_file(lock)[0] if lock.exists() else None + ) + return identity + + +def _run_dir(plan: plan_lib.Plan) -> Path: + return Path(plan_lib.RUNS_MOUNT) / "runs" / plan.run_id + + +def _hf_download(ref: plan_lib.InputRef, dest_dir: Path) -> Path: + from huggingface_hub import hf_hub_download + + return Path( + hf_hub_download( + repo_id=ref.repo_id, + filename=ref.path_in_repo, + revision=ref.revision, + repo_type=ref.repo_type, + local_dir=str(dest_dir), + ) + ) + + +def _hf_lfs_sha256(ref: plan_lib.InputRef) -> tuple[str | None, int | None]: + """The Hub's recorded LFS sha256 and size, without downloading.""" + + from huggingface_hub import HfApi + + infos = HfApi().get_paths_info( + ref.repo_id, [ref.path_in_repo], revision=ref.revision, repo_type=ref.repo_type + ) + if not infos: + raise FileNotFoundError(f"{ref.uri} not found on the Hub") + info = infos[0] + lfs = getattr(info, "lfs", None) + if lfs is None: + return None, getattr(info, "size", None) + sha = lfs.get("sha256") if isinstance(lfs, dict) else getattr(lfs, "sha256", None) + size = lfs.get("size") if isinstance(lfs, dict) else getattr(lfs, "size", None) + return sha, size + + +def _stage_input(ref: plan_lib.InputRef, work: Path) -> dict[str, object]: + """Copy or download one input to its stable local path; verify its digest.""" + + local = Path(plan_lib.input_local_path(str(work), ref)) + if ref.kind == "volume": + source = Path(plan_lib.INPUTS_MOUNT) / ref.volume_path + if not source.is_file(): + raise plan_lib.PlanError( + f"input {ref.name!r}: {ref.uri} is not on {plan_lib.INPUTS_VOLUME}" + ) + sha, size = plan_lib.copy_with_sha256(source, local) + else: + downloaded = _hf_download(ref, work / "hf" / ref.name) + local.parent.mkdir(parents=True, exist_ok=True) + os.replace(downloaded, local) + sha, size = plan_lib.sha256_file(local) + plan_lib.verify_digest(ref, sha) + return {"name": ref.name, "uri": ref.uri, "sha256": sha, "bytes": size} + + +def _check_input(ref: plan_lib.InputRef, work: Path) -> dict[str, object]: + row: dict[str, object] = {"name": ref.name, "uri": ref.uri, "expected": ref.sha256} + try: + if ref.kind == "volume": + source = Path(plan_lib.INPUTS_MOUNT) / ref.volume_path + if not source.is_file(): + row["problem"] = f"not on {plan_lib.INPUTS_VOLUME}" + return row + sha, size = plan_lib.sha256_file(source) + row["how"] = "hashed on the volume" + else: + sha, size = _hf_lfs_sha256(ref) + row["how"] = "Hub LFS metadata" + if sha is None: # small non-LFS file: download it and hash + path = _hf_download(ref, work / "hf-check" / ref.name) + sha, size = plan_lib.sha256_file(path) + row["how"] = "downloaded and hashed" + row.update(sha256=sha, bytes=size) + if sha != ref.sha256: + row["problem"] = "sha256 mismatch" + except Exception as error: # noqa: BLE001 - reported, not swallowed + row["problem"] = f"{type(error).__name__}: {error}" + return row + + +_PARSE_SNIPPET = """ +import importlib.util, json, sys +spec = importlib.util.spec_from_file_location("stage_tool", sys.argv[1]) +module = importlib.util.module_from_spec(spec) +spec.loader.exec_module(module) +args = getattr(module, sys.argv[2])(sys.argv[3:]) +print(json.dumps({"stages": list(getattr(args, "stages", [])) or None})) +""" + + +def _parse_check(plan: plan_lib.Plan, argv: list[str]) -> dict[str, object]: + """Run the pinned tool's own argument parser on the built argv.""" + + python, _, script, *tool_argv = argv + if plan.tool.parse_function is None: + cmd = [python, "-B", script, "--help"] + else: + cmd = [python, "-B", "-c", _PARSE_SNIPPET, script, plan.tool.parse_function] + cmd += tool_argv + proc = subprocess.run( + cmd, cwd=plan_lib.IMAGE_REPO_ROOT, text=True, capture_output=True, timeout=900 + ) + return { + "returncode": proc.returncode, + "stdout_tail": proc.stdout[-2000:], + "stderr_tail": proc.stderr[-4000:], + } + + +def _load_json(path: Path) -> dict | None: + try: + return json.loads(path.read_text()) + except (OSError, ValueError): + return None + + +# --------------------------------------------------------------------------- # +# Functions # +# --------------------------------------------------------------------------- # + + +@app.function( + image=image, + volumes=VOLUMES, + secrets=_secrets(), + cpu=plan_lib.CHECK.cpu, + memory=plan_lib.CHECK.memory_mib, + timeout=plan_lib.CHECK.timeout_s, + retries=0, +) +def check_stage(plan_data: dict) -> dict: + """Validate image, clone, environment, argv and inputs; run nothing.""" + + plan = plan_lib.parse_plan(plan_data) + work = Path(plan_lib.WORK_ROOT) + report: dict[str, object] = {"summary": plan_lib.summarize(plan)} + problems: list[str] = [] + + git = _git_state(plan) + report["git"] = git + if not git["head_matches_plan"]: + problems.append(f"clone HEAD {git['head']} is not the plan commit") + if not git["tree_clean"]: + problems.append("the image's clone is dirty") + if not git["tool_present"]: + problems.append(f"{plan.tool.script} is not in commit {plan.commit}") + + env_probe = subprocess.run( + [ + f"{plan_lib.IMAGE_VENV}/bin/python", + "-c", + "import json, importlib.metadata as m; import microcosm.build; " + "print(json.dumps({p: m.version(p) for p in " + "('policyengine-us', 'policyengine-core', 'microcosm-build')}))", + ], + cwd=plan_lib.IMAGE_REPO_ROOT, + text=True, + capture_output=True, + timeout=600, + ) + report["environment"] = ( + json.loads(env_probe.stdout) + if env_probe.returncode == 0 + else {"error": env_probe.stderr[-2000:]} + ) + if env_probe.returncode != 0: + problems.append("the synced environment does not import microcosm.build") + + argv = plan_lib.planned_argv(plan) + parse = _parse_check(plan, argv) if git["tool_present"] else {"returncode": None} + report["tool_argument_parse"] = parse + if parse["returncode"] != 0: + problems.append("the pinned tool's parser refuses the built argv") + + inputs = [_check_input(ref, work) for ref in plan.inputs.values()] + report["inputs"] = inputs + problems += [ + f"input {row['name']}: {row['problem']}" for row in inputs if "problem" in row + ] + + runs_volume.reload() + identity = _load_json( + _run_dir(plan) / "state" / "checkpoints" / "run_identity.json" + ) + report["prior_run_identity"] = identity + problems += plan_lib.prior_state_problems(plan, identity) + + report["problems"] = problems + report["ok"] = not problems + return report + + +def _run_stage(plan_data: dict) -> dict: + plan = plan_lib.parse_plan(plan_data) + git = _git_state(plan) + if not (git["head_matches_plan"] and git["tree_clean"] and git["tool_present"]): + raise plan_lib.PlanError(f"image clone does not match the plan: {git}") + + work = Path(plan_lib.WORK_ROOT) + state = work / "state" + run_dir = _run_dir(plan) + runs_volume.reload() + + # 1. Refuse a foreign or missing predecessor before paying for inputs. + identity = _load_json(run_dir / "state" / "checkpoints" / "run_identity.json") + problems = plan_lib.prior_state_problems(plan, identity) + if problems: + raise plan_lib.PlanError("; ".join(problems)) + + # 2. Inputs to stable local paths, each digest verified; then this run's + # prior state (checkpoints, calibrated H5) onto local disk. + inputs_verified = [_stage_input(ref, work) for ref in plan.inputs.values()] + pulled = plan_lib.mirror_tree(run_dir / "state", state) + receipts_dir = run_dir / "receipts" + prior = [] + for path in sorted(receipts_dir.glob("*.json")) if receipts_dir.exists() else []: + prior.append({"file": path.name, "sha256": plan_lib.sha256_file(path)[0]}) + + # 3. The stage itself, in the pinned tree, logged to the state directory. + argv = plan_lib.planned_argv(plan) + started_at, started = _now(), time.time() + log_path = state / "logs" / f"{plan.stage}-{started_at.replace(':', '')}.log" + log_path.parent.mkdir(parents=True, exist_ok=True) + env = {**os.environ, "HF_HUB_OFFLINE": "1", **plan.env} + print(f"$ {' '.join(argv)}", flush=True) + with open(log_path, "w") as log: + proc = subprocess.Popen( + argv, + cwd=plan_lib.IMAGE_REPO_ROOT, + env=env, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + assert proc.stdout is not None + for line in proc.stdout: + log.write(line) + print(line, end="", flush=True) + returncode = proc.wait() + wall = time.time() - started + peak_kib = resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss + finished_at = _now() + + # 4. Outputs: hash the whole state tree, mirror it to the volume, receipt. + outputs = plan_lib.hash_tree(state) + pushed = plan_lib.mirror_tree(state, run_dir / "state") + receipt = plan_lib.build_receipt( + plan, + plan_data, + argv=argv, + returncode=returncode, + started_at=started_at, + finished_at=finished_at, + wall_seconds=wall, + peak_rss_bytes=peak_kib * 1024, + inputs_verified=inputs_verified, + outputs=outputs, + git=git, + runner={**_runner_identity(), "state_pulled": pulled, "state_pushed": pushed}, + prior_receipts=prior, + ) + receipts_dir.mkdir(parents=True, exist_ok=True) + name = f"{plan.stage}-{started_at.replace(':', '')}.json" + (receipts_dir / name).write_text(json.dumps(receipt, indent=2) + "\n") + runs_volume.commit() + receipt["receipt_path"] = f"runs/{plan.run_id}/receipts/{name}" + shutil.rmtree(work / "hf", ignore_errors=True) + return receipt + + +@app.function( + image=image, + volumes=VOLUMES, + secrets=_secrets(), + cpu=plan_lib.HEAVY.cpu, + memory=plan_lib.HEAVY.memory_mib, + timeout=plan_lib.HEAVY.timeout_s, + retries=0, +) +def run_stage_heavy(plan_data: dict) -> dict: + return _run_stage(plan_data) + + +@app.function( + image=image, + volumes=VOLUMES, + secrets=_secrets(), + cpu=plan_lib.LIGHT.cpu, + memory=plan_lib.LIGHT.memory_mib, + timeout=plan_lib.LIGHT.timeout_s, + retries=0, +) +def run_stage_light(plan_data: dict) -> dict: + return _run_stage(plan_data) + + +RUNNERS = {"heavy": run_stage_heavy, "light": run_stage_light} + + +@app.local_entrypoint() +def main(run: bool = False) -> None: + """Check the plan (default) or, with --run, execute the stage.""" + + assert _LOADED is not None + plan_data, plan = _LOADED + if not run: + report = check_stage.remote(plan_data) + print(json.dumps(report, indent=2)) + if not report["ok"]: + raise SystemExit(f"CHECK FAILED: {report['problems']}") + print("CHECK OK") + return + runner = RUNNERS[plan.resources.name] + receipt = runner.remote(plan_data) + brief = { + key: receipt[key] + for key in ( + "status", + "stage", + "run_id", + "wall_seconds", + "peak_rss_bytes", + "returncode", + "estimated_usd_at_list_price", + "receipt_path", + ) + } + brief["outputs"] = len(receipt["outputs"]) + print(json.dumps(brief, indent=2)) + if receipt["returncode"] != 0: + raise SystemExit(f"STAGE FAILED (receipt {receipt['receipt_path']})") diff --git a/tools/modal_us_stage_plan.py b/tools/modal_us_stage_plan.py new file mode 100644 index 000000000..302c3a884 --- /dev/null +++ b/tools/modal_us_stage_plan.py @@ -0,0 +1,931 @@ +#!/usr/bin/env python3 +"""Plans, argv and sha256 receipts for running a US build stage on Modal. + +The Modal app is ``tools/modal_us_stage.py``; this module is its pure half. +It uses only the standard library, so the Modal client's own interpreter, the +container and the unit tests all import it, and nothing here imports +``modal``, touches the network or runs a build. + +A *plan* is a small JSON file naming one stage of one registered tool, the +pushed commit whose tree runs it, the run it belongs to, and every input by +URI and sha256. The runner refuses a plan that pins anything loosely: the +commit is a full 40-hex sha, every input carries its digest, options and +environment overrides come from allowlists, and the flags the runner owns +(input paths, checkpoint and output locations) cannot be passed through. + +Command line (no Modal needed):: + + python3 tools/modal_us_stage_plan.py validate PLAN.json + python3 tools/modal_us_stage_plan.py digest FILE [FILE ...] + python3 tools/modal_us_stage_plan.py verify-receipt RECEIPT.json \\ + --state-root DIR + +See ``docs/us-modal-stage-runbook.md``. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import re +import shlex +import shutil +import sys +from collections.abc import Callable, Iterable, Mapping, Sequence +from dataclasses import dataclass, field +from pathlib import Path, PurePosixPath + +PLAN_SCHEMA = "microcosm-modal-us-stage-plan/1" +RECEIPT_SCHEMA = "microcosm-modal-us-stage-receipt/1" +DEFAULT_REPO_URL = "https://github.com/PolicyEngine/microcosm.git" + +# Named Modal volumes. Inputs live content-addressed under ``cas/sha256/``; +# each run keeps its tool state and receipts under ``runs//``. +INPUTS_VOLUME = "microcosm-us-stage-inputs" +RUNS_VOLUME = "microcosm-us-stage-runs" +INPUTS_MOUNT = "/vol/inputs" +RUNS_MOUNT = "/vol/runs" + +# Container layout. The build tree is cloned at the plan's commit into +# IMAGE_REPO_ROOT and synced into IMAGE_VENV from its own uv.lock; the stage +# works on ephemeral local disk under WORK_ROOT and mirrors its state to the +# runs volume when it finishes. +IMAGE_REPO_ROOT = "/root/microcosm" +IMAGE_VENV = "/opt/venv" +# Local US builds run on 3.14 (runtime.python 3.14.4 in the #974 build +# manifest); the image matches so a Modal replay differs only by platform. +IMAGE_PYTHON_VERSION = "3.14" +IMAGE_UV_VERSION = "0.11.7" +RUNNER_HF_HUB_VERSION = "1.18.0" +WORK_ROOT = "/work" + +# Modal list prices for standard (non-sandbox) compute, read from +# https://modal.com/pricing on 2026-09-22. Modal bills the higher of the +# request and actual use, so an estimate from the request is a floor only +# when the stage stays inside it. +CPU_USD_PER_CORE_SECOND = 0.0000131 +MEMORY_USD_PER_GIB_SECOND = 0.00000222 + +_SHA256 = re.compile(r"^[0-9a-f]{64}$") +_COMMIT = re.compile(r"^[0-9a-f]{40}$") +_RUN_ID = re.compile(r"^[a-z0-9][a-z0-9._-]{2,79}$") +_BRANCH = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._/-]{0,99}$") +_FILENAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,199}$") +_HF_REPO_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$") +_HF_REVISION = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,199}$") +_ENV_KEY = re.compile( + r"^(?:(?:MICROCOSM|POPULACE)_[A-Z0-9_]+|(?:OMP|MKL|OPENBLAS|NUMEXPR)_NUM_THREADS)$" +) +_GITHUB_URL = re.compile(r"^https://github\.com/[A-Za-z0-9._-]+/[A-Za-z0-9._-]+\.git$") + + +class PlanError(ValueError): + """A plan the runner refuses to execute.""" + + +# --------------------------------------------------------------------------- # +# Resources # +# --------------------------------------------------------------------------- # + + +@dataclass(frozen=True) +class Resources: + """One Modal resource class: a fixed request per decorated function.""" + + name: str + cpu: float + memory_mib: int + timeout_s: int + + @property + def memory_gib(self) -> float: + return self.memory_mib / 1024 + + def estimated_usd(self, wall_seconds: float) -> float: + """List-price cost of holding this request for ``wall_seconds``.""" + + per_second = ( + self.cpu * CPU_USD_PER_CORE_SECOND + + self.memory_gib * MEMORY_USD_PER_GIB_SECOND + ) + return round(per_second * wall_seconds, 2) + + +# Sized from measured local peaks (see MEASURED below): the heavy class holds +# the ACS materialize stage on the state SOI surface (~94 GB reported on +# 2026-09-22, 74.8 GB measured on the totals surface) with headroom; the +# engine pass is single-threaded (4,990 CPU-s over 5,067 wall-s), so four +# cores are for hashing and BLAS, not the microsimulation. +CHECK = Resources("check", cpu=2.0, memory_mib=8 * 1024, timeout_s=30 * 60) +HEAVY = Resources("heavy", cpu=4.0, memory_mib=128 * 1024, timeout_s=8 * 3600) +LIGHT = Resources("light", cpu=2.0, memory_mib=48 * 1024, timeout_s=4 * 3600) +RESOURCE_CLASSES = {item.name: item for item in (CHECK, HEAVY, LIGHT)} + + +@dataclass(frozen=True) +class Measured: + """A local measurement a resource class was sized from.""" + + peak_rss_bytes: int + wall_seconds: float + cpu_seconds: float + source: str + + +_ACS_974 = ( + "experiments/us-acs-local-hours-rebuild-20260922/" + "run-resources-and-staging-excerpt.json (#974, totals SOI surface)" +) +MEASURED = { + ("us-acs-local-release", "materialize"): Measured( + 74_760_110_080, 5067.1, 4990.4, _ACS_974 + ), + ("us-acs-local-release", "calibrate"): Measured( + 67_617_161_216, 415.1, 471.9, _ACS_974 + ), + ("us-acs-local-release", "qa"): Measured(21_867_954_176, 172.2, 168.1, _ACS_974), + ("us-acs-local-release", "finalize"): Measured( + 23_333_388_288, 78.9, 74.6, _ACS_974 + ), + ("us-acs-local-release", "package"): Measured(21_777_367_040, 78.9, 73.7, _ACS_974), +} + + +# --------------------------------------------------------------------------- # +# Tool registry # +# --------------------------------------------------------------------------- # + + +@dataclass(frozen=True) +class OptionFlag: + flag: str + kind: type + + +@dataclass(frozen=True) +class StageSpec: + name: str + resources: Resources + required_inputs: tuple[str, ...] + + +@dataclass(frozen=True) +class ToolSpec: + name: str + script: str + inputs: tuple[str, ...] + stages: Mapping[str, StageSpec] + options: Mapping[str, OptionFlag] + owned_flags: frozenset[str] + argv_builder: Callable[[Plan, Mapping[str, str], str], list[str]] + # Name of the tool's own argparse entry point, called by the check to + # validate the built argv against the pinned commit without running. + parse_function: str | None = None + + +ACS_LOCAL_RELEASE_ARTIFACT = "populace_us_2024_acs_local.h5" + + +def _acs_local_release_argv( + plan: Plan, input_paths: Mapping[str, str], state_dir: str +) -> list[str]: + state = PurePosixPath(state_dir) + argv = [ + "--stage", + plan.stage, + "--staging-h5", + input_paths["staging_h5"], + "--staging-summary", + input_paths["staging_summary"], + "--ladder", + input_paths["ladder"], + ] + if "feed" in plan.inputs: + argv += [ + "--feed", + input_paths["feed"], + "--feed-sha256", + plan.inputs["feed"].sha256, + ] + argv += [ + "--checkpoint-dir", + str(state / "checkpoints"), + "--out-h5", + str(state / ACS_LOCAL_RELEASE_ARTIFACT), + ] + if plan.stage in {"package", "all"}: + argv += ["--out", str(state / "out")] + return argv + + +_ACS_BASE_INPUTS = ("staging_h5", "staging_summary", "ladder") +US_ACS_LOCAL_RELEASE = ToolSpec( + name="us-acs-local-release", + script="tools/build_us_acs_local_release.py", + inputs=("staging_h5", "staging_summary", "feed", "ladder"), + stages={ + "materialize": StageSpec("materialize", HEAVY, (*_ACS_BASE_INPUTS, "feed")), + "calibrate": StageSpec("calibrate", HEAVY, _ACS_BASE_INPUTS), + "qa": StageSpec("qa", LIGHT, _ACS_BASE_INPUTS), + "finalize": StageSpec("finalize", LIGHT, _ACS_BASE_INPUTS), + "package": StageSpec("package", LIGHT, _ACS_BASE_INPUTS), + "all": StageSpec("all", HEAVY, (*_ACS_BASE_INPUTS, "feed")), + }, + options={ + "families": OptionFlag("--families", str), + "geographies": OptionFlag("--geographies", str), + "soi_mode": OptionFlag("--soi-mode", str), + "epochs": OptionFlag("--epochs", int), + "epoch_batch": OptionFlag("--epoch-batch", int), + "max_weight_ratio": OptionFlag("--max-weight-ratio", float), + "target_loss_cap": OptionFlag("--target-loss-cap", float), + "l2_lambda": OptionFlag("--l2-lambda", float), + "seed": OptionFlag("--seed", int), + "resume": OptionFlag("--resume", bool), + "batch": OptionFlag("--batch", int), + "hh_chunk": OptionFlag("--hh-chunk", int), + "allow_partial_geography": OptionFlag("--allow-partial-geography", bool), + }, + owned_flags=frozenset( + { + "--stage", + "--staging-h5", + "--staging-summary", + "--feed", + "--feed-sha256", + "--ladder", + "--checkpoint-dir", + "--out-h5", + "--out-summary", + "--gate-report", + "--out", + # A Modal run always builds a clean clone of a pushed commit. + "--allow-dirty", + } + ), + argv_builder=_acs_local_release_argv, + parse_function="_parse_args", +) + +TOOLS: dict[str, ToolSpec] = {US_ACS_LOCAL_RELEASE.name: US_ACS_LOCAL_RELEASE} + + +# --------------------------------------------------------------------------- # +# Inputs # +# --------------------------------------------------------------------------- # + + +def cas_volume_path(sha256: str, filename: str) -> str: + """Content-addressed path of an uploaded input inside INPUTS_VOLUME.""" + + _require_sha256(sha256, "cas digest") + _require_filename(filename) + return f"cas/sha256/{sha256}/{filename}" + + +@dataclass(frozen=True) +class InputRef: + name: str + uri: str + sha256: str + kind: str + filename: str + volume_path: str | None = None + repo_type: str | None = None + repo_id: str | None = None + revision: str | None = None + path_in_repo: str | None = None + + def to_json(self) -> dict[str, str]: + return {"uri": self.uri, "sha256": self.sha256} + + +def _require_sha256(value: object, what: str) -> str: + if not isinstance(value, str) or not _SHA256.match(value): + raise PlanError(f"{what}: expected 64 lowercase hex characters, got {value!r}") + return value + + +def _require_filename(name: str) -> str: + if not _FILENAME.match(name): + raise PlanError(f"unsafe file name {name!r}") + return name + + +def _safe_relative_path(text: str, what: str) -> PurePosixPath: + path = PurePosixPath(text) + if ( + not text + or path.is_absolute() + or any(part in {"", ".", ".."} for part in text.split("/")) + ): + raise PlanError(f"{what}: {text!r} must be a clean relative path") + for part in path.parts: + _require_filename(part) + return path + + +def parse_input(name: str, spec: object) -> InputRef: + """Parse one plan input: ``{"uri": ..., "sha256": ...}``. + + ``volume://`` names a file in INPUTS_VOLUME; a path under + ``cas/sha256//`` must carry the input's own digest. + ``hf:////@/`` names a file + in a Hugging Face repo at an explicit revision (a public repo needs no + token). The sha256 is verified after staging either way. + """ + + if not isinstance(spec, Mapping) or set(spec) != {"uri", "sha256"}: + raise PlanError(f"input {name!r}: expected exactly {{'uri', 'sha256'}}") + uri = spec["uri"] + sha256 = _require_sha256(spec["sha256"], f"input {name!r} sha256") + if not isinstance(uri, str): + raise PlanError(f"input {name!r}: uri must be a string") + if uri.startswith("volume://"): + rel = _safe_relative_path(uri.removeprefix("volume://"), f"input {name!r}") + parts = rel.parts + if parts[:2] == ("cas", "sha256"): + if len(parts) != 4 or parts[2] != sha256: + raise PlanError( + f"input {name!r}: content-addressed path {uri!r} does not " + f"name its own digest {sha256}" + ) + return InputRef( + name=name, + uri=uri, + sha256=sha256, + kind="volume", + filename=rel.name, + volume_path=str(rel), + ) + if uri.startswith("hf://"): + body = uri.removeprefix("hf://") + repo_type_part, _, rest = body.partition("/") + repo_type = {"datasets": "dataset", "models": "model"}.get(repo_type_part) + if repo_type is None: + raise PlanError( + f"input {name!r}: hf URI must start hf://datasets/ or hf://models/" + ) + org, _, rest = rest.partition("/") + name_and_rev, _, path_in_repo = rest.partition("/") + repo_name, at, revision = name_and_rev.partition("@") + repo_id = f"{org}/{repo_name}" + if not at or not _HF_REVISION.match(revision): + raise PlanError( + f"input {name!r}: hf URI needs an explicit @revision " + "(tag, branch without '/', or commit)" + ) + if not _HF_REPO_ID.match(repo_id): + raise PlanError(f"input {name!r}: bad Hugging Face repo id {repo_id!r}") + rel = _safe_relative_path(path_in_repo, f"input {name!r} path") + return InputRef( + name=name, + uri=uri, + sha256=sha256, + kind="hf", + filename=rel.name, + repo_type=repo_type, + repo_id=repo_id, + revision=revision, + path_in_repo=str(rel), + ) + raise PlanError(f"input {name!r}: unsupported uri {uri!r} (volume:// or hf://)") + + +def input_local_path(work_root: str, ref: InputRef) -> str: + """Where the container stages an input; stable across a run's stages.""" + + return str(PurePosixPath(work_root) / "inputs" / ref.name / ref.filename) + + +# --------------------------------------------------------------------------- # +# Plans # +# --------------------------------------------------------------------------- # + + +@dataclass(frozen=True) +class Plan: + tool: ToolSpec + stage: str + run_id: str + repo_url: str + commit: str + branch: str + inputs: Mapping[str, InputRef] + options: Mapping[str, object] = field(default_factory=dict) + env: Mapping[str, str] = field(default_factory=dict) + + @property + def stage_spec(self) -> StageSpec: + return self.tool.stages[self.stage] + + @property + def resources(self) -> Resources: + return self.stage_spec.resources + + +_PLAN_KEYS = {"schema", "tool", "stage", "run_id", "source", "inputs", "options", "env"} + + +def parse_plan(data: object) -> Plan: + """Validate a plan mapping; refuse anything loosely pinned.""" + + if not isinstance(data, Mapping): + raise PlanError("plan must be a JSON object") + unknown = set(data) - _PLAN_KEYS + if unknown: + raise PlanError(f"unknown plan keys: {sorted(unknown)}") + if data.get("schema") != PLAN_SCHEMA: + raise PlanError(f"plan schema must be {PLAN_SCHEMA!r}") + tool = TOOLS.get(data.get("tool")) # type: ignore[arg-type] + if tool is None: + raise PlanError(f"unknown tool {data.get('tool')!r}; known: {sorted(TOOLS)}") + stage = data.get("stage") + if stage not in tool.stages: + raise PlanError( + f"tool {tool.name!r} has no stage {stage!r}; known: {sorted(tool.stages)}" + ) + run_id = data.get("run_id") + if not isinstance(run_id, str) or not _RUN_ID.match(run_id): + raise PlanError( + f"run_id {run_id!r}: 3-80 chars of lowercase letters, digits, . _ -" + ) + + source = data.get("source") + if not isinstance(source, Mapping) or not {"commit", "branch"} <= set(source): + raise PlanError("source needs 'commit' and 'branch' (and optional 'repo_url')") + if set(source) - {"commit", "branch", "repo_url"}: + raise PlanError( + f"unknown source keys: {sorted(set(source) - {'commit', 'branch', 'repo_url'})}" + ) + commit = source["commit"] + if not isinstance(commit, str) or not _COMMIT.match(commit): + raise PlanError("source.commit must be a full 40-hex lowercase sha") + branch = source["branch"] + if not isinstance(branch, str) or not _BRANCH.match(branch) or ".." in branch: + raise PlanError(f"source.branch {branch!r} is not a safe branch name") + repo_url = source.get("repo_url", DEFAULT_REPO_URL) + if not isinstance(repo_url, str) or not _GITHUB_URL.match(repo_url): + raise PlanError("source.repo_url must be https://github.com//.git") + + raw_inputs = data.get("inputs") + if not isinstance(raw_inputs, Mapping): + raise PlanError("inputs must be an object of name -> {uri, sha256}") + unknown_inputs = set(raw_inputs) - set(tool.inputs) + if unknown_inputs: + raise PlanError( + f"tool {tool.name!r} takes no inputs {sorted(unknown_inputs)}; " + f"known: {list(tool.inputs)}" + ) + missing = [ + name for name in tool.stages[stage].required_inputs if name not in raw_inputs + ] + if missing: + raise PlanError(f"stage {stage!r} requires inputs {missing}") + inputs = {name: parse_input(name, raw_inputs[name]) for name in sorted(raw_inputs)} + local_paths = [input_local_path(WORK_ROOT, ref) for ref in inputs.values()] + if len(set(local_paths)) != len(local_paths): + raise PlanError("two inputs would stage to the same path") + + raw_options = data.get("options", {}) + if not isinstance(raw_options, Mapping): + raise PlanError("options must be an object") + options: dict[str, object] = {} + for key in sorted(raw_options): + option = tool.options.get(key) + if option is None: + raise PlanError( + f"option {key!r} is not allowlisted for {tool.name!r}; " + f"known: {sorted(tool.options)}" + ) + value = raw_options[key] + if option.kind is bool: + ok = isinstance(value, bool) + elif option.kind is int: + ok = isinstance(value, int) and not isinstance(value, bool) + elif option.kind is float: + ok = isinstance(value, int | float) and not isinstance(value, bool) + else: + ok = isinstance(value, str) and value != "" and not value.startswith("-") + if not ok: + raise PlanError( + f"option {key!r} must be {option.kind.__name__}, got {value!r}" + ) + options[key] = value + + raw_env = data.get("env", {}) + if not isinstance(raw_env, Mapping): + raise PlanError("env must be an object") + env: dict[str, str] = {} + for key in sorted(raw_env): + if not _ENV_KEY.match(key): + raise PlanError( + f"env {key!r} is not allowlisted (MICROCOSM_*, POPULACE_*, *_NUM_THREADS)" + ) + if not isinstance(raw_env[key], str): + raise PlanError(f"env {key!r} must be a string") + env[key] = raw_env[key] + + return Plan( + tool=tool, + stage=stage, + run_id=run_id, + repo_url=repo_url, + commit=commit, + branch=branch, + inputs=inputs, + options=options, + env=env, + ) + + +def load_plan(path: Path) -> tuple[dict, Plan]: + data = json.loads(Path(path).read_text()) + return data, parse_plan(data) + + +def canonical_json(data: object) -> str: + return json.dumps(data, sort_keys=True, separators=(",", ":")) + + +def plan_digest(data: object) -> str: + return hashlib.sha256(canonical_json(data).encode()).hexdigest() + + +# --------------------------------------------------------------------------- # +# Image and argv # +# --------------------------------------------------------------------------- # + + +def image_build_commands(plan: Plan, repo_root: str = IMAGE_REPO_ROOT) -> list[str]: + """Shell steps that pin the image's tree and environment to the plan. + + The tree is a shallow clone of the pushed commit (so ``git rev-parse`` + and ``git status`` inside the tools report the real code vintage), and + the environment is synced from that tree's own uv.lock with --frozen. + """ + + root = shlex.quote(repo_root) + return [ + f"git init -q {root}", + f"git -C {root} remote add origin {shlex.quote(plan.repo_url)}", + f"git -C {root} fetch -q --depth 1 origin {plan.commit}", + f"git -C {root} checkout -q -B {shlex.quote(plan.branch)} FETCH_HEAD", + f'test "$(git -C {root} rev-parse HEAD)" = {plan.commit}', + f"cd {root} && uv sync --all-packages --extra us --frozen", + f'test -z "$(git -C {root} status --porcelain)"', + ] + + +def option_argv(plan: Plan) -> list[str]: + argv: list[str] = [] + for key, value in sorted(plan.options.items()): + option = plan.tool.options[key] + if option.flag in plan.tool.owned_flags: # registry invariant + raise PlanError(f"option {key!r} maps to runner-owned flag {option.flag}") + if option.kind is bool: + if value: + argv.append(option.flag) + else: + argv += [option.flag, str(value)] + return argv + + +def build_stage_argv( + plan: Plan, + *, + python: str, + input_paths: Mapping[str, str], + state_dir: str, +) -> list[str]: + """The exact argv the container runs (cwd = the cloned tree).""" + + missing = sorted(set(plan.inputs) - set(input_paths)) + if missing: + raise PlanError(f"no staged path for inputs {missing}") + tool_argv = plan.tool.argv_builder(plan, input_paths, state_dir) + tool_argv += option_argv(plan) + return [python, "-B", plan.tool.script, *tool_argv] + + +def planned_argv(plan: Plan, work_root: str = WORK_ROOT) -> list[str]: + paths = { + name: input_local_path(work_root, ref) for name, ref in plan.inputs.items() + } + return build_stage_argv( + plan, + python=f"{IMAGE_VENV}/bin/python", + input_paths=paths, + state_dir=str(PurePosixPath(work_root) / "state"), + ) + + +# --------------------------------------------------------------------------- # +# Hashing, mirroring and receipts # +# --------------------------------------------------------------------------- # + +_CHUNK = 8 * 1024 * 1024 + + +def sha256_file(path: Path | str) -> tuple[str, int]: + digest = hashlib.sha256() + size = 0 + with open(path, "rb") as handle: + while chunk := handle.read(_CHUNK): + digest.update(chunk) + size += len(chunk) + return digest.hexdigest(), size + + +def copy_with_sha256(src: Path | str, dst: Path | str) -> tuple[str, int]: + """Copy ``src`` to ``dst`` hashing in the same pass; atomic rename.""" + + dst = Path(dst) + dst.parent.mkdir(parents=True, exist_ok=True) + tmp = dst.with_name(dst.name + ".partial") + digest = hashlib.sha256() + size = 0 + with open(src, "rb") as reader, open(tmp, "wb") as writer: + while chunk := reader.read(_CHUNK): + digest.update(chunk) + writer.write(chunk) + size += len(chunk) + os.replace(tmp, dst) + return digest.hexdigest(), size + + +def verify_digest(ref: InputRef, actual_sha256: str) -> None: + if actual_sha256 != ref.sha256: + raise PlanError( + f"input {ref.name!r} ({ref.uri}) is sha256 {actual_sha256}, " + f"plan pins {ref.sha256}" + ) + + +def tree_listing(root: Path | str) -> dict[str, tuple[int, int]]: + """``{relative posix path: (bytes, mtime_ns)}`` for every regular file.""" + + root = Path(root) + listing: dict[str, tuple[int, int]] = {} + if not root.exists(): + return listing + for path in sorted(root.rglob("*")): + if path.is_file() and not path.is_symlink(): + stat = path.stat() + listing[path.relative_to(root).as_posix()] = ( + stat.st_size, + stat.st_mtime_ns, + ) + return listing + + +def mirror_actions( + source: Mapping[str, tuple[int, int]], destination: Mapping[str, tuple[int, int]] +) -> tuple[list[str], list[str]]: + """Files to copy and to delete so ``destination`` mirrors ``source``. + + A file is copied when it is new or its size or mtime differs (copies + preserve mtime, so an untouched checkpoint is not re-sent); a file the + stage deleted locally is deleted from the destination. + """ + + copy = sorted( + path for path, meta in source.items() if destination.get(path) != meta + ) + delete = sorted(path for path in destination if path not in source) + return copy, delete + + +def mirror_tree(source: Path | str, destination: Path | str) -> dict[str, int]: + source, destination = Path(source), Path(destination) + copy, delete = mirror_actions(tree_listing(source), tree_listing(destination)) + for rel in copy: + target = destination / rel + target.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source / rel, target) + for rel in delete: + (destination / rel).unlink() + return {"copied": len(copy), "deleted": len(delete)} + + +def hash_tree(root: Path | str) -> list[dict[str, object]]: + root = Path(root) + outputs = [] + for rel in tree_listing(root): + sha, size = sha256_file(root / rel) + outputs.append({"path": rel, "bytes": size, "sha256": sha}) + return outputs + + +def build_receipt( + plan: Plan, + plan_data: Mapping, + *, + argv: Sequence[str], + returncode: int, + started_at: str, + finished_at: str, + wall_seconds: float, + peak_rss_bytes: int | None, + inputs_verified: Sequence[Mapping[str, object]], + outputs: Sequence[Mapping[str, object]], + git: Mapping[str, object], + runner: Mapping[str, object], + prior_receipts: Sequence[Mapping[str, object]] = (), +) -> dict[str, object]: + resources = plan.resources + return { + "schema": RECEIPT_SCHEMA, + "status": "COMPLETED" if returncode == 0 else "FAILED", + "tool": plan.tool.name, + "stage": plan.stage, + "run_id": plan.run_id, + "plan_sha256": plan_digest(plan_data), + "plan": plan_data, + "source": { + "repo_url": plan.repo_url, + "commit": plan.commit, + "branch": plan.branch, + **dict(git), + }, + "runner": dict(runner), + "resources": { + "class": resources.name, + "cpu": resources.cpu, + "memory_mib": resources.memory_mib, + "timeout_s": resources.timeout_s, + }, + "argv": list(argv), + "returncode": returncode, + "started_at": started_at, + "finished_at": finished_at, + "wall_seconds": round(wall_seconds, 1), + "peak_rss_bytes": peak_rss_bytes, + "estimated_usd_at_list_price": resources.estimated_usd(wall_seconds), + "inputs": [dict(item) for item in inputs_verified], + "prior_receipts": [dict(item) for item in prior_receipts], + "outputs": [dict(item) for item in outputs], + } + + +def verify_receipt( + receipt: Mapping, state_root: Path | str, *, strict: bool = False +) -> list[str]: + """Re-hash a fetched state tree against a receipt; return the problems.""" + + if receipt.get("schema") != RECEIPT_SCHEMA: + return [f"not a {RECEIPT_SCHEMA} receipt"] + root = Path(state_root) + problems: list[str] = [] + declared = set() + for item in receipt.get("outputs", []): + rel = str(item["path"]) + declared.add(rel) + path = root / rel + if not path.is_file(): + problems.append(f"missing: {rel}") + continue + sha, size = sha256_file(path) + if size != item["bytes"]: + problems.append( + f"size mismatch: {rel} is {size} bytes, receipt {item['bytes']}" + ) + elif sha != item["sha256"]: + problems.append(f"sha256 mismatch: {rel}") + if strict: + extra = sorted(set(tree_listing(root)) - declared) + problems += [f"not in receipt: {rel}" for rel in extra] + return problems + + +def prior_state_problems(plan: Plan, run_identity: Mapping | None) -> list[str]: + """Refusals for a stage whose predecessor state is missing or foreign. + + Every ACS stage after materialize re-verifies the staging digest itself + (``_verify_run_identity``); this catches the same mismatch in the cheap + check, before a paid container starts. + """ + + if plan.tool is not US_ACS_LOCAL_RELEASE or plan.stage in {"materialize", "all"}: + return [] + if not run_identity: + return [ + f"run {plan.run_id!r} has no checkpoints/run_identity.json on " + f"{RUNS_VOLUME}; run materialize for this run first" + ] + problems = [] + if run_identity.get("staging_sha256") != plan.inputs["staging_h5"].sha256: + problems.append( + "staging_h5 sha256 differs from the run's materialized identity " + f"({run_identity.get('staging_sha256')})" + ) + if run_identity.get("ladder_sha256") not in (None, plan.inputs["ladder"].sha256): + problems.append( + "ladder sha256 differs from the run's materialized identity " + f"({run_identity.get('ladder_sha256')})" + ) + return problems + + +def summarize(plan: Plan) -> dict[str, object]: + resources = plan.resources + measured = MEASURED.get((plan.tool.name, plan.stage)) + summary: dict[str, object] = { + "tool": plan.tool.name, + "stage": plan.stage, + "run_id": plan.run_id, + "commit": plan.commit, + "branch": plan.branch, + "resources": { + "class": resources.name, + "cpu": resources.cpu, + "memory_gib": resources.memory_gib, + "timeout_h": resources.timeout_s / 3600, + }, + "inputs": {name: ref.to_json() for name, ref in plan.inputs.items()}, + "argv": planned_argv(plan), + "image_build_commands": image_build_commands(plan), + } + if measured is not None: + summary["measured_locally"] = { + "peak_rss_gb": round(measured.peak_rss_bytes / 1e9, 1), + "wall_seconds": measured.wall_seconds, + "source": measured.source, + } + summary["estimated_usd_at_measured_wall"] = resources.estimated_usd( + measured.wall_seconds + ) + return summary + + +# --------------------------------------------------------------------------- # +# CLI # +# --------------------------------------------------------------------------- # + + +def _digest_lines(paths: Iterable[Path]) -> list[str]: + lines = [] + for path in paths: + sha, size = sha256_file(path) + remote = cas_volume_path(sha, path.name) + lines.append( + json.dumps( + { + "file": str(path), + "bytes": size, + "input": {"uri": f"volume://{remote}", "sha256": sha}, + "upload": ( + f"modal volume put {INPUTS_VOLUME} " + f"{shlex.quote(str(path))} {remote}" + ), + } + ) + ) + return lines + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + sub = parser.add_subparsers(dest="command", required=True) + validate = sub.add_parser("validate", help="validate a plan; print argv and sizing") + validate.add_argument("plan", type=Path) + digest = sub.add_parser("digest", help="sha256 + CAS upload command per file") + digest.add_argument("paths", nargs="+", type=Path) + verify = sub.add_parser("verify-receipt", help="re-hash a fetched state tree") + verify.add_argument("receipt", type=Path) + verify.add_argument("--state-root", type=Path, required=True) + verify.add_argument("--strict", action="store_true") + args = parser.parse_args(argv) + + if args.command == "validate": + try: + _, plan = load_plan(args.plan) + except PlanError as error: + print(f"REFUSED: {error}", file=sys.stderr) + return 2 + print(json.dumps(summarize(plan), indent=2)) + return 0 + if args.command == "digest": + for line in _digest_lines(args.paths): + print(line) + return 0 + receipt = json.loads(args.receipt.read_text()) + problems = verify_receipt(receipt, args.state_root, strict=args.strict) + for problem in problems: + print(problem, file=sys.stderr) + print( + json.dumps( + { + "verified": not problems, + "outputs": len(receipt.get("outputs", [])), + "problems": len(problems), + } + ) + ) + return 0 if not problems else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) From 732191d428bbdaa74d05b8f0df72bb28896c637a Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 22 Sep 2026 23:20:25 -0400 Subject: [PATCH 02/19] Add the runbook for running a US build stage on Modal How a plan pins a run (commit, inputs by digest, allowlisted options), what the image, input staging, state mirroring and receipts do, the resource classes and their list-price cost from the #974 measured peaks, the exact commands (digest, upload, validate, check, run, fetch, verify), the #974 materialize replay as the first acceptance test, data placement, and what is not covered yet. Co-Authored-By: Claude Opus 5.5 --- docs/us-modal-stage-runbook.md | 208 +++++++++++++++++++++++++++++++++ 1 file changed, 208 insertions(+) create mode 100644 docs/us-modal-stage-runbook.md diff --git a/docs/us-modal-stage-runbook.md b/docs/us-modal-stage-runbook.md new file mode 100644 index 000000000..b4c1928b4 --- /dev/null +++ b/docs/us-modal-stage-runbook.md @@ -0,0 +1,208 @@ +# Running a US build stage on Modal + +Epic #956 acceleration item E: heavy US stages should not have to queue on +the one 128 GiB build machine. This runbook covers the smallest working path: +one registered tool (`tools/build_us_acs_local_release.py`, stages +`materialize`, `calibrate`, `qa`, `finalize`, `package`, or `all`) run on +Modal from a pinned commit, with inputs fetched by digest and outputs listed +with sha256 receipts. + +Two files do the work: + +- `tools/modal_us_stage.py` is the Modal app. It builds the image, stages the + inputs, runs the tool and writes the receipt. +- `tools/modal_us_stage_plan.py` is its pure half, standard library only. + It validates plans, builds the argv, sizes resources, hashes and mirrors + files, and writes and verifies receipts. Unit tests: + `packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py`. + +Nothing here uploads to the Hugging Face Hub, touches `latest.json` or +notifies anyone. The furthest a stage goes is `package`, which writes a +release directory onto the runs volume. Publication stays the human step in +`tools/publish_release.sh`. + +## How a run works + +1. **A plan pins everything.** A plan is a JSON file + (`microcosm-modal-us-stage-plan/1`; example: + `docs/us-modal-stage-example-plan.json`). It names the tool, the stage, a + `run_id`, a full 40-hex commit and the branch it was pushed on, and every + input as `{uri, sha256}`. Options (`soi_mode`, `hh_chunk`, `epochs` and + so on) and environment overrides (`MICROCOSM_*`, `POPULACE_*`, + `*_NUM_THREADS`) come from allowlists. The runner sets the flags for + input paths, checkpoints and outputs, and a plan cannot pass them. It + also cannot pass `--allow-dirty`: every run builds from a clean clone. +2. **The image is the commit.** It starts from `debian_slim` with Python + 3.14, the version the local US builds record (`runtime.python` 3.14.4 in + the #974 build manifest). It adds git and uv 0.11.7, makes a shallow + clone of the plan's commit from GitHub, checks out the plan's branch name, + and asserts `HEAD` equals the commit. It then runs + `uv sync --all-packages --extra us --frozen` against that tree's own + `uv.lock` into `/opt/venv` and asserts `git status --porcelain` is empty. + The release tool's `_repo_code_identity` therefore records the real sha + and branch. A commit that is not on GitHub fails the build. The runner + code comes from your checkout, not from the pinned commit, so any pushed + commit whose tree has the tool can run, including commits older than the + runner. The receipt records the sha256 of both runner files. +3. **Inputs are verified before the tool starts.** A `volume://` input is a + path on the `microcosm-us-stage-inputs` volume. Local files go there + content-addressed, at `cas/sha256//`, and the plan refuses a + CAS path whose digest is not the input's own. An `hf://` input is a Hub + file at an explicit revision: + `hf://datasets//@/`. Public repos need no + token. Each input is copied or downloaded to a stable local path, + `/work/inputs//`, hashed in the same pass, and refused on + mismatch. +4. **State lives on the runs volume.** Each `run_id` keeps + `runs//state/` on `microcosm-us-stage-runs`. That directory holds + the checkpoints, the calibrated H5, the release root for `package` and + the stage logs. A stage pulls the state to local disk and runs the tool + there, since materialize writes a memory-mapped matrix. It then mirrors + the state back: changed files are copied and files the tool deleted are + removed. Later stages of the same run pick up from that state. Before + paying for inputs, a stage after materialize checks that the run's + `checkpoints/run_identity.json` exists and pins the same staging and + ladder digests. The tool re-verifies the staging digest itself + (`_verify_run_identity`). +5. **Every output has a receipt.** `runs//receipts/-.json` + (`microcosm-modal-us-stage-receipt/1`) records: + - the plan and its sha256; + - the commit, the clone's `HEAD`, clean state and branch; + - the sha256 of the runner files and of `uv.lock`; + - the resource class, the exact argv, the return code and wall seconds; + - peak child RSS and the list-price cost estimate; + - every verified input; + - the sha256 of the earlier receipts in the run; + - every file in the state tree with its bytes and sha256. + + A failed stage still mirrors its state, so `calibrate` can resume from + `weights_latest.npz`, and its receipt says `FAILED`. + +## Resources and cost + +The heavy and light classes are sized from the #974 measured peaks +(`experiments/us-acs-local-hours-rebuild-20260922/run-resources-and-staging-excerpt.json`, +totals SOI surface). The overnight build of 22 September was expected to +peak near 94 GB for materialize on the state surface. That number was +reported with the task and has not been measured by this runner. + +| Stage | Measured locally | Class | Request | List price at measured wall | +| --- | --- | --- | --- | --- | +| materialize | 74.8 GB, 5,067 s wall, 4,990 CPU-s | heavy | 4 cores, 128 GiB, 8 h timeout | about $1.71 | +| calibrate | 67.6 GB, 415 s | heavy | 4 cores, 128 GiB | about $0.14 | +| qa | 21.9 GB, 172 s | light | 2 cores, 48 GiB, 4 h timeout | about $0.02 | +| finalize | 23.3 GB, 79 s | light | 2 cores, 48 GiB | about $0.01 | +| package | 21.8 GB, 79 s | light | 2 cores, 48 GiB | about $0.01 | +| check | n/a | check | 2 cores, 8 GiB, 30 min timeout | cents | + +The prices are Modal's list prices for standard compute, read from +modal.com/pricing on 22 September 2026: $0.0000131 per core-second and +$0.00000222 per GiB-second. Modal bills the higher of the request and actual +use. The heavy class costs about $1.21 an hour, so an 8-hour timeout costs at +most about $9.70. The engine pass in materialize is single-threaded (CPU +seconds roughly equal wall seconds), so extra cores would not speed it up. +The table leaves out volume storage and image builds. + +## Commands + +Run everything from a checkout of this branch. The Modal CLI must be on +`PATH` and logged in to the `policyengine` workspace (`modal profile +current`). The plan-module commands run on plain `python3` and need neither +Modal nor the workspace environment. + +```bash +# 0. Once per workspace. The app also creates these volumes on first use. +modal volume create microcosm-us-stage-inputs +modal volume create microcosm-us-stage-runs + +# 1. Put local inputs on the inputs volume, content-addressed. `digest` +# prints each file's sha256, the plan input, and the exact upload command. +python3 tools/modal_us_stage_plan.py digest \ + /acs_multispine_staging.h5 /acs_multispine_staging.summary.json \ + /consumer_facts.jsonl /us_puma_ladder_2020.npz +modal volume put microcosm-us-stage-inputs cas/sha256// + +# 2. Write the plan (copy docs/us-modal-stage-example-plan.json) and +# validate it locally. This prints the argv, the resource class and the +# estimate, and exits 2 with REFUSED on a loose plan. +python3 tools/modal_us_stage_plan.py validate plan.json + +# 3. Check on Modal (default mode; 2 cores / 8 GiB). It builds the image, +# verifies the clone, imports the environment, runs the pinned tool's own +# _parse_args on the built argv, checks every input digest (Hub inputs +# by their LFS sha256, without downloading) and the run's prior state. +MICROCOSM_MODAL_PLAN=plan.json modal run tools/modal_us_stage.py + +# 4. Run the stage. --detach keeps it running if this terminal goes away; +# the receipt lands on the runs volume either way. +MICROCOSM_MODAL_PLAN=plan.json modal run --detach tools/modal_us_stage.py --run + +# 5. Fetch the state and verify it against the receipt. +modal volume get microcosm-us-stage-runs runs/ ./modal-runs/ +python3 tools/modal_us_stage_plan.py verify-receipt \ + ./modal-runs//receipts/-.json \ + --state-root ./modal-runs//state +``` + +Next stage: copy the plan, change `stage`, keep the `run_id`, and drop `feed` +if you like, since only materialize reads it. Run steps 3 and 4 again. Run +one stage of a run at a time, because two concurrent stages would race on the +same state directory. + +For a gated or private Hub input, set `MICROCOSM_MODAL_HF_SECRET` to the name +of a Modal secret that holds `HF_TOKEN`, for example `huggingface-token` in +the `policyengine` workspace. Otherwise no secret is attached, and the tool +runs with `HF_HUB_OFFLINE=1`. + +## First acceptance: replay #974 materialize + +`docs/us-modal-stage-example-plan.json` replays the materialize stage of the +full-scale local build of 22 September (#974). The plan pins the build +commit `cadaf418` (branch `local-acs-hours-rebuild-20260921`), that run's +staging H5 and staging summary, the `chronicle_us_b571381` feed and the PUMA +ladder, all by the digests the #974 receipt records. A Modal replay should +reproduce these values from that receipt's `run_identity`: + +- `targets_sha256` `0f447ce92b4279881382fdd4006be47cfbdd6972be438f967a7cca97c29a4f81` +- `n_targets` 1,247, from 760 admin specs, all compiled +- 1,588,854 households +- no dropped population cells + +Compare `runs/replay-974-materialize/state/checkpoints/run_identity.json` +with +`experiments/us-acs-local-hours-rebuild-20260922/build_manifest.json`. A +match shows that the Modal image and platform reproduce the local target +compile. + +## Data placement + +Inputs go to the `policyengine` Modal workspace. Upload only files that are +allowed there. No registered stage takes the restricted IRS PUF files +(`--puf-source-year-csv` in `tools/build_us_puf_support_base.py`), and the +base stage is not registered. Logs hold only what the tool prints. Receipts +hold digests, sizes and paths, never file contents. + +## What this does not cover yet + +- **The graph-native line (#893).** The DAG executor lives on #893's branch, + not on main. Registering it means one more `ToolSpec` with its CLI, inputs + and state layout. The executor runs nodes one after another, and the 1/15 + run peaked at 46.7 GB (#956). A full-source run needs the byte-transport + and parallel-executor work first, plus a resource class sized from a + measured full-scale peak. +- **The base stage (`tools/build_us_puf_support_base.py`)** is not + registered. Its peak was reported at 72 GB; this runner has not measured + it. It requires the processed PUF (`--puf-h5`), and + `--puf-source-year-csv` (the restricted TY2015 IRS PUF CSV) whenever the + processed PUF is nonzero. Nobody has decided whether those files may sit + on Modal volumes. +- **The staging stage (`tools/build_us_acs_multispine_base.py`)** is not + registered either. It takes a directory input (`--inputs-dir`, the ACS + PUMS archive cache), which the plan format does not support. Supporting it + would take an archive digest plus extraction. +- **Preemption.** Functions run with `retries=0` on Modal's default + (preemptible) placement. A preempted materialize has to be run again, and + state is only mirrored when a stage ends. +- **Certification.** A receipt proves which bytes a stage produced. It does + not certify a release. Preflight (`tools/preflight_us_release_gates.py`) + and certification still run on the output as before. From 5fdc81a3b0420ee667b9b5c41124c3df46f5a71b Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 22 Sep 2026 23:26:08 -0400 Subject: [PATCH 03/19] Add a near-free runner smoke and a wall-clock budget to Modal stage plans runner-smoke is an inline tool (no file in the pinned tree, so it runs on any pushed commit) that imports the synced environment, reads every staged input and writes one state file, which proves the whole run path: inputs staged and verified from the volume and the Hub, state mirrored to the runs volume, receipt written. It runs in the check-sized class. max_wall_seconds bounds a stage's cost below the resource class's hard timeout: the runner stops the tool when it passes and the receipt records stopped_at_budget and FAILED. Co-Authored-By: Claude Opus 5.5 --- docs/us-modal-stage-smoke-plan.json | 25 +++++ .../tests/test_us_modal_stage_plan_tool.py | 85 +++++++++++++++++ tools/modal_us_stage.py | 51 +++++++++- tools/modal_us_stage_plan.py | 92 ++++++++++++++++++- 4 files changed, 246 insertions(+), 7 deletions(-) create mode 100644 docs/us-modal-stage-smoke-plan.json diff --git a/docs/us-modal-stage-smoke-plan.json b/docs/us-modal-stage-smoke-plan.json new file mode 100644 index 000000000..e3fe7d520 --- /dev/null +++ b/docs/us-modal-stage-smoke-plan.json @@ -0,0 +1,25 @@ +{ + "schema": "microcosm-modal-us-stage-plan/1", + "tool": "runner-smoke", + "stage": "smoke", + "run_id": "runner-smoke-cadaf418", + "source": { + "commit": "cadaf418097964038ae539a3d5ce5acedaf49912", + "branch": "local-acs-hours-rebuild-20260921" + }, + "inputs": { + "ladder": { + "uri": "volume://cas/sha256/39a2ab2abeab07a88362af7ab2940e0e1d50a297c919e4bbc6fb65bab51147d8/us_puma_ladder_2020.npz", + "sha256": "39a2ab2abeab07a88362af7ab2940e0e1d50a297c919e4bbc6fb65bab51147d8" + }, + "staging_summary": { + "uri": "volume://cas/sha256/ebee1ac61811810082936b827a75e6f95ab0db1b139715f1d663bcb812c9e541/acs_multispine_staging.summary.json", + "sha256": "ebee1ac61811810082936b827a75e6f95ab0db1b139715f1d663bcb812c9e541" + }, + "hub_file": { + "uri": "hf://datasets/policyengine/populace-us@85a1ccb0b61e46fede793e3da6d5e4dd2e307374/latest.json", + "sha256": "7a5021f34d91ae79961ab93c6d90dc2d1b9322e9cc485606f92994a1ddd60fc7" + } + }, + "max_wall_seconds": 600 +} diff --git a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py index ded601186..a7ef7f96c 100644 --- a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py +++ b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py @@ -461,3 +461,88 @@ def test_example_plan_in_the_runbook_parses() -> None: example = ROOT / "docs" / "us-modal-stage-example-plan.json" plan = plan_lib.parse_plan(json.loads(example.read_text())) assert plan.tool is plan_lib.US_ACS_LOCAL_RELEASE + + +def test_every_stage_uses_a_resource_class_the_app_can_run() -> None: + # tools/modal_us_stage.py defines one Modal function per class. + for tool in plan_lib.TOOLS.values(): + for stage in tool.stages.values(): + assert stage.resources is plan_lib.RESOURCE_CLASSES[stage.resources.name] + + +@pytest.mark.parametrize( + ("value", "ok"), + [(60, True), (10_800, True), (59, False), (8 * 3600, False), (True, False)], +) +def test_max_wall_seconds_is_bounded_below_the_hard_timeout(value, ok) -> None: + data = _plan_data(max_wall_seconds=value) + if ok: + assert plan_lib.parse_plan(data).max_wall_seconds == value + else: + with pytest.raises(plan_lib.PlanError, match="max_wall_seconds"): + plan_lib.parse_plan(data) + + +def test_budget_stop_marks_the_receipt_failed(tmp_path: Path) -> None: + data = _plan_data(max_wall_seconds=3600) + plan = plan_lib.parse_plan(data) + receipt = plan_lib.build_receipt( + plan, + data, + argv=["python"], + returncode=0, + started_at="t0", + finished_at="t1", + wall_seconds=3600.0, + peak_rss_bytes=None, + inputs_verified=[], + outputs=[], + git={}, + runner={}, + stopped_at_budget=True, + ) + assert receipt["status"] == "FAILED" + assert (receipt["max_wall_seconds"], receipt["stopped_at_budget"]) == (3600, True) + + +def _smoke_plan_data() -> dict: + return { + "schema": plan_lib.PLAN_SCHEMA, + "tool": "runner-smoke", + "stage": "smoke", + "run_id": "runner-smoke-20260923", + "source": {"commit": COMMIT, "branch": "us-modal-stage-runner"}, + "inputs": { + "ladder": { + "uri": f"volume://cas/sha256/{LADDER_SHA}/us_puma_ladder_2020.npz", + "sha256": LADDER_SHA, + }, + }, + } + + +def test_runner_smoke_is_inline_small_and_runs_the_synced_environment( + tmp_path: Path, +) -> None: + plan = plan_lib.parse_plan(_smoke_plan_data()) + assert plan.resources is plan_lib.CHECK + argv = plan_lib.planned_argv(plan) + assert argv[:3] == ["/opt/venv/bin/python", "-B", "-c"] + assert argv[4:] == [ + "/work/state", + "/work/inputs/ladder/us_puma_ladder_2020.npz", + ] + # The inline code runs as written: it writes the state file the receipt + # lists (executed here against a local stand-in input). + code = argv[3] + assert "import microcosm.build" in code + ladder = tmp_path / "inputs" / "ladder" / "us_puma_ladder_2020.npz" + ladder.parent.mkdir(parents=True) + ladder.write_bytes(b"npz") + state = tmp_path / "state" + import subprocess + + subprocess.run([sys.executable, "-c", code, str(state), str(ladder)], check=True) + payload = json.loads((state / "smoke" / "inputs.json").read_text()) + assert payload["inputs"] == [{"input": "ladder", "bytes": 3}] + assert payload["policyengine_us"] diff --git a/tools/modal_us_stage.py b/tools/modal_us_stage.py index 3152f965a..b588d19ee 100644 --- a/tools/modal_us_stage.py +++ b/tools/modal_us_stage.py @@ -36,6 +36,7 @@ import shutil import subprocess import sys +import threading import time from datetime import UTC, datetime from pathlib import Path @@ -137,7 +138,8 @@ def _git_state(plan: plan_lib.Plan) -> dict[str, object]: "head_matches_plan": head == plan.commit, "tree_clean": clean, "branch_checked_out": _git("branch", "--show-current"), - "tool_present": (Path(plan_lib.IMAGE_REPO_ROOT) / plan.tool.script).is_file(), + "tool_present": plan.tool.script is None + or (Path(plan_lib.IMAGE_REPO_ROOT) / plan.tool.script).is_file(), } @@ -253,6 +255,8 @@ def _check_input(ref: plan_lib.InputRef, work: Path) -> dict[str, object]: def _parse_check(plan: plan_lib.Plan, argv: list[str]) -> dict[str, object]: """Run the pinned tool's own argument parser on the built argv.""" + if plan.tool.script is None: + return {"returncode": 0, "skipped": "inline tool; nothing to parse"} python, _, script, *tool_argv = argv if plan.tool.parse_function is None: cmd = [python, "-B", script, "--help"] @@ -269,6 +273,28 @@ def _parse_check(plan: plan_lib.Plan, argv: list[str]) -> dict[str, object]: } +class _BudgetWatch(threading.Thread): + """Stop the tool once the plan's max_wall_seconds has passed.""" + + def __init__(self, proc: subprocess.Popen, max_wall_seconds: int | None) -> None: + super().__init__(daemon=True) + self.proc = proc + self.max_wall_seconds = max_wall_seconds + self.done = threading.Event() + self.fired = False + + def run(self) -> None: + if self.max_wall_seconds is None or self.done.wait(self.max_wall_seconds): + return + self.fired = True + print(f"BUDGET: stopping the tool after {self.max_wall_seconds}s", flush=True) + self.proc.terminate() + try: + self.proc.wait(timeout=60) + except subprocess.TimeoutExpired: + self.proc.kill() + + def _load_json(path: Path) -> dict | None: try: return json.loads(path.read_text()) @@ -394,11 +420,14 @@ def _run_stage(plan_data: dict) -> dict: stdout=subprocess.PIPE, stderr=subprocess.STDOUT, ) + budget = _BudgetWatch(proc, plan.max_wall_seconds) + budget.start() assert proc.stdout is not None for line in proc.stdout: log.write(line) print(line, end="", flush=True) returncode = proc.wait() + budget.done.set() wall = time.time() - started peak_kib = resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss finished_at = _now() @@ -420,6 +449,7 @@ def _run_stage(plan_data: dict) -> dict: git=git, runner={**_runner_identity(), "state_pulled": pulled, "state_pushed": pushed}, prior_receipts=prior, + stopped_at_budget=budget.fired, ) receipts_dir.mkdir(parents=True, exist_ok=True) name = f"{plan.stage}-{started_at.replace(':', '')}.json" @@ -456,7 +486,24 @@ def run_stage_light(plan_data: dict) -> dict: return _run_stage(plan_data) -RUNNERS = {"heavy": run_stage_heavy, "light": run_stage_light} +@app.function( + image=image, + volumes=VOLUMES, + secrets=_secrets(), + cpu=plan_lib.CHECK.cpu, + memory=plan_lib.CHECK.memory_mib, + timeout=plan_lib.CHECK.timeout_s, + retries=0, +) +def run_stage_small(plan_data: dict) -> dict: + return _run_stage(plan_data) + + +RUNNERS = { + "heavy": run_stage_heavy, + "light": run_stage_light, + "check": run_stage_small, +} @app.local_entrypoint() diff --git a/tools/modal_us_stage_plan.py b/tools/modal_us_stage_plan.py index 302c3a884..f41362bc5 100644 --- a/tools/modal_us_stage_plan.py +++ b/tools/modal_us_stage_plan.py @@ -174,7 +174,8 @@ class StageSpec: @dataclass(frozen=True) class ToolSpec: name: str - script: str + # Path of the tool in the pinned tree; None for an inline ``-c`` tool. + script: str | None inputs: tuple[str, ...] stages: Mapping[str, StageSpec] options: Mapping[str, OptionFlag] @@ -269,7 +270,54 @@ def _acs_local_release_argv( parse_function="_parse_args", ) -TOOLS: dict[str, ToolSpec] = {US_ACS_LOCAL_RELEASE.name: US_ACS_LOCAL_RELEASE} +# A near-free end-to-end proof of the run path (inputs staged and verified, +# the synced environment imported, a state file written, mirrored to the +# runs volume and listed in a receipt) on any pushed commit, because the +# code is inline rather than a file in the pinned tree. +_RUNNER_SMOKE_CODE = """\ +import importlib.metadata as metadata +import json +import pathlib +import sys + +import microcosm.build # noqa: F401 - proves the synced environment imports + +state = pathlib.Path(sys.argv[1]) +rows = [ + {"input": pathlib.Path(p).parent.name, "bytes": pathlib.Path(p).stat().st_size} + for p in sys.argv[2:] +] +payload = {"inputs": rows, "policyengine_us": metadata.version("policyengine-us")} +(state / "smoke").mkdir(parents=True, exist_ok=True) +(state / "smoke" / "inputs.json").write_text(json.dumps(payload, indent=2) + "\\n") +print("runner smoke ok:", json.dumps(payload)) +""" + + +def _runner_smoke_argv( + plan: Plan, input_paths: Mapping[str, str], state_dir: str +) -> list[str]: + return [ + "-c", + _RUNNER_SMOKE_CODE, + state_dir, + *(input_paths[name] for name in sorted(plan.inputs)), + ] + + +RUNNER_SMOKE = ToolSpec( + name="runner-smoke", + script=None, + inputs=("feed", "ladder", "staging_summary", "hub_file"), + stages={"smoke": StageSpec("smoke", CHECK, ())}, + options={}, + owned_flags=frozenset(), + argv_builder=_runner_smoke_argv, +) + +TOOLS: dict[str, ToolSpec] = { + tool.name: tool for tool in (US_ACS_LOCAL_RELEASE, RUNNER_SMOKE) +} # --------------------------------------------------------------------------- # @@ -416,6 +464,9 @@ class Plan: inputs: Mapping[str, InputRef] options: Mapping[str, object] = field(default_factory=dict) env: Mapping[str, str] = field(default_factory=dict) + # Runner-side budget: the tool is stopped after this many seconds, so a + # stage's cost is bounded below the resource class's hard timeout. + max_wall_seconds: int | None = None @property def stage_spec(self) -> StageSpec: @@ -426,7 +477,19 @@ def resources(self) -> Resources: return self.stage_spec.resources -_PLAN_KEYS = {"schema", "tool", "stage", "run_id", "source", "inputs", "options", "env"} +_PLAN_KEYS = { + "schema", + "tool", + "stage", + "run_id", + "source", + "inputs", + "options", + "env", + "max_wall_seconds", +} +# Time the container keeps for staging inputs, hashing and mirroring state. +_RUNNER_OVERHEAD_SECONDS = 15 * 60 def parse_plan(data: object) -> Plan: @@ -528,6 +591,18 @@ def parse_plan(data: object) -> Plan: raise PlanError(f"env {key!r} must be a string") env[key] = raw_env[key] + max_wall = data.get("max_wall_seconds") + ceiling = tool.stages[stage].resources.timeout_s - _RUNNER_OVERHEAD_SECONDS + if max_wall is not None and ( + not isinstance(max_wall, int) + or isinstance(max_wall, bool) + or not 60 <= max_wall <= ceiling + ): + raise PlanError( + f"max_wall_seconds must be an integer from 60 to {ceiling} for " + f"stage {stage!r}" + ) + return Plan( tool=tool, stage=stage, @@ -538,6 +613,7 @@ def parse_plan(data: object) -> Plan: inputs=inputs, options=options, env=env, + max_wall_seconds=max_wall, ) @@ -607,7 +683,8 @@ def build_stage_argv( raise PlanError(f"no staged path for inputs {missing}") tool_argv = plan.tool.argv_builder(plan, input_paths, state_dir) tool_argv += option_argv(plan) - return [python, "-B", plan.tool.script, *tool_argv] + script = [] if plan.tool.script is None else [plan.tool.script] + return [python, "-B", *script, *tool_argv] def planned_argv(plan: Plan, work_root: str = WORK_ROOT) -> list[str]: @@ -734,11 +811,14 @@ def build_receipt( git: Mapping[str, object], runner: Mapping[str, object], prior_receipts: Sequence[Mapping[str, object]] = (), + stopped_at_budget: bool = False, ) -> dict[str, object]: resources = plan.resources return { "schema": RECEIPT_SCHEMA, - "status": "COMPLETED" if returncode == 0 else "FAILED", + "status": "COMPLETED" + if returncode == 0 and not stopped_at_budget + else "FAILED", "tool": plan.tool.name, "stage": plan.stage, "run_id": plan.run_id, @@ -759,6 +839,8 @@ def build_receipt( }, "argv": list(argv), "returncode": returncode, + "max_wall_seconds": plan.max_wall_seconds, + "stopped_at_budget": stopped_at_budget, "started_at": started_at, "finished_at": finished_at, "wall_seconds": round(wall_seconds, 1), From 9ff62f28cd47b818b5552a0d15a767acc3dddf2d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 22 Sep 2026 23:28:41 -0400 Subject: [PATCH 04/19] Record the Modal check and smoke run in the runbook What ran on Modal on 22 September (volumes, the #974 replay check with its one correct refusal, the smoke run and its locally verified receipt), the max_wall_seconds budget, the smoke plan, and the exact commands for the heavy replay that has not run yet. Co-Authored-By: Claude Opus 5.5 --- docs/us-modal-stage-runbook.md | 58 ++++++++++++++++++++++++++++++++-- 1 file changed, 55 insertions(+), 3 deletions(-) diff --git a/docs/us-modal-stage-runbook.md b/docs/us-modal-stage-runbook.md index b4c1928b4..536af92ef 100644 --- a/docs/us-modal-stage-runbook.md +++ b/docs/us-modal-stage-runbook.md @@ -33,8 +33,8 @@ release directory onto the runs volume. Publication stays the human step in input paths, checkpoints and outputs, and a plan cannot pass them. It also cannot pass `--allow-dirty`: every run builds from a clean clone. 2. **The image is the commit.** It starts from `debian_slim` with Python - 3.14, the version the local US builds record (`runtime.python` 3.14.4 in - the #974 build manifest). It adds git and uv 0.11.7, makes a shallow + 3.14, the minor version the local US builds record (`runtime.python` + 3.14.4 in the #974 build manifest; Modal served 3.14.2). It adds git and uv 0.11.7, makes a shallow clone of the plan's commit from GitHub, checks out the plan's branch name, and asserts `HEAD` equals the commit. It then runs `uv sync --all-packages --extra us --frozen` against that tree's own @@ -134,16 +134,26 @@ python3 tools/modal_us_stage_plan.py validate plan.json MICROCOSM_MODAL_PLAN=plan.json modal run tools/modal_us_stage.py # 4. Run the stage. --detach keeps it running if this terminal goes away; -# the receipt lands on the runs volume either way. +# the receipt lands on the runs volume either way. Set "max_wall_seconds" +# in the plan to cap the cost below the class's hard timeout: the runner +# stops the tool then, and the receipt says FAILED, stopped_at_budget. MICROCOSM_MODAL_PLAN=plan.json modal run --detach tools/modal_us_stage.py --run # 5. Fetch the state and verify it against the receipt. +mkdir -p modal-runs modal volume get microcosm-us-stage-runs runs/ ./modal-runs/ python3 tools/modal_us_stage_plan.py verify-receipt \ ./modal-runs//receipts/-.json \ --state-root ./modal-runs//state ``` +To prove the run path on a new commit or workspace for well under a cent, +run `docs/us-modal-stage-smoke-plan.json` (tool `runner-smoke`). It is an +inline script, so it needs no file in the pinned tree. It imports the synced +environment, reads every staged input (two volume files and one Hub file) +and writes one state file, which goes through the same staging, mirroring +and receipt code as a real stage. + Next stage: copy the plan, change `stage`, keep the `run_id`, and drop `feed` if you like, since only materialize reads it. Run steps 3 and 4 again. Run one stage of a run at a time, because two concurrent stages would race on the @@ -174,6 +184,48 @@ with match shows that the Modal image and platform reproduce the local target compile. +## Verified on Modal, 22 September 2026 + +These runs were in the `policyengine` workspace. The main checkout, the +overnight build and the Hub were not touched. + +- **Volumes.** `microcosm-us-stage-inputs` and `microcosm-us-stage-runs` + were created. Three small inputs of the #974 build went to `cas/sha256/`: + the PUMA ladder (446,791 bytes), the staging summary (498,113 bytes) and + the `chronicle_us_b571381` feed (164,624,488 bytes). The 10.7 GB staging + H5 was not uploaded. +- **Check of the #974 replay plan, which reported one problem, as + intended.** The image built from `cadaf418`. `HEAD` matched, the tree was + clean, and the branch was checked out. The environment synced from that + tree's lock: policyengine-us 2.2.1 and policyengine-core 3.32.5, the + versions #974 records. The pinned tool's `_parse_args` accepted the built + argv and returned `["materialize"]`. The three uploaded inputs were + verified by sha256 on the volume. The only problem reported was + `staging_h5` not being on the volume, which is the correct refusal. +- **Smoke run (`--run`).** `runner-smoke-cadaf418` completed in 7 seconds. + It staged and verified the two volume inputs and the public Hub file + `policyengine/populace-us@85a1ccb0…/latest.json`, and wrote + `smoke/inputs.json`. The receipt is + `runs/runner-smoke-cadaf418/receipts/smoke-2026-09-23T032440Z.json`. After + `modal volume get`, `verify-receipt --strict` passed locally with 2 + outputs and 0 problems. Modal served Python 3.14.2. The local builds + recorded 3.14.4. + +The heavy materialize replay has not been run. It needs the staging H5 on +the inputs volume, which is about 10.7 GB to upload, plus one heavy run at +about $1.71 list price for the measured wall time. That run is the next +step: + +```bash +# On the build machine, the #974 staging H5 is under +# _recovered/scratch-backup/893/local-hours-rebuild-20260921/full-staging-a3/. +python3 tools/modal_us_stage_plan.py digest /acs_multispine_staging.h5 +# expect f335f573…; then run the printed `modal volume put …` line +MICROCOSM_MODAL_PLAN=docs/us-modal-stage-example-plan.json modal run tools/modal_us_stage.py +MICROCOSM_MODAL_PLAN=docs/us-modal-stage-example-plan.json \ + modal run --detach tools/modal_us_stage.py --run +``` + ## Data placement Inputs go to the `policyengine` Modal workspace. Upload only files that are From af18bd65b1696a7578a88ed078c98e29fbb32518 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 22 Sep 2026 23:28:57 -0400 Subject: [PATCH 05/19] Note that a preempted or OOM-killed container leaves no receipt Co-Authored-By: Claude Opus 5.5 --- docs/us-modal-stage-runbook.md | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/us-modal-stage-runbook.md b/docs/us-modal-stage-runbook.md index 536af92ef..30511022a 100644 --- a/docs/us-modal-stage-runbook.md +++ b/docs/us-modal-stage-runbook.md @@ -252,9 +252,11 @@ hold digests, sizes and paths, never file contents. registered either. It takes a directory input (`--inputs-dir`, the ACS PUMS archive cache), which the plan format does not support. Supporting it would take an archive digest plus extraction. -- **Preemption.** Functions run with `retries=0` on Modal's default - (preemptible) placement. A preempted materialize has to be run again, and - state is only mirrored when a stage ends. +- **Preemption and out-of-memory kills.** Functions run with `retries=0` on + Modal's default (preemptible) placement. The heavy class sets a memory + request but no hard limit. State is only mirrored and the receipt only + written when the tool exits, so a container that is preempted or killed + for memory leaves no receipt, and the stage has to be run again. - **Certification.** A receipt proves which bytes a stage produced. It does not certify a release. Preflight (`tools/preflight_us_release_gates.py`) and certification still run on the output as before. From 09373886e5dacf771319cdbffa0a9ad04aa762c7 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Tue, 22 Sep 2026 23:30:37 -0400 Subject: [PATCH 06/19] Skip the smoke-code execution test without the us extra The inline smoke records the policyengine-us version, which the engine-free fast and wheels lanes do not install; the argv-shape test still runs there. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_us_modal_stage_plan_tool.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py index a7ef7f96c..7c57c3fa1 100644 --- a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py +++ b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py @@ -12,6 +12,7 @@ import importlib.util import json import os +import subprocess import sys from pathlib import Path @@ -521,27 +522,30 @@ def _smoke_plan_data() -> dict: } -def test_runner_smoke_is_inline_small_and_runs_the_synced_environment( - tmp_path: Path, -) -> None: +def test_runner_smoke_is_inline_and_check_sized() -> None: plan = plan_lib.parse_plan(_smoke_plan_data()) assert plan.resources is plan_lib.CHECK argv = plan_lib.planned_argv(plan) assert argv[:3] == ["/opt/venv/bin/python", "-B", "-c"] + assert "import microcosm.build" in argv[3] assert argv[4:] == [ "/work/state", "/work/inputs/ladder/us_puma_ladder_2020.npz", ] + + +@pytest.mark.skipif( + importlib.util.find_spec("policyengine_us") is None, + reason="the smoke records the policyengine-us version (the us extra)", +) +def test_runner_smoke_code_writes_its_state_file(tmp_path: Path) -> None: # The inline code runs as written: it writes the state file the receipt # lists (executed here against a local stand-in input). - code = argv[3] - assert "import microcosm.build" in code + code = plan_lib.planned_argv(plan_lib.parse_plan(_smoke_plan_data()))[3] ladder = tmp_path / "inputs" / "ladder" / "us_puma_ladder_2020.npz" ladder.parent.mkdir(parents=True) ladder.write_bytes(b"npz") state = tmp_path / "state" - import subprocess - subprocess.run([sys.executable, "-c", code, str(state), str(ladder)], check=True) payload = json.loads((state / "smoke" / "inputs.json").read_text()) assert payload["inputs"] == [{"input": "ladder", "bytes": 3}] From b2e34fe4ff8e6dbf0da42d28c5d8e034de66566d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 23 Sep 2026 01:55:10 -0400 Subject: [PATCH 07/19] Add the plan for the Modal acceptance run of the 23 September materialize Materialize on the state SOI surface at build commit 767312d6 (branch overnight-acs-local-20260923), with that night's staging H5 and summary, the chronicle_us_b571381 feed and the PUMA ladder pinned by digest, the local run's peak-limit environment, and a 26,400 s wall budget (about $8.90 at list price for the heavy class, under $10 with staging and mirroring). Co-Authored-By: Claude Opus 5.5 --- ...-modal-stage-acceptance-20260923-plan.json | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) create mode 100644 docs/us-modal-stage-acceptance-20260923-plan.json diff --git a/docs/us-modal-stage-acceptance-20260923-plan.json b/docs/us-modal-stage-acceptance-20260923-plan.json new file mode 100644 index 000000000..807f5f602 --- /dev/null +++ b/docs/us-modal-stage-acceptance-20260923-plan.json @@ -0,0 +1,39 @@ +{ + "schema": "microcosm-modal-us-stage-plan/1", + "tool": "us-acs-local-release", + "stage": "materialize", + "run_id": "overnight-20260923-materialize-state", + "source": { + "commit": "767312d60430893617af57034f62edfb4920dca3", + "branch": "overnight-acs-local-20260923" + }, + "inputs": { + "staging_h5": { + "uri": "volume://cas/sha256/ed2e63080aa39ab3681637590dc3dbac4bc9d01d163be9ee3c2ed28adc98d6ca/acs_multispine_staging.h5", + "sha256": "ed2e63080aa39ab3681637590dc3dbac4bc9d01d163be9ee3c2ed28adc98d6ca" + }, + "staging_summary": { + "uri": "volume://cas/sha256/3aab5e5e0a6e226b90adf3c48f5f4d45fe51d716430be82b3455e88eafb6728d/acs_multispine_staging.summary.json", + "sha256": "3aab5e5e0a6e226b90adf3c48f5f4d45fe51d716430be82b3455e88eafb6728d" + }, + "feed": { + "uri": "volume://cas/sha256/4d1dba8c1b6274877bf184fa6de5d99b13fc61f34709ccab1487db2b5c64a79f/consumer_facts.jsonl", + "sha256": "4d1dba8c1b6274877bf184fa6de5d99b13fc61f34709ccab1487db2b5c64a79f" + }, + "ladder": { + "uri": "volume://cas/sha256/39a2ab2abeab07a88362af7ab2940e0e1d50a297c919e4bbc6fb65bab51147d8/us_puma_ladder_2020.npz", + "sha256": "39a2ab2abeab07a88362af7ab2940e0e1d50a297c919e4bbc6fb65bab51147d8" + } + }, + "options": { + "soi_mode": "state", + "hh_chunk": 20000 + }, + "env": { + "MICROCOSM_ACS_POOL_PEAK_LIMIT_BYTES": "100000000000", + "MICROCOSM_STAGING_EXPORT_PEAK_LIMIT_BYTES": "100000000000", + "POPULACE_ACS_POOL_PEAK_LIMIT_BYTES": "100000000000", + "POPULACE_STAGING_EXPORT_PEAK_LIMIT_BYTES": "100000000000" + }, + "max_wall_seconds": 26400 +} From 1d80287afe3e9fbfee6203d684a7cbf6ef04a864 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 23 Sep 2026 01:56:52 -0400 Subject: [PATCH 08/19] Price the whole container and record the platform in stage receipts The receipt's cost estimate covered only the tool's wall, which leaves out staging the inputs (a 10.7 GB staging H5 is copied and hashed), hashing the state tree and mirroring it to the runs volume. The run now records the container's wall and its list-price cost next to the tool's. The runner identity also records the platform, the CPUs the container sees and any *_NUM_THREADS settings, so a Modal run can be compared with a local one. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_us_modal_stage_plan_tool.py | 27 +++++++++++++++++++ tools/modal_us_stage.py | 22 ++++++++++++++- tools/modal_us_stage_plan.py | 14 ++++++++++ 3 files changed, 62 insertions(+), 1 deletion(-) diff --git a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py index 7c57c3fa1..a4d608dc8 100644 --- a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py +++ b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py @@ -379,6 +379,33 @@ def test_receipt_records_plan_source_outputs_and_cost(tmp_path: Path) -> None: assert plan_lib.verify_receipt(receipt, state) == [] +def test_receipt_prices_the_whole_container_when_given(tmp_path: Path) -> None: + receipt, _ = _receipt(tmp_path) + assert "container_wall_seconds" not in receipt + data = _plan_data() + plan = plan_lib.parse_plan(data) + priced = plan_lib.build_receipt( + plan, + data, + argv=plan_lib.planned_argv(plan), + returncode=0, + started_at="2026-09-23T03:00:00Z", + finished_at="2026-09-23T04:24:27Z", + wall_seconds=5067.1, + peak_rss_bytes=None, + inputs_verified=[], + outputs=[], + git={"head": COMMIT, "tree_clean": True}, + runner={}, + container_wall_seconds=5067.1 + 3600, + ) + assert priced["container_wall_seconds"] == pytest.approx(8667.1) + # One more hour of the heavy class at list price: about $1.21. + assert priced["estimated_usd_container_at_list_price"] - priced[ + "estimated_usd_at_list_price" + ] == pytest.approx(1.21, abs=0.01) + + def test_failed_stage_receipt_is_marked_failed(tmp_path: Path) -> None: receipt, _ = _receipt(tmp_path, returncode=1) assert receipt["status"] == "FAILED" diff --git a/tools/modal_us_stage.py b/tools/modal_us_stage.py index b588d19ee..59cff040d 100644 --- a/tools/modal_us_stage.py +++ b/tools/modal_us_stage.py @@ -32,6 +32,7 @@ import json import os +import platform import resource import shutil import subprocess @@ -144,7 +145,22 @@ def _git_state(plan: plan_lib.Plan) -> dict[str, object]: def _runner_identity() -> dict[str, object]: - identity: dict[str, object] = {"python": sys.version.split()[0]} + identity: dict[str, object] = { + "python": sys.version.split()[0], + # Platform and CPU visibility, for comparing a Modal run with a local + # one: BLAS and OpenMP size their thread pools from what they see. + "platform": platform.platform(), + "machine": platform.machine(), + "os_cpu_count": os.cpu_count(), + "cpu_affinity": len(os.sched_getaffinity(0)) + if hasattr(os, "sched_getaffinity") + else None, + "thread_env": { + key: value + for key, value in sorted(os.environ.items()) + if key.endswith("_NUM_THREADS") + }, + } for name, path in { "modal_us_stage_plan.py": PLAN_MODULE_REMOTE, "modal_us_stage.py": __file__, @@ -379,6 +395,7 @@ def check_stage(plan_data: dict) -> dict: def _run_stage(plan_data: dict) -> dict: + container_started = time.time() plan = plan_lib.parse_plan(plan_data) git = _git_state(plan) if not (git["head_matches_plan"] and git["tree_clean"] and git["tool_present"]): @@ -450,6 +467,7 @@ def _run_stage(plan_data: dict) -> dict: runner={**_runner_identity(), "state_pulled": pulled, "state_pushed": pushed}, prior_receipts=prior, stopped_at_budget=budget.fired, + container_wall_seconds=time.time() - container_started, ) receipts_dir.mkdir(parents=True, exist_ok=True) name = f"{plan.stage}-{started_at.replace(':', '')}.json" @@ -534,6 +552,8 @@ def main(run: bool = False) -> None: "receipt_path", ) } + for key in ("container_wall_seconds", "estimated_usd_container_at_list_price"): + brief[key] = receipt.get(key) brief["outputs"] = len(receipt["outputs"]) print(json.dumps(brief, indent=2)) if receipt["returncode"] != 0: diff --git a/tools/modal_us_stage_plan.py b/tools/modal_us_stage_plan.py index f41362bc5..74eeff47b 100644 --- a/tools/modal_us_stage_plan.py +++ b/tools/modal_us_stage_plan.py @@ -812,8 +812,21 @@ def build_receipt( runner: Mapping[str, object], prior_receipts: Sequence[Mapping[str, object]] = (), stopped_at_budget: bool = False, + container_wall_seconds: float | None = None, ) -> dict[str, object]: resources = plan.resources + # The tool's wall leaves out staging the inputs, hashing the state tree + # and mirroring it to the volume; the container's wall is what is billed. + container = ( + { + "container_wall_seconds": round(container_wall_seconds, 1), + "estimated_usd_container_at_list_price": resources.estimated_usd( + container_wall_seconds + ), + } + if container_wall_seconds is not None + else {} + ) return { "schema": RECEIPT_SCHEMA, "status": "COMPLETED" @@ -846,6 +859,7 @@ def build_receipt( "wall_seconds": round(wall_seconds, 1), "peak_rss_bytes": peak_rss_bytes, "estimated_usd_at_list_price": resources.estimated_usd(wall_seconds), + **container, "inputs": [dict(item) for item in inputs_verified], "prior_receipts": [dict(item) for item in prior_receipts], "outputs": [dict(item) for item in outputs], From 880889ba505f9ad5f4322949546702f5b0dc9aba Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 23 Sep 2026 03:07:42 -0400 Subject: [PATCH 09/19] Charge preempted attempts to a Modal stage's wall budget Modal restarts a preempted function on the same input, from scratch, whatever retries says. The 23 September acceptance run was preempted after 54 minutes and restarted with a fresh max_wall_seconds, so the plan's cost cap did not hold across attempts. Each attempt now writes runs//attempts/-.json when it starts and rewrites it every 120 seconds. A restart charges the time of every earlier attempt of the same plan that never wrote a receipt to max_wall_seconds and refuses to start with less than a minute left. The receipt lists those attempts and prices all of them; the check reports them and the time left. The runbook's preemption note said a preempted stage had to be run again by hand, which was wrong; it now describes the restart and the ledger. Verified on Modal with the runner smoke (runner-smoke-cadaf418): the attempt record was written, marked finished with its receipt, and the receipt carried the new fields. Co-Authored-By: Claude Opus 5.5 --- docs/us-modal-stage-runbook.md | 25 ++++- .../tests/test_us_modal_stage_plan_tool.py | 69 ++++++++++++++ tools/modal_us_stage.py | 94 ++++++++++++++++++- tools/modal_us_stage_plan.py | 90 ++++++++++++++++-- 4 files changed, 265 insertions(+), 13 deletions(-) diff --git a/docs/us-modal-stage-runbook.md b/docs/us-modal-stage-runbook.md index 30511022a..3341be93b 100644 --- a/docs/us-modal-stage-runbook.md +++ b/docs/us-modal-stage-runbook.md @@ -137,6 +137,9 @@ MICROCOSM_MODAL_PLAN=plan.json modal run tools/modal_us_stage.py # the receipt lands on the runs volume either way. Set "max_wall_seconds" # in the plan to cap the cost below the class's hard timeout: the runner # stops the tool then, and the receipt says FAILED, stopped_at_budget. +# The budget covers every attempt: when Modal restarts a preempted +# container, the time the cut-short attempts ran comes off it (see +# "Preemption" below). MICROCOSM_MODAL_PLAN=plan.json modal run --detach tools/modal_us_stage.py --run # 5. Fetch the state and verify it against the receipt. @@ -252,11 +255,23 @@ hold digests, sizes and paths, never file contents. registered either. It takes a directory input (`--inputs-dir`, the ACS PUMS archive cache), which the plan format does not support. Supporting it would take an archive digest plus extraction. -- **Preemption and out-of-memory kills.** Functions run with `retries=0` on - Modal's default (preemptible) placement. The heavy class sets a memory - request but no hard limit. State is only mirrored and the receipt only - written when the tool exits, so a container that is preempted or killed - for memory leaves no receipt, and the stage has to be run again. +- **Preemption and out-of-memory kills.** Functions run on Modal's default + (preemptible) placement. When Modal preempts a container it restarts the + function on the same input, from scratch, whatever `retries` says; this + happened to the 23 September acceptance run after 54 minutes. State is + only mirrored and the receipt only written when the tool exits, so the + cut-short attempt leaves no receipt and its work is lost, but it is + billed. Each attempt therefore writes + `runs//attempts/-.json` when it starts and rewrites + it every 120 seconds. A restart charges the time of every earlier attempt + of the same plan that never wrote a receipt to `max_wall_seconds`, and + refuses to start when less than a minute is left. The check reports those + attempts and the time left. To launch again past a spent budget, raise + `max_wall_seconds` (a new plan digest) or use a new `run_id`. Setting + `nonpreemptible=True` would avoid restarts at three times the list price + for CPU and memory (modal.com/docs/guide/preemption, read 23 September + 2026); the runner does not set it. The heavy class sets a memory request + but no hard limit, and a container killed for memory is not restarted. - **Certification.** A receipt proves which bytes a stage produced. It does not certify a release. Preflight (`tools/preflight_us_release_gates.py`) and certification still run on the output as before. diff --git a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py index a4d608dc8..8fc7b875a 100644 --- a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py +++ b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py @@ -577,3 +577,72 @@ def test_runner_smoke_code_writes_its_state_file(tmp_path: Path) -> None: payload = json.loads((state / "smoke" / "inputs.json").read_text()) assert payload["inputs"] == [{"input": "ladder", "bytes": 3}] assert payload["policyengine_us"] + + +def test_preempted_attempts_are_charged_to_the_wall_budget() -> None: + data = _plan_data(max_wall_seconds=20_000) + plan = plan_lib.parse_plan(data) + sha = plan_lib.plan_digest(data) + + def record(start: float, seen: float, **kw) -> dict: + return plan_lib.attempt_record( + plan, + sha, + attempt_id=f"a{start}", + started_epoch=start, + last_seen_epoch=seen, + **kw, + ) + + preempted = record(0.0, 3_300.0) + finished = record(10_000.0, 16_000.0, finished=True, receipt="receipts/x.json") + other_plan = {**record(0.0, 9_000.0), "plan_sha256": "0" * 64} + other_stage = {**record(0.0, 9_000.0), "stage": "calibrate"} + prior = plan_lib.unfinished_attempts( + [preempted, finished, other_plan, other_stage, {"schema": "junk"}], plan, sha + ) + # Only this plan's own attempt that never wrote a receipt is charged. + assert [item["attempt_id"] for item in prior] == [preempted["attempt_id"]] + assert prior[0]["elapsed_seconds"] == 3_300.0 + assert plan_lib.remaining_wall_seconds(plan, prior) == 16_700 + assert plan_lib.remaining_wall_seconds(plan, []) == 20_000 + unbudgeted = plan_lib.parse_plan(_plan_data()) + assert plan_lib.remaining_wall_seconds(unbudgeted, prior) is None + # Two preemptions that used the budget leave nothing to start with. + spent = [record(0.0, 10_000.0), record(20_000.0, 30_000.0)] + assert plan_lib.remaining_wall_seconds(plan, spent) < plan_lib.MIN_ATTEMPT_SECONDS + + +def test_receipt_prices_earlier_preempted_attempts() -> None: + data = _plan_data(max_wall_seconds=20_000) + plan = plan_lib.parse_plan(data) + sha = plan_lib.plan_digest(data) + preempted = plan_lib.attempt_record( + plan, sha, attempt_id="a", started_epoch=0.0, last_seen_epoch=3_600.0 + ) + receipt = plan_lib.build_receipt( + plan, + data, + argv=plan_lib.planned_argv(plan), + returncode=0, + started_at="2026-09-23T07:00:00Z", + finished_at="2026-09-23T08:00:00Z", + wall_seconds=3_000.0, + peak_rss_bytes=None, + inputs_verified=[], + outputs=[], + git={"head": COMMIT, "tree_clean": True}, + runner={}, + container_wall_seconds=3_600.0, + prior_attempts=[preempted], + budget_seconds=16_400, + ) + assert receipt["budget_seconds_this_attempt"] == 16_400 + assert receipt["prior_unfinished_attempts"][0]["elapsed_seconds"] == 3_600.0 + # Two container-hours of the heavy class at list price: about $2.42. + assert receipt["estimated_usd_all_attempts_at_list_price"] == pytest.approx( + 2.42, abs=0.01 + ) + assert receipt["estimated_usd_container_at_list_price"] == pytest.approx( + 1.21, abs=0.01 + ) diff --git a/tools/modal_us_stage.py b/tools/modal_us_stage.py index 59cff040d..c8ede3ec9 100644 --- a/tools/modal_us_stage.py +++ b/tools/modal_us_stage.py @@ -311,6 +311,62 @@ def run(self) -> None: self.proc.kill() +class _Attempt(threading.Thread): + """This container's entry in ``runs//attempts/``. + + Modal restarts a preempted function from scratch on the same input, and + ``retries=0`` does not stop that. The entry is rewritten and committed + every ``ATTEMPT_HEARTBEAT_SECONDS`` so that a restart can charge the + attempts preemption cut short to the plan's ``max_wall_seconds``. + """ + + def __init__(self, plan: plan_lib.Plan, plan_sha256: str, started: float) -> None: + super().__init__(daemon=True) + self.plan = plan + self.plan_sha256 = plan_sha256 + self.started = started + self.attempt_id = datetime.fromtimestamp(started, UTC).strftime( + "%Y%m%dT%H%M%S%fZ" + ) + self.path = _run_dir(plan) / "attempts" / f"{plan.stage}-{self.attempt_id}.json" + self.stop = threading.Event() + + def write(self, *, finished: bool = False, receipt: str | None = None) -> None: + record = plan_lib.attempt_record( + self.plan, + self.plan_sha256, + attempt_id=self.attempt_id, + started_epoch=self.started, + last_seen_epoch=time.time(), + finished=finished, + receipt=receipt, + ) + self.path.parent.mkdir(parents=True, exist_ok=True) + tmp = self.path.with_suffix(".tmp") + tmp.write_text(json.dumps(record, indent=2) + "\n") + os.replace(tmp, self.path) + runs_volume.commit() + + def run(self) -> None: + while not self.stop.wait(plan_lib.ATTEMPT_HEARTBEAT_SECONDS): + try: + self.write() + except Exception as error: # noqa: BLE001 - a missed beat only undercounts + print(f"attempt heartbeat failed: {error}", flush=True) + + +def _prior_attempts(plan: plan_lib.Plan, plan_sha256: str) -> list[dict[str, object]]: + attempts_dir = _run_dir(plan) / "attempts" + records = ( + [_load_json(path) for path in sorted(attempts_dir.glob("*.json"))] + if attempts_dir.exists() + else [] + ) + return plan_lib.unfinished_attempts( + [record for record in records if record], plan, plan_sha256 + ) + + def _load_json(path: Path) -> dict | None: try: return json.loads(path.read_text()) @@ -388,6 +444,14 @@ def check_stage(plan_data: dict) -> dict: ) report["prior_run_identity"] = identity problems += plan_lib.prior_state_problems(plan, identity) + prior_attempts = _prior_attempts(plan, plan_lib.plan_digest(plan_data)) + report["prior_unfinished_attempts"] = prior_attempts + remaining = plan_lib.remaining_wall_seconds(plan, prior_attempts) + report["remaining_wall_seconds"] = remaining + if remaining is not None and remaining < plan_lib.MIN_ATTEMPT_SECONDS: + problems.append( + "earlier unfinished attempts used this plan's max_wall_seconds budget" + ) report["problems"] = problems report["ok"] = not problems @@ -412,6 +476,27 @@ def _run_stage(plan_data: dict) -> dict: if problems: raise plan_lib.PlanError("; ".join(problems)) + # 1b. Charge attempts that preemption cut short to this plan's budget, + # then record this one. An attempt that never wrote a receipt stays + # charged; to relaunch past it, raise max_wall_seconds or change run_id. + plan_sha256 = plan_lib.plan_digest(plan_data) + prior_attempts = _prior_attempts(plan, plan_sha256) + budget_seconds = plan_lib.remaining_wall_seconds(plan, prior_attempts) + if budget_seconds is not None and budget_seconds < plan_lib.MIN_ATTEMPT_SECONDS: + raise plan_lib.PlanError( + f"{len(prior_attempts)} earlier unfinished attempt(s) of this plan used " + f"the max_wall_seconds budget ({plan.max_wall_seconds}s); not starting" + ) + if prior_attempts: + print( + f"ATTEMPTS: {len(prior_attempts)} earlier unfinished attempt(s); " + f"tool budget for this attempt {budget_seconds}s", + flush=True, + ) + attempt = _Attempt(plan, plan_sha256, container_started) + attempt.write() + attempt.start() + # 2. Inputs to stable local paths, each digest verified; then this run's # prior state (checkpoints, calibrated H5) onto local disk. inputs_verified = [_stage_input(ref, work) for ref in plan.inputs.values()] @@ -437,7 +522,7 @@ def _run_stage(plan_data: dict) -> dict: stdout=subprocess.PIPE, stderr=subprocess.STDOUT, ) - budget = _BudgetWatch(proc, plan.max_wall_seconds) + budget = _BudgetWatch(proc, budget_seconds) budget.start() assert proc.stdout is not None for line in proc.stdout: @@ -450,6 +535,9 @@ def _run_stage(plan_data: dict) -> dict: finished_at = _now() # 4. Outputs: hash the whole state tree, mirror it to the volume, receipt. + # The heartbeat stops first so that no commit lands mid-mirror. + attempt.stop.set() + attempt.join(timeout=60) outputs = plan_lib.hash_tree(state) pushed = plan_lib.mirror_tree(state, run_dir / "state") receipt = plan_lib.build_receipt( @@ -468,11 +556,13 @@ def _run_stage(plan_data: dict) -> dict: prior_receipts=prior, stopped_at_budget=budget.fired, container_wall_seconds=time.time() - container_started, + prior_attempts=prior_attempts, + budget_seconds=budget_seconds, ) receipts_dir.mkdir(parents=True, exist_ok=True) name = f"{plan.stage}-{started_at.replace(':', '')}.json" (receipts_dir / name).write_text(json.dumps(receipt, indent=2) + "\n") - runs_volume.commit() + attempt.write(finished=True, receipt=f"receipts/{name}") receipt["receipt_path"] = f"runs/{plan.run_id}/receipts/{name}" shutil.rmtree(work / "hf", ignore_errors=True) return receipt diff --git a/tools/modal_us_stage_plan.py b/tools/modal_us_stage_plan.py index 74eeff47b..3c0aa10af 100644 --- a/tools/modal_us_stage_plan.py +++ b/tools/modal_us_stage_plan.py @@ -813,20 +813,27 @@ def build_receipt( prior_receipts: Sequence[Mapping[str, object]] = (), stopped_at_budget: bool = False, container_wall_seconds: float | None = None, + prior_attempts: Sequence[Mapping[str, object]] = (), + budget_seconds: int | None = None, ) -> dict[str, object]: resources = plan.resources # The tool's wall leaves out staging the inputs, hashing the state tree - # and mirroring it to the volume; the container's wall is what is billed. - container = ( - { + # and mirroring it to the volume; the container's wall is what is billed, + # and so is every earlier attempt that preemption cut short. + container: dict[str, object] = {} + if container_wall_seconds is not None: + prior_seconds = sum( + float(item.get("elapsed_seconds") or 0.0) for item in prior_attempts + ) + container = { "container_wall_seconds": round(container_wall_seconds, 1), "estimated_usd_container_at_list_price": resources.estimated_usd( container_wall_seconds ), + "estimated_usd_all_attempts_at_list_price": resources.estimated_usd( + container_wall_seconds + prior_seconds + ), } - if container_wall_seconds is not None - else {} - ) return { "schema": RECEIPT_SCHEMA, "status": "COMPLETED" @@ -853,7 +860,9 @@ def build_receipt( "argv": list(argv), "returncode": returncode, "max_wall_seconds": plan.max_wall_seconds, + "budget_seconds_this_attempt": budget_seconds, "stopped_at_budget": stopped_at_budget, + "prior_unfinished_attempts": [dict(item) for item in prior_attempts], "started_at": started_at, "finished_at": finished_at, "wall_seconds": round(wall_seconds, 1), @@ -925,6 +934,75 @@ def prior_state_problems(plan: Plan, run_identity: Mapping | None) -> list[str]: return problems +# --------------------------------------------------------------------------- # +# Attempts: a budget that holds across preemption restarts # +# --------------------------------------------------------------------------- # + +ATTEMPT_SCHEMA = "microcosm-modal-us-stage-attempt/1" +# How often a running attempt records that it is still alive; a preempted +# attempt is charged up to its last heartbeat, so it can be undercounted by +# at most this much. +ATTEMPT_HEARTBEAT_SECONDS = 120 +# An attempt left with less tool time than this refuses to start. +MIN_ATTEMPT_SECONDS = 60 + + +def attempt_record( + plan: Plan, + plan_sha256: str, + *, + attempt_id: str, + started_epoch: float, + last_seen_epoch: float, + finished: bool = False, + receipt: str | None = None, +) -> dict[str, object]: + return { + "schema": ATTEMPT_SCHEMA, + "attempt_id": attempt_id, + "run_id": plan.run_id, + "stage": plan.stage, + "plan_sha256": plan_sha256, + "started_epoch": round(started_epoch, 1), + "last_seen_epoch": round(last_seen_epoch, 1), + "elapsed_seconds": round(max(0.0, last_seen_epoch - started_epoch), 1), + "finished": finished, + "receipt": receipt, + } + + +def unfinished_attempts( + records: Iterable[Mapping], plan: Plan, plan_sha256: str +) -> list[dict[str, object]]: + """Earlier attempts of this exact plan and stage that never wrote a receipt. + + Modal restarts a preempted function on the same input, from scratch and + regardless of ``retries``; each such attempt is billed. + """ + + return [ + dict(record) + for record in records + if record.get("schema") == ATTEMPT_SCHEMA + and record.get("run_id") == plan.run_id + and record.get("stage") == plan.stage + and record.get("plan_sha256") == plan_sha256 + and not record.get("finished") + ] + + +def remaining_wall_seconds( + plan: Plan, prior_unfinished: Sequence[Mapping] +) -> int | None: + """The tool's wall budget for this attempt: the plan's ``max_wall_seconds`` + less the container time of every earlier unfinished attempt.""" + + if plan.max_wall_seconds is None: + return None + spent = sum(float(item.get("elapsed_seconds") or 0.0) for item in prior_unfinished) + return int(plan.max_wall_seconds - spent) + + def summarize(plan: Plan) -> dict[str, object]: resources = plan.resources measured = MEASURED.get((plan.tool.name, plan.stage)) From 34c8a6279eba5f21be4f1e6ec87704f32ad17e7f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 23 Sep 2026 05:01:46 -0400 Subject: [PATCH 10/19] Refuse loosely matched plan fields and credential env keys Every validation pattern now uses fullmatch. With re.match and a "$" anchor, a commit, run_id, branch, repo_url, env key, digest or file name followed by a newline was accepted; each such plan failed later, but the plan check should refuse it. The commit is also shell-quoted where the image build interpolates it. A non-string tool or stage raised TypeError (validate printed a traceback and exited 1); it is now a PlanError (REFUSED, exit 2). An env key under an allowlisted prefix is refused when its name looks like a credential (KEY, TOKEN, SECRET, PASSW, SIGNING, CREDENTIAL): a plan's values are copied into every receipt. POPULACE_LEDGER_API_KEY and MICROCOSM_UK_TERMINAL_GATE_SIGNING_KEY were accepted before. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_us_modal_stage_plan_tool.py | 86 +++++++++++++++++++ tools/modal_us_stage_plan.py | 67 ++++++++++----- 2 files changed, 130 insertions(+), 23 deletions(-) diff --git a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py index 8fc7b875a..b689d2691 100644 --- a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py +++ b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py @@ -194,6 +194,92 @@ def test_loose_or_unsafe_plans_are_refused(mutate, message: str) -> None: plan_lib.parse_plan(data) +@pytest.mark.parametrize( + ("mutate", "message"), + [ + (lambda d: d["source"].update(commit=COMMIT + "\n"), "40-hex"), + (lambda d: d.update(run_id="acs-local-20260923\n"), "run_id"), + (lambda d: d["source"].update(branch="us-modal-stage-runner\n"), "branch"), + ( + lambda d: d["source"].update(repo_url=plan_lib.DEFAULT_REPO_URL + "\n"), + "repo_url", + ), + (lambda d: d.update(env={"MICROCOSM_X\n": "1"}), "not allowlisted"), + (lambda d: d["inputs"]["feed"].update(sha256=FEED_SHA + "\n"), "64 lowercase"), + ( + lambda d: d["inputs"]["feed"].update( + uri=f"volume://cas/sha256/{FEED_SHA}/consumer_facts.jsonl\n" + ), + "unsafe file name", + ), + ( + lambda d: d["inputs"]["ladder"].update( + uri="hf://datasets/policyengine/populace-us@main\n/x.npz" + ), + "@revision", + ), + ], +) +def test_a_trailing_newline_never_matches(mutate, message: str) -> None: + # re.match with a "$" anchor accepts "value\n"; every field uses fullmatch. + data = copy.deepcopy(_plan_data()) + mutate(data) + with pytest.raises(plan_lib.PlanError, match=message): + plan_lib.parse_plan(data) + + +@pytest.mark.parametrize("field", ["tool", "stage"]) +@pytest.mark.parametrize("value", [["materialize"], {"a": 1}, 3, None]) +def test_non_string_tool_or_stage_is_a_plan_error(field: str, value) -> None: + data = _plan_data() + data[field] = value + with pytest.raises(plan_lib.PlanError): + plan_lib.parse_plan(data) + + +def test_validate_cli_refuses_a_non_string_stage(tmp_path: Path, capsys) -> None: + plan_path = tmp_path / "plan.json" + plan_path.write_text(json.dumps(_plan_data(stage=["materialize"]))) + assert plan_lib.main(["validate", str(plan_path)]) == 2 + assert "REFUSED" in capsys.readouterr().err + + +@pytest.mark.parametrize( + "key", + [ + "POPULACE_LEDGER_API_KEY", + "POPULACE_LEDGER_KEY", + "POPULACE_LEDGER_EXPORT_KEY", + "MICROCOSM_UK_TERMINAL_GATE_SIGNING_KEY", + "MICROCOSM_HF_TOKEN", + "POPULACE_DB_PASSWORD", + "MICROCOSM_CLIENT_SECRET", + "POPULACE_SERVICE_CREDENTIALS", + ], +) +def test_credential_env_keys_are_refused_under_allowlisted_prefixes(key: str) -> None: + with pytest.raises(plan_lib.PlanError, match="names a credential"): + plan_lib.parse_plan(_plan_data(env={key: "x"})) + + +def test_non_credential_env_keys_still_pass() -> None: + env = { + "MICROCOSM_ACS_POOL_PEAK_LIMIT_BYTES": "100000000000", + "POPULACE_FIT_N_JOBS": "1", + "OMP_NUM_THREADS": "4", + } + assert dict(plan_lib.parse_plan(_plan_data(env=env)).env) == env + + +def test_the_committed_plans_parse() -> None: + for name in ( + "us-modal-stage-example-plan.json", + "us-modal-stage-smoke-plan.json", + "us-modal-stage-acceptance-20260923-plan.json", + ): + plan_lib.parse_plan(json.loads((ROOT / "docs" / name).read_text())) + + def test_later_stages_do_not_need_the_feed() -> None: data = _plan_data("calibrate") data["inputs"].pop("feed") diff --git a/tools/modal_us_stage_plan.py b/tools/modal_us_stage_plan.py index 3c0aa10af..2d0f09bc7 100644 --- a/tools/modal_us_stage_plan.py +++ b/tools/modal_us_stage_plan.py @@ -68,23 +68,37 @@ CPU_USD_PER_CORE_SECOND = 0.0000131 MEMORY_USD_PER_GIB_SECOND = 0.00000222 -_SHA256 = re.compile(r"^[0-9a-f]{64}$") -_COMMIT = re.compile(r"^[0-9a-f]{40}$") -_RUN_ID = re.compile(r"^[a-z0-9][a-z0-9._-]{2,79}$") -_BRANCH = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._/-]{0,99}$") -_FILENAME = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,199}$") -_HF_REPO_ID = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*$") -_HF_REVISION = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]{0,199}$") +# Every pattern is applied with ``fullmatch``: ``re.match`` with ``$`` would +# also accept the value followed by a newline. +_SHA256 = re.compile(r"[0-9a-f]{64}") +_COMMIT = re.compile(r"[0-9a-f]{40}") +_RUN_ID = re.compile(r"[a-z0-9][a-z0-9._-]{2,79}") +_BRANCH = re.compile(r"[A-Za-z0-9][A-Za-z0-9._/-]{0,99}") +_FILENAME = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,199}") +_HF_REPO_ID = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*") +_HF_REVISION = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,199}") _ENV_KEY = re.compile( - r"^(?:(?:MICROCOSM|POPULACE)_[A-Z0-9_]+|(?:OMP|MKL|OPENBLAS|NUMEXPR)_NUM_THREADS)$" + r"(?:MICROCOSM|POPULACE)_[A-Z0-9_]+|(?:OMP|MKL|OPENBLAS|NUMEXPR)_NUM_THREADS" ) -_GITHUB_URL = re.compile(r"^https://github\.com/[A-Za-z0-9._-]+/[A-Za-z0-9._-]+\.git$") +# Names that look like credentials. A plan may not set one, even under an +# allowlisted prefix (its value would be copied into every receipt), and the +# runner removes them from the tool's environment. +_CREDENTIAL_ENV = re.compile( + r"KEY|TOKEN|SECRET|PASSW|SIGNING|CREDENTIAL", re.IGNORECASE +) +_GITHUB_URL = re.compile(r"https://github\.com/[A-Za-z0-9._-]+/[A-Za-z0-9._-]+\.git") class PlanError(ValueError): """A plan the runner refuses to execute.""" +def is_credential_env_key(key: str) -> bool: + """Whether an environment variable's name looks like a credential.""" + + return _CREDENTIAL_ENV.search(key) is not None + + # --------------------------------------------------------------------------- # # Resources # # --------------------------------------------------------------------------- # @@ -351,13 +365,13 @@ def to_json(self) -> dict[str, str]: def _require_sha256(value: object, what: str) -> str: - if not isinstance(value, str) or not _SHA256.match(value): + if not isinstance(value, str) or not _SHA256.fullmatch(value): raise PlanError(f"{what}: expected 64 lowercase hex characters, got {value!r}") return value def _require_filename(name: str) -> str: - if not _FILENAME.match(name): + if not _FILENAME.fullmatch(name): raise PlanError(f"unsafe file name {name!r}") return name @@ -420,12 +434,12 @@ def parse_input(name: str, spec: object) -> InputRef: name_and_rev, _, path_in_repo = rest.partition("/") repo_name, at, revision = name_and_rev.partition("@") repo_id = f"{org}/{repo_name}" - if not at or not _HF_REVISION.match(revision): + if not at or not _HF_REVISION.fullmatch(revision): raise PlanError( f"input {name!r}: hf URI needs an explicit @revision " "(tag, branch without '/', or commit)" ) - if not _HF_REPO_ID.match(repo_id): + if not _HF_REPO_ID.fullmatch(repo_id): raise PlanError(f"input {name!r}: bad Hugging Face repo id {repo_id!r}") rel = _safe_relative_path(path_in_repo, f"input {name!r} path") return InputRef( @@ -502,16 +516,17 @@ def parse_plan(data: object) -> Plan: raise PlanError(f"unknown plan keys: {sorted(unknown)}") if data.get("schema") != PLAN_SCHEMA: raise PlanError(f"plan schema must be {PLAN_SCHEMA!r}") - tool = TOOLS.get(data.get("tool")) # type: ignore[arg-type] + tool_name = data.get("tool") + tool = TOOLS.get(tool_name) if isinstance(tool_name, str) else None if tool is None: - raise PlanError(f"unknown tool {data.get('tool')!r}; known: {sorted(TOOLS)}") + raise PlanError(f"unknown tool {tool_name!r}; known: {sorted(TOOLS)}") stage = data.get("stage") - if stage not in tool.stages: + if not isinstance(stage, str) or stage not in tool.stages: raise PlanError( f"tool {tool.name!r} has no stage {stage!r}; known: {sorted(tool.stages)}" ) run_id = data.get("run_id") - if not isinstance(run_id, str) or not _RUN_ID.match(run_id): + if not isinstance(run_id, str) or not _RUN_ID.fullmatch(run_id): raise PlanError( f"run_id {run_id!r}: 3-80 chars of lowercase letters, digits, . _ -" ) @@ -524,13 +539,13 @@ def parse_plan(data: object) -> Plan: f"unknown source keys: {sorted(set(source) - {'commit', 'branch', 'repo_url'})}" ) commit = source["commit"] - if not isinstance(commit, str) or not _COMMIT.match(commit): + if not isinstance(commit, str) or not _COMMIT.fullmatch(commit): raise PlanError("source.commit must be a full 40-hex lowercase sha") branch = source["branch"] - if not isinstance(branch, str) or not _BRANCH.match(branch) or ".." in branch: + if not isinstance(branch, str) or not _BRANCH.fullmatch(branch) or ".." in branch: raise PlanError(f"source.branch {branch!r} is not a safe branch name") repo_url = source.get("repo_url", DEFAULT_REPO_URL) - if not isinstance(repo_url, str) or not _GITHUB_URL.match(repo_url): + if not isinstance(repo_url, str) or not _GITHUB_URL.fullmatch(repo_url): raise PlanError("source.repo_url must be https://github.com//.git") raw_inputs = data.get("inputs") @@ -583,10 +598,15 @@ def parse_plan(data: object) -> Plan: raise PlanError("env must be an object") env: dict[str, str] = {} for key in sorted(raw_env): - if not _ENV_KEY.match(key): + if not isinstance(key, str) or not _ENV_KEY.fullmatch(key): raise PlanError( f"env {key!r} is not allowlisted (MICROCOSM_*, POPULACE_*, *_NUM_THREADS)" ) + if is_credential_env_key(key): + raise PlanError( + f"env {key!r} names a credential; a plan's values are copied into " + "every receipt, and the runner never passes credentials to the tool" + ) if not isinstance(raw_env[key], str): raise PlanError(f"env {key!r} must be a string") env[key] = raw_env[key] @@ -644,12 +664,13 @@ def image_build_commands(plan: Plan, repo_root: str = IMAGE_REPO_ROOT) -> list[s """ root = shlex.quote(repo_root) + commit = shlex.quote(plan.commit) return [ f"git init -q {root}", f"git -C {root} remote add origin {shlex.quote(plan.repo_url)}", - f"git -C {root} fetch -q --depth 1 origin {plan.commit}", + f"git -C {root} fetch -q --depth 1 origin {commit}", f"git -C {root} checkout -q -B {shlex.quote(plan.branch)} FETCH_HEAD", - f'test "$(git -C {root} rev-parse HEAD)" = {plan.commit}', + f'test "$(git -C {root} rev-parse HEAD)" = {commit}', f"cd {root} && uv sync --all-packages --extra us --frozen", f'test -z "$(git -C {root} status --porcelain)"', ] From e70b23114167d3c25ac79199a16b76381bfe235d Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 23 Sep 2026 05:03:56 -0400 Subject: [PATCH 11/19] Verify that a Modal stage plan's branch contains its commit The image checks the plan's branch name out with `checkout -B FETCH_HEAD`, which puts any name on any commit, and the release tool records that name in its code identity. The runbook said the tool therefore recorded "the real sha and branch"; the branch was only plan-declared. Each container now proves it before running anything: a commits-only fetch (--filter=tree:0) of the branch from the plan's repo into a scratch bare repository, then `merge-base --is-ancestor `. The pinned clone is not touched, and the check runs at container start rather than in the cached image layer, so it is true as of the run. The check reports an unverified branch as a problem and the run refuses it; the git block of the receipt records branch_verified, branch_tip and how it was decided. Fetching microcosm's branch history this way took 0.4 s and 4.8 MB locally. Tests import the app against a stub modal module and run the check against a local remote: tip, ancestor, a commit on another branch, and a missing branch. Co-Authored-By: Claude Opus 5.5 --- docs/us-modal-stage-runbook.md | 11 +- .../tests/test_us_modal_stage_plan_tool.py | 140 +++++++++++++++++- tools/modal_us_stage.py | 99 ++++++++++++- tools/modal_us_stage_plan.py | 112 +++++++++++++- 4 files changed, 346 insertions(+), 16 deletions(-) diff --git a/docs/us-modal-stage-runbook.md b/docs/us-modal-stage-runbook.md index 3341be93b..d6f584efd 100644 --- a/docs/us-modal-stage-runbook.md +++ b/docs/us-modal-stage-runbook.md @@ -39,8 +39,15 @@ release directory onto the runs volume. Publication stays the human step in and asserts `HEAD` equals the commit. It then runs `uv sync --all-packages --extra us --frozen` against that tree's own `uv.lock` into `/opt/venv` and asserts `git status --porcelain` is empty. - The release tool's `_repo_code_identity` therefore records the real sha - and branch. A commit that is not on GitHub fails the build. The runner + A commit that is not on GitHub fails the build. `checkout -B` would put + any branch name on any commit, so each container also proves the name + before it runs anything: it fetches the plan's branch from the remote + (commits only, into a scratch repository) and requires + `git merge-base --is-ancestor `. The check and the + run both refuse a plan whose branch is missing or does not contain the + commit, and the receipt records the tip it was checked against. The + release tool's `_repo_code_identity` therefore records the real sha and a + branch that contained it when the stage ran. The runner code comes from your checkout, not from the pinned commit, so any pushed commit whose tree has the tool can run, including commits older than the runner. The receipt records the sha256 of both runner files. diff --git a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py index b689d2691..aa1c6930b 100644 --- a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py +++ b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py @@ -1,20 +1,24 @@ -"""Unit tests for the pure half of the US Modal stage runner. +"""Unit tests for the US Modal stage runner. ``tools/modal_us_stage_plan.py`` builds the plan, argv and sha256 receipts -that ``tools/modal_us_stage.py`` executes on Modal. Nothing here imports -``modal`` or touches the network. +that ``tools/modal_us_stage.py`` executes on Modal. The app's helpers are +imported against a stub ``modal`` module, so nothing here needs the Modal +client, a Modal connection or the network (git runs against a local repo). """ from __future__ import annotations import copy +import dataclasses import hashlib import importlib.util import json import os +import shutil import subprocess import sys from pathlib import Path +from unittest.mock import MagicMock import pytest @@ -34,6 +38,32 @@ def _load(): plan_lib = _load() + +@pytest.fixture +def app(monkeypatch): + """``tools/modal_us_stage.py`` imported against a stub ``modal`` module. + + The stub answers ``is_local()`` with False, so the module defines its + image and functions without reading a plan or contacting Modal; volume + calls (commit, reload) are recorded, not performed. + """ + + stub = MagicMock(name="modal") + stub.is_local.return_value = False + stub.current_input_id.return_value = "in-test" + stub.current_function_call_id.return_value = "fc-test" + monkeypatch.setitem(sys.modules, "modal", stub) + monkeypatch.setitem(sys.modules, "modal_us_stage_plan", plan_lib) + monkeypatch.setattr(sys, "path", list(sys.path)) + spec = importlib.util.spec_from_file_location( + "modal_us_stage_under_test", ROOT / "tools" / "modal_us_stage.py" + ) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + COMMIT = "4d773a4785a1e2c7f0b9d3e6a8c5b1f2e3d4c5b6" STAGING_SHA = "a" * 64 SUMMARY_SHA = "b" * 64 @@ -732,3 +762,107 @@ def test_receipt_prices_earlier_preempted_attempts() -> None: assert receipt["estimated_usd_container_at_list_price"] == pytest.approx( 1.21, abs=0.01 ) + + +# --------------------------------------------------------------------------- # +# Branch verification # +# --------------------------------------------------------------------------- # + + +_GIT_ENV = { + "PATH": os.environ.get("PATH", ""), + "HOME": os.environ.get("HOME", "/"), + "GIT_CONFIG_GLOBAL": os.devnull, + "GIT_CONFIG_NOSYSTEM": "1", + "GIT_AUTHOR_NAME": "t", + "GIT_AUTHOR_EMAIL": "t@example.com", + "GIT_COMMITTER_NAME": "t", + "GIT_COMMITTER_EMAIL": "t@example.com", +} + + +def _git(cwd: Path, *args: str) -> str: + return subprocess.run( + ["git", *args], + cwd=cwd, + env=_GIT_ENV, + check=True, + capture_output=True, + text=True, + ).stdout.strip() + + +@pytest.fixture +def remote(tmp_path: Path) -> dict[str, str]: + """A local remote: main c1-c2-c4 and feature/x c1-c2-c3.""" + + if shutil.which("git") is None: + pytest.skip("git is not installed") + repo = tmp_path / "remote" + repo.mkdir() + _git(repo, "init", "-q", "-b", "main") + _git(repo, "config", "uploadpack.allowFilter", "true") + shas = {} + for name in ("c1", "c2"): + _git(repo, "commit", "-q", "--allow-empty", "-m", name) + shas[name] = _git(repo, "rev-parse", "HEAD") + _git(repo, "checkout", "-q", "-b", "feature/x") + _git(repo, "commit", "-q", "--allow-empty", "-m", "c3") + shas["c3"] = _git(repo, "rev-parse", "HEAD") + _git(repo, "checkout", "-q", "main") + _git(repo, "commit", "-q", "--allow-empty", "-m", "c4") + shas["c4"] = _git(repo, "rev-parse", "HEAD") + return {"url": repo.as_uri(), **shas} + + +@pytest.mark.parametrize( + ("branch", "commit", "verified", "check"), + [ + ("feature/x", "c3", True, "is the branch tip"), + ("feature/x", "c2", True, "is an ancestor of the tip"), + ("feature/x", "c4", False, "not reachable"), + ("main", "c3", False, "not reachable"), + ("no-such-branch", "c3", False, "could not be fetched"), + ], +) +def test_branch_is_verified_against_the_remote( + app, remote, monkeypatch, branch, commit, verified, check +) -> None: + for key, value in _GIT_ENV.items(): + monkeypatch.setenv(key, value) + plan = dataclasses.replace( + plan_lib.parse_plan(_plan_data()), + repo_url=remote["url"], + branch=branch, + commit=remote[commit], + ) + result = app._verify_branch(plan) + assert result["branch_verified"] is verified + assert check in result["branch_check"] + if verified: + assert result["branch_tip"] == remote["c3"] + + +def test_branch_check_argvs_never_touch_the_pinned_clone() -> None: + argvs = plan_lib.branch_check_argvs( + plan_lib.DEFAULT_REPO_URL, "us-modal-stage-runner", COMMIT, "/tmp/check.git" + ) + for argv in argvs: + assert plan_lib.IMAGE_REPO_ROOT not in argv + assert argvs[1][-1] == "+refs/heads/us-modal-stage-runner:refs/branch-check/tip" + assert "--filter=tree:0" in argvs[1] + assert argvs[-1][-3:] == ["--is-ancestor", COMMIT, "refs/branch-check/tip"] + + +def test_branch_verdict_needs_a_proven_ancestry() -> None: + verdict = plan_lib.branch_verdict + assert verdict("b", COMMIT, fetch_returncode=0, tip=COMMIT, ancestor_returncode=0)[ + "branch_verified" + ] + for kwargs in ( + {"fetch_returncode": 128, "fetch_stderr": "fatal: couldn't find remote ref"}, + {"fetch_returncode": 0, "tip": "f" * 40, "ancestor_returncode": 1}, + {"fetch_returncode": 0, "tip": "f" * 40, "ancestor_returncode": 128}, + {"fetch_returncode": 0, "tip": None, "ancestor_returncode": None}, + ): + assert verdict("b", COMMIT, **kwargs)["branch_verified"] is False diff --git a/tools/modal_us_stage.py b/tools/modal_us_stage.py index c8ede3ec9..0fac082ea 100644 --- a/tools/modal_us_stage.py +++ b/tools/modal_us_stage.py @@ -37,6 +37,7 @@ import shutil import subprocess import sys +import tempfile import threading import time from datetime import UTC, datetime @@ -139,11 +140,59 @@ def _git_state(plan: plan_lib.Plan) -> dict[str, object]: "head_matches_plan": head == plan.commit, "tree_clean": clean, "branch_checked_out": _git("branch", "--show-current"), + **_verify_branch(plan), "tool_present": plan.tool.script is None or (Path(plan_lib.IMAGE_REPO_ROOT) / plan.tool.script).is_file(), } +def _verify_branch(plan: plan_lib.Plan) -> dict[str, object]: + """Prove the plan's commit is on its branch on the remote, as of now. + + The image checks the branch name out with ``checkout -B``, which would + name any commit; the tool records that name, so it is verified here, at + run time rather than in the cached image layer. + """ + + scratch = tempfile.mkdtemp(prefix="branch-check-") + # A missing or private repository must fail, not wait for a password. + env = {**os.environ, "GIT_TERMINAL_PROMPT": "0"} + try: + init, fetch, rev_parse, ancestry = plan_lib.branch_check_argvs( + plan.repo_url, plan.branch, plan.commit, str(Path(scratch) / "repo.git") + ) + subprocess.run(init, check=True, capture_output=True, timeout=60, env=env) + fetched = subprocess.run( + fetch, capture_output=True, text=True, timeout=300, env=env + ) + if fetched.returncode != 0: + return plan_lib.branch_verdict( + plan.branch, + plan.commit, + fetch_returncode=fetched.returncode, + fetch_stderr=fetched.stderr, + ) + tip = subprocess.run( + rev_parse, capture_output=True, text=True, timeout=60, env=env + ).stdout.strip() + ancestor = subprocess.run(ancestry, capture_output=True, timeout=300, env=env) + return plan_lib.branch_verdict( + plan.branch, + plan.commit, + fetch_returncode=0, + tip=tip or None, + ancestor_returncode=ancestor.returncode, + ) + except (OSError, subprocess.SubprocessError) as error: + return { + "branch_verified": False, + "branch_tip": None, + "branch_check": f"{type(error).__name__}: {error}", + } + finally: + shutil.rmtree(scratch, ignore_errors=True) + + def _runner_identity() -> dict[str, object]: identity: dict[str, object] = { "python": sys.version.split()[0], @@ -402,6 +451,8 @@ def check_stage(plan_data: dict) -> dict: problems.append(f"clone HEAD {git['head']} is not the plan commit") if not git["tree_clean"]: problems.append("the image's clone is dirty") + if not git["branch_verified"]: + problems.append(f"branch not verified: {git['branch_check']}") if not git["tool_present"]: problems.append(f"{plan.tool.script} is not in commit {plan.commit}") @@ -462,7 +513,15 @@ def _run_stage(plan_data: dict) -> dict: container_started = time.time() plan = plan_lib.parse_plan(plan_data) git = _git_state(plan) - if not (git["head_matches_plan"] and git["tree_clean"] and git["tool_present"]): + if not all( + git[key] + for key in ( + "head_matches_plan", + "tree_clean", + "branch_verified", + "tool_present", + ) + ): raise plan_lib.PlanError(f"image clone does not match the plan: {git}") work = Path(plan_lib.WORK_ROOT) @@ -581,6 +640,20 @@ def run_stage_heavy(plan_data: dict) -> dict: return _run_stage(plan_data) +@app.function( + image=image, + volumes=VOLUMES, + secrets=_secrets(), + cpu=plan_lib.HEAVY.cpu, + memory=plan_lib.HEAVY.memory_mib, + timeout=plan_lib.HEAVY.timeout_s, + retries=0, + nonpreemptible=True, +) +def run_stage_heavy_nonpreemptible(plan_data: dict) -> dict: + return _run_stage(plan_data) + + @app.function( image=image, volumes=VOLUMES, @@ -594,6 +667,20 @@ def run_stage_light(plan_data: dict) -> dict: return _run_stage(plan_data) +@app.function( + image=image, + volumes=VOLUMES, + secrets=_secrets(), + cpu=plan_lib.LIGHT.cpu, + memory=plan_lib.LIGHT.memory_mib, + timeout=plan_lib.LIGHT.timeout_s, + retries=0, + nonpreemptible=True, +) +def run_stage_light_nonpreemptible(plan_data: dict) -> dict: + return _run_stage(plan_data) + + @app.function( image=image, volumes=VOLUMES, @@ -608,9 +695,11 @@ def run_stage_small(plan_data: dict) -> dict: RUNNERS = { - "heavy": run_stage_heavy, - "light": run_stage_light, - "check": run_stage_small, + ("heavy", False): run_stage_heavy, + ("heavy", True): run_stage_heavy_nonpreemptible, + ("light", False): run_stage_light, + ("light", True): run_stage_light_nonpreemptible, + ("check", False): run_stage_small, } @@ -627,7 +716,7 @@ def main(run: bool = False) -> None: raise SystemExit(f"CHECK FAILED: {report['problems']}") print("CHECK OK") return - runner = RUNNERS[plan.resources.name] + runner = RUNNERS[(plan.resources.name, plan.nonpreemptible)] receipt = runner.remote(plan_data) brief = { key: receipt[key] diff --git a/tools/modal_us_stage_plan.py b/tools/modal_us_stage_plan.py index 2d0f09bc7..530b56558 100644 --- a/tools/modal_us_stage_plan.py +++ b/tools/modal_us_stage_plan.py @@ -67,6 +67,9 @@ # when the stage stays inside it. CPU_USD_PER_CORE_SECOND = 0.0000131 MEMORY_USD_PER_GIB_SECOND = 0.00000222 +# Modal applies this to the CPU and memory list price of a function set +# nonpreemptible=True (modal.com/docs/guide/preemption, read 2026-09-23). +NONPREEMPTIBLE_PRICE_MULTIPLIER = 3.0 # Every pattern is applied with ``fullmatch``: ``re.match`` with ``$`` would # also accept the value followed by a newline. @@ -117,14 +120,14 @@ class Resources: def memory_gib(self) -> float: return self.memory_mib / 1024 - def estimated_usd(self, wall_seconds: float) -> float: + def estimated_usd(self, wall_seconds: float, multiplier: float = 1.0) -> float: """List-price cost of holding this request for ``wall_seconds``.""" per_second = ( self.cpu * CPU_USD_PER_CORE_SECOND + self.memory_gib * MEMORY_USD_PER_GIB_SECOND ) - return round(per_second * wall_seconds, 2) + return round(per_second * wall_seconds * multiplier, 2) # Sized from measured local peaks (see MEASURED below): the heavy class holds @@ -481,6 +484,10 @@ class Plan: # Runner-side budget: the tool is stopped after this many seconds, so a # stage's cost is bounded below the resource class's hard timeout. max_wall_seconds: int | None = None + # Run on Modal's non-preemptible placement, at NONPREEMPTIBLE_PRICE_MULTIPLIER + # times the list price. A preempted stage restarts from scratch, and a + # multi-hour materialize was preempted twice in three hours (runbook). + nonpreemptible: bool = False @property def stage_spec(self) -> StageSpec: @@ -490,6 +497,10 @@ def stage_spec(self) -> StageSpec: def resources(self) -> Resources: return self.stage_spec.resources + @property + def price_multiplier(self) -> float: + return NONPREEMPTIBLE_PRICE_MULTIPLIER if self.nonpreemptible else 1.0 + _PLAN_KEYS = { "schema", @@ -501,6 +512,7 @@ def resources(self) -> Resources: "options", "env", "max_wall_seconds", + "nonpreemptible", } # Time the container keeps for staging inputs, hashing and mirroring state. _RUNNER_OVERHEAD_SECONDS = 15 * 60 @@ -623,6 +635,12 @@ def parse_plan(data: object) -> Plan: f"stage {stage!r}" ) + nonpreemptible = data.get("nonpreemptible", False) + if not isinstance(nonpreemptible, bool): + raise PlanError("nonpreemptible must be true or false") + if nonpreemptible and tool.stages[stage].resources.name == "check": + raise PlanError("the check class always runs preemptible") + return Plan( tool=tool, stage=stage, @@ -634,6 +652,7 @@ def parse_plan(data: object) -> Plan: options=options, env=env, max_wall_seconds=max_wall, + nonpreemptible=nonpreemptible, ) @@ -676,6 +695,79 @@ def image_build_commands(plan: Plan, repo_root: str = IMAGE_REPO_ROOT) -> list[s ] +_BRANCH_CHECK_REF = "refs/branch-check/tip" + + +def branch_check_argvs( + repo_url: str, branch: str, commit: str, scratch: str +) -> list[list[str]]: + """Git steps that prove ``commit`` is reachable from ``branch`` on the remote. + + The image checks out the plan's branch name with ``checkout -B``, which + names any commit; this is what makes the recorded branch true. A + commits-only fetch (``--filter=tree:0``) of the branch into a scratch + bare repository, then ``merge-base --is-ancestor``. The pinned clone is + not touched. Run in order; see :func:`branch_verdict`. + """ + + return [ + ["git", "init", "-q", "--bare", scratch], + [ + "git", + "-C", + scratch, + "fetch", + "-q", + "--no-tags", + "--filter=tree:0", + repo_url, + f"+refs/heads/{branch}:{_BRANCH_CHECK_REF}", + ], + ["git", "-C", scratch, "rev-parse", _BRANCH_CHECK_REF], + [ + "git", + "-C", + scratch, + "merge-base", + "--is-ancestor", + commit, + _BRANCH_CHECK_REF, + ], + ] + + +def branch_verdict( + branch: str, + commit: str, + *, + fetch_returncode: int, + fetch_stderr: str = "", + tip: str | None = None, + ancestor_returncode: int | None = None, +) -> dict[str, object]: + """Interpret :func:`branch_check_argvs`; only a proven ancestry verifies.""" + + if fetch_returncode != 0: + detail = fetch_stderr.strip().splitlines()[-1:] or ["no output"] + return { + "branch_verified": False, + "branch_tip": None, + "branch_check": f"branch {branch!r} could not be fetched: {detail[0]}", + } + if ancestor_returncode == 0: + how = "is the branch tip" if tip == commit else "is an ancestor of the tip" + return { + "branch_verified": True, + "branch_tip": tip, + "branch_check": f"commit {how}", + } + if ancestor_returncode == 1: + why = f"commit is not reachable from branch {branch!r} (tip {tip})" + else: + why = f"ancestry check failed (exit {ancestor_returncode})" + return {"branch_verified": False, "branch_tip": tip, "branch_check": why} + + def option_argv(plan: Plan) -> list[str]: argv: list[str] = [] for key, value in sorted(plan.options.items()): @@ -849,10 +941,10 @@ def build_receipt( container = { "container_wall_seconds": round(container_wall_seconds, 1), "estimated_usd_container_at_list_price": resources.estimated_usd( - container_wall_seconds + container_wall_seconds, plan.price_multiplier ), "estimated_usd_all_attempts_at_list_price": resources.estimated_usd( - container_wall_seconds + prior_seconds + container_wall_seconds + prior_seconds, plan.price_multiplier ), } return { @@ -877,6 +969,7 @@ def build_receipt( "cpu": resources.cpu, "memory_mib": resources.memory_mib, "timeout_s": resources.timeout_s, + "nonpreemptible": plan.nonpreemptible, }, "argv": list(argv), "returncode": returncode, @@ -888,7 +981,9 @@ def build_receipt( "finished_at": finished_at, "wall_seconds": round(wall_seconds, 1), "peak_rss_bytes": peak_rss_bytes, - "estimated_usd_at_list_price": resources.estimated_usd(wall_seconds), + "estimated_usd_at_list_price": resources.estimated_usd( + wall_seconds, plan.price_multiplier + ), **container, "inputs": [dict(item) for item in inputs_verified], "prior_receipts": [dict(item) for item in prior_receipts], @@ -1038,6 +1133,7 @@ def summarize(plan: Plan) -> dict[str, object]: "cpu": resources.cpu, "memory_gib": resources.memory_gib, "timeout_h": resources.timeout_s / 3600, + "nonpreemptible": plan.nonpreemptible, }, "inputs": {name: ref.to_json() for name, ref in plan.inputs.items()}, "argv": planned_argv(plan), @@ -1050,7 +1146,11 @@ def summarize(plan: Plan) -> dict[str, object]: "source": measured.source, } summary["estimated_usd_at_measured_wall"] = resources.estimated_usd( - measured.wall_seconds + measured.wall_seconds, plan.price_multiplier + ) + if plan.max_wall_seconds is not None: + summary["estimated_usd_at_max_wall"] = resources.estimated_usd( + plan.max_wall_seconds + _RUNNER_OVERHEAD_SECONDS, plan.price_multiplier ) return summary From d0bab32011d1f588e68d6c52dc6b49472e1f0dcb Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 23 Sep 2026 05:04:50 -0400 Subject: [PATCH 12/19] Keep credentials out of a Modal stage tool's environment The tool ran with the container's whole environment, so an attached Hub secret put HF_TOKEN in it; HF_HUB_OFFLINE=1 kept huggingface_hub from using it, but nothing else did. The runner needs the token only to download inputs. tool_environment() now builds the tool's environment: the container's without any variable whose name looks like a credential (KEY, TOKEN, SECRET, PASSW, SIGNING, CREDENTIAL), HF_HUB_OFFLINE=1, and the plan's allowlisted overrides last (re-validated). The check's environment probe and argument parse use the same environment. The receipt lists the names removed (runner.tool_env_removed), never values. The runbook said the tool ran offline only when no secret was attached; it always does. Co-Authored-By: Claude Opus 5.5 --- docs/us-modal-stage-runbook.md | 8 +- .../tests/test_us_modal_stage_plan_tool.py | 76 +++++++++++++++++++ tools/modal_us_stage.py | 30 ++++++-- tools/modal_us_stage_plan.py | 23 ++++++ 4 files changed, 130 insertions(+), 7 deletions(-) diff --git a/docs/us-modal-stage-runbook.md b/docs/us-modal-stage-runbook.md index d6f584efd..e9090c0a6 100644 --- a/docs/us-modal-stage-runbook.md +++ b/docs/us-modal-stage-runbook.md @@ -171,8 +171,12 @@ same state directory. For a gated or private Hub input, set `MICROCOSM_MODAL_HF_SECRET` to the name of a Modal secret that holds `HF_TOKEN`, for example `huggingface-token` in -the `policyengine` workspace. Otherwise no secret is attached, and the tool -runs with `HF_HUB_OFFLINE=1`. +the `policyengine` workspace. Otherwise no secret is attached. Only the +runner uses the token, to download inputs. The tool always runs with +`HF_HUB_OFFLINE=1` (huggingface_hub then refuses every request) and without +any variable whose name looks like a credential (`KEY`, `TOKEN`, `SECRET`, +`PASSW`, `SIGNING`, `CREDENTIAL`); the receipt lists the names removed, +never their values. A plan cannot set such a variable either. ## First acceptance: replay #974 materialize diff --git a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py index aa1c6930b..37313ac2a 100644 --- a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py +++ b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py @@ -485,6 +485,7 @@ def test_receipt_records_plan_source_outputs_and_cost(tmp_path: Path) -> None: "cpu": 4.0, "memory_mib": 131072, "timeout_s": 28800, + "nonpreemptible": False, } paths = [item["path"] for item in receipt["outputs"]] assert paths == ["checkpoints/run_identity.json", "populace_us_2024_acs_local.h5"] @@ -866,3 +867,78 @@ def test_branch_verdict_needs_a_proven_ancestry() -> None: {"fetch_returncode": 0, "tip": None, "ancestor_returncode": None}, ): assert verdict("b", COMMIT, **kwargs)["branch_verified"] is False + + +def test_nonpreemptible_placement_is_opt_in_and_priced_at_three_times_list() -> None: + plain = plan_lib.parse_plan(_plan_data(max_wall_seconds=3600)) + assert plain.nonpreemptible is False + assert plain.price_multiplier == 1.0 + data = _plan_data(max_wall_seconds=3600, nonpreemptible=True) + plan = plan_lib.parse_plan(data) + assert plan.nonpreemptible is True + summary = plan_lib.summarize(plan) + assert summary["resources"]["nonpreemptible"] is True + # The heavy class for the #974 materialize wall: $1.71 at list, x3. + assert summary["estimated_usd_at_measured_wall"] == pytest.approx(5.12, abs=0.01) + # The budget plus the runner's 15 minutes of staging and mirroring. + assert plan_lib.summarize(plain)["estimated_usd_at_max_wall"] == pytest.approx( + 1.51, abs=0.01 + ) + assert summary["estimated_usd_at_max_wall"] == pytest.approx(4.54, abs=0.01) + receipt = plan_lib.build_receipt( + plan, + data, + argv=plan_lib.planned_argv(plan), + returncode=0, + started_at="2026-09-23T07:00:00Z", + finished_at="2026-09-23T08:00:00Z", + wall_seconds=3_600.0, + peak_rss_bytes=None, + inputs_verified=[], + outputs=[], + git={"head": COMMIT, "tree_clean": True}, + runner={}, + container_wall_seconds=3_600.0, + ) + assert receipt["resources"]["nonpreemptible"] is True + assert receipt["estimated_usd_at_list_price"] == pytest.approx(3.63, abs=0.01) + with pytest.raises(plan_lib.PlanError, match="true or false"): + plan_lib.parse_plan(_plan_data(nonpreemptible="yes")) + + +# --------------------------------------------------------------------------- # +# The tool's environment # +# --------------------------------------------------------------------------- # + + +def test_tool_environment_strips_credentials_and_stays_offline() -> None: + base = { + "PATH": "/usr/bin", + "HOME": "/root", + "HF_TOKEN": "hf_secret", + "HUGGING_FACE_HUB_TOKEN": "hf_secret", + "HF_HUB_OFFLINE": "0", + "AWS_SECRET_ACCESS_KEY": "x", + "MODAL_IDENTITY_TOKEN": "x", + "OMP_NUM_THREADS": "8", + } + plan_env = {"OMP_NUM_THREADS": "4", "MICROCOSM_ACS_POOL_PEAK_LIMIT_BYTES": "1"} + env, removed = plan_lib.tool_environment(base, plan_env) + assert removed == [ + "AWS_SECRET_ACCESS_KEY", + "HF_TOKEN", + "HUGGING_FACE_HUB_TOKEN", + "MODAL_IDENTITY_TOKEN", + ] + assert not set(removed) & set(env) + assert "hf_secret" not in env.values() + assert env["HF_HUB_OFFLINE"] == "1" + assert env["OMP_NUM_THREADS"] == "4" # the plan's override wins + assert env["PATH"] == "/usr/bin" + assert env["MICROCOSM_ACS_POOL_PEAK_LIMIT_BYTES"] == "1" + + +@pytest.mark.parametrize("key", ["HF_HUB_OFFLINE", "HF_TOKEN", "POPULACE_LEDGER_KEY"]) +def test_tool_environment_refuses_an_unvalidated_plan_env(key: str) -> None: + with pytest.raises(plan_lib.PlanError, match="may not be passed"): + plan_lib.tool_environment({}, {key: "x"}) diff --git a/tools/modal_us_stage.py b/tools/modal_us_stage.py index 0fac082ea..dbe1a6f80 100644 --- a/tools/modal_us_stage.py +++ b/tools/modal_us_stage.py @@ -317,7 +317,9 @@ def _check_input(ref: plan_lib.InputRef, work: Path) -> dict[str, object]: """ -def _parse_check(plan: plan_lib.Plan, argv: list[str]) -> dict[str, object]: +def _parse_check( + plan: plan_lib.Plan, argv: list[str], env: dict[str, str] +) -> dict[str, object]: """Run the pinned tool's own argument parser on the built argv.""" if plan.tool.script is None: @@ -329,7 +331,12 @@ def _parse_check(plan: plan_lib.Plan, argv: list[str]) -> dict[str, object]: cmd = [python, "-B", "-c", _PARSE_SNIPPET, script, plan.tool.parse_function] cmd += tool_argv proc = subprocess.run( - cmd, cwd=plan_lib.IMAGE_REPO_ROOT, text=True, capture_output=True, timeout=900 + cmd, + cwd=plan_lib.IMAGE_REPO_ROOT, + env=env, + text=True, + capture_output=True, + timeout=900, ) return { "returncode": proc.returncode, @@ -456,6 +463,9 @@ def check_stage(plan_data: dict) -> dict: if not git["tool_present"]: problems.append(f"{plan.tool.script} is not in commit {plan.commit}") + # The probe and the parser run in the environment the stage would get. + tool_env, env_removed = plan_lib.tool_environment(os.environ, plan.env) + report["tool_env_removed"] = env_removed env_probe = subprocess.run( [ f"{plan_lib.IMAGE_VENV}/bin/python", @@ -465,6 +475,7 @@ def check_stage(plan_data: dict) -> dict: "('policyengine-us', 'policyengine-core', 'microcosm-build')}))", ], cwd=plan_lib.IMAGE_REPO_ROOT, + env=tool_env, text=True, capture_output=True, timeout=600, @@ -478,7 +489,11 @@ def check_stage(plan_data: dict) -> dict: problems.append("the synced environment does not import microcosm.build") argv = plan_lib.planned_argv(plan) - parse = _parse_check(plan, argv) if git["tool_present"] else {"returncode": None} + parse = ( + _parse_check(plan, argv, tool_env) + if git["tool_present"] + else {"returncode": None} + ) report["tool_argument_parse"] = parse if parse["returncode"] != 0: problems.append("the pinned tool's parser refuses the built argv") @@ -570,7 +585,7 @@ def _run_stage(plan_data: dict) -> dict: started_at, started = _now(), time.time() log_path = state / "logs" / f"{plan.stage}-{started_at.replace(':', '')}.log" log_path.parent.mkdir(parents=True, exist_ok=True) - env = {**os.environ, "HF_HUB_OFFLINE": "1", **plan.env} + env, env_removed = plan_lib.tool_environment(os.environ, plan.env) print(f"$ {' '.join(argv)}", flush=True) with open(log_path, "w") as log: proc = subprocess.Popen( @@ -611,7 +626,12 @@ def _run_stage(plan_data: dict) -> dict: inputs_verified=inputs_verified, outputs=outputs, git=git, - runner={**_runner_identity(), "state_pulled": pulled, "state_pushed": pushed}, + runner={ + **_runner_identity(), + "state_pulled": pulled, + "state_pushed": pushed, + "tool_env_removed": env_removed, + }, prior_receipts=prior, stopped_at_budget=budget.fired, container_wall_seconds=time.time() - container_started, diff --git a/tools/modal_us_stage_plan.py b/tools/modal_us_stage_plan.py index 530b56558..14d5b8d12 100644 --- a/tools/modal_us_stage_plan.py +++ b/tools/modal_us_stage_plan.py @@ -812,6 +812,29 @@ def planned_argv(plan: Plan, work_root: str = WORK_ROOT) -> list[str]: ) +def tool_environment( + base: Mapping[str, str], plan_env: Mapping[str, str] +) -> tuple[dict[str, str], list[str]]: + """The stage tool's environment, and the names removed from ``base``. + + The container's environment without any credential-looking variable + (``HF_TOKEN`` from an attached Hub secret, Modal's own tokens), with + ``HF_HUB_OFFLINE=1`` so the tool cannot reach the Hub, and with the + plan's allowlisted overrides last. Only names are returned, for the + receipt; never values. + """ + + for key in plan_env: + if not _ENV_KEY.fullmatch(key) or is_credential_env_key(key): + raise PlanError(f"env {key!r} may not be passed to the tool") + removed = sorted(key for key in base if is_credential_env_key(key)) + dropped = set(removed) + env = {key: value for key, value in base.items() if key not in dropped} + env["HF_HUB_OFFLINE"] = "1" + env.update(plan_env) + return env, removed + + # --------------------------------------------------------------------------- # # Hashing, mirroring and receipts # # --------------------------------------------------------------------------- # From 9fa2b158996685d591b427265aaa65c81fdb4609 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 23 Sep 2026 05:06:57 -0400 Subject: [PATCH 13/19] Record the Modal acceptance attempt of the 23 September materialize The state-surface materialize at 767312d6 ran on Modal next to the local run from the same inputs. The check passed, and the Modal run did not finish. Modal preempted it at 54 minutes and again at 2 hours 4 minutes, restarting from zero each time, and it was stopped by hand at about $3.62 (workspace billing report). Per chunk it ran two to five times slower than local (149 to 355 s against 66 to 80 s), CPU-bound on one core under gVisor, and held 15 to 24 GB more RSS. The runbook records these numbers and the local run identity and checkpoint digests for the next comparison. It also replaces the unmeasured 94 GB estimate with the 77.9 GB local peak. It documents the opt-in "nonpreemptible" plan field, priced at three times list. That code (NONPREEMPTIBLE_PRICE_MULTIPLIER, Plan.nonpreemptible, run_stage_heavy_nonpreemptible and run_stage_light_nonpreemptible, and their tests) landed in e70b23114. A concurrent commit in the same worktree picked it up from this session's uncommitted edit, so that commit's message does not describe it. The runbook also drops an unobserved claim about out-of-memory restarts. Co-Authored-By: Claude Opus 5.5 --- docs/us-modal-stage-runbook.md | 95 +++++++++++++++++++++++++++++++--- 1 file changed, 87 insertions(+), 8 deletions(-) diff --git a/docs/us-modal-stage-runbook.md b/docs/us-modal-stage-runbook.md index e9090c0a6..84c90fe43 100644 --- a/docs/us-modal-stage-runbook.md +++ b/docs/us-modal-stage-runbook.md @@ -89,9 +89,12 @@ release directory onto the runs volume. Publication stays the human step in The heavy and light classes are sized from the #974 measured peaks (`experiments/us-acs-local-hours-rebuild-20260922/run-resources-and-staging-excerpt.json`, -totals SOI surface). The overnight build of 22 September was expected to -peak near 94 GB for materialize on the state surface. That number was -reported with the task and has not been measured by this runner. +totals SOI surface). On the state SOI surface, the overnight build of 23 +September measured materialize locally at a 77.9 GB peak, 5,569 s of tool +wall and 5,472 CPU-s (4,459 targets; see the acceptance attempt below). On +Modal the same stage ran two to five times slower per chunk and held 15 to +24 GB more RSS at the same point, so the local wall times in this table +understate Modal's wall and cost. | Stage | Measured locally | Class | Request | List price at measured wall | | --- | --- | --- | --- | --- | @@ -240,6 +243,80 @@ MICROCOSM_MODAL_PLAN=docs/us-modal-stage-example-plan.json \ modal run --detach tools/modal_us_stage.py --run ``` +## Acceptance attempt: 23 September materialize on the state surface + +The plan `docs/us-modal-stage-acceptance-20260923-plan.json` ran the +materialize stage of the overnight build of 23 September on Modal +(`run_id` `overnight-20260923-materialize-state`). It pinned build commit +`767312d6` on branch `overnight-acs-local-20260923`, `soi_mode` `state`, +`hh_chunk` 20,000, the local run's peak-limit environment and a +26,400-second wall budget. The same stage ran locally at the same time from +the same inputs. The Modal run did not finish, so its outputs could not be +compared. + +- **Inputs.** The staging H5 (`ed2e6308…`, 10,685,765,051 bytes) went to + the inputs volume in 255 seconds, about 42 MB/s, and the staging summary + (`3aab5e5e…`) went up too. The feed and ladder were already there. +- **Check.** It passed with no problems. The image built from `767312d6`, + the clone was clean on its branch, and the environment synced to + policyengine-us 2.2.1 and policyengine-core 3.32.5, the same versions as + the local environment. The tool's parser accepted the argv, and all four + inputs matched their digests on the volume. +- **Run.** The container started at 06:03:06 UTC and the tool's first + log line came at 06:06:22, so staging and verifying the inputs took under + 3.5 minutes. The tool's own timings, Modal against local: + + | Step | Modal | Local | + | --- | --- | --- | + | Start to staging frame loaded (hashing, specs, load) | 255 to 281 s | 79 s | + | One chunk of 20,000 households | 149 to 355 s | 66 to 80 s | + | peak RSS after chunk 1 | 66.3 GB | 51.0 GB | + | peak RSS after chunk 29 | 74.5 GB | 51.0 GB | + + The container ran gVisor (`Linux-4.19.0-gvisor-x86_64`) with Python + 3.14.2. In the check-class container Modal set `OMP_NUM_THREADS`, + `OPENBLAS_NUM_THREADS`, `MKL_NUM_THREADS` and `BLIS_NUM_THREADS` to that + class's CPU request, 2. In the heavy container `nproc` reported 4. The + local run used macOS arm64 with Python 3.14.4, and its config set none of + those variables. + A sample inside the container, taken with `modal container exec`, showed + the tool process using about 97% of one core (59.4 CPU-s in 61 s). The + stage was CPU-bound on one core, not throttled. +- **Preempted twice.** Modal preempted the first container after 54 + minutes (after chunk 12 of 80) and the second after 2 hours 4 minutes + (after chunk 29). Each time it restarted the function from zero. The run + was launched before the attempt ledger existed, so the budget restarted + too. It was stopped by hand (`modal app stop`) 34 seconds into a third + attempt, because at the observed pace a third attempt needed 6.5 to 7 + more hours and would have taken the total past $10. +- **Cost.** Modal's workspace billing report + (`modal.billing.workspace_billing_report`, hourly) shows $3.57 for the + run through 09:00 UTC. The last 2.5 minutes add about $0.05, for about + $3.62 in total. The check cost $0.02. Billing was at the request, about + $1.21 an hour. +- **Local result for the next comparison.** The local run finished in + 5,569 s with a 77.9 GB peak. `run_identity`: staging `ed2e6308…`, ladder + `39a2ab2a…`, 1,588,854 households, 4,459 targets (3,972 admin specs + declared and compiled, plus 487 state and CD population cells), no + dropped cells, `targets_sha256` + `d843209746bdcf96a831fa90d6fa2fe7aa2969fa32290db35afa9889f7c5d377`. + Checkpoint sha256: `held_back_columns.json` `026c455b…`, + `reviewed_null_fills.json` `727f39c1…`, `targets.json` `d8432097…` + (the same as `targets_sha256`), and `target_frame_lean.h5` + (28,749,924,512 bytes) `ae7d0ae7…`. HDF5 files can differ byte for byte + when their stored values match, so the lean H5 should also be compared + dataset by dataset. The build machine keeps per-dataset digests of the + local file next to the overnight run. + +At the pace observed, one uninterrupted Modal materialize on the state +surface takes about 3.5 to 7 hours. That is about $4 to $9 on preemptible +placement, or $13 to $26 with `"nonpreemptible": true`, against 1.5 hours +locally. The Modal path works, but it does not yet make this stage cheaper +or faster. The run did show that the stage runs as one long single-core +loop over 80 chunks of households. If the chunks are independent, which +this runbook has not checked in `materialize_chunked`, fanning them out +across containers would let preemption lose one chunk instead of the run. + ## Data placement Inputs go to the `policyengine` Modal workspace. Upload only files that are @@ -278,11 +355,13 @@ hold digests, sizes and paths, never file contents. of the same plan that never wrote a receipt to `max_wall_seconds`, and refuses to start when less than a minute is left. The check reports those attempts and the time left. To launch again past a spent budget, raise - `max_wall_seconds` (a new plan digest) or use a new `run_id`. Setting - `nonpreemptible=True` would avoid restarts at three times the list price - for CPU and memory (modal.com/docs/guide/preemption, read 23 September - 2026); the runner does not set it. The heavy class sets a memory request - but no hard limit, and a container killed for memory is not restarted. + `max_wall_seconds` (a new plan digest) or use a new `run_id`. A plan can + set `"nonpreemptible": true` to run on Modal's non-preemptible placement, + which avoids restarts at three times the list price for CPU and memory + (modal.com/docs/guide/preemption, read 23 September 2026); `validate` and + the receipt price it that way. The heavy class sets a memory request but + no hard limit. What Modal does with a container killed for memory has not + been observed here. - **Certification.** A receipt proves which bytes a stage produced. It does not certify a release. Preflight (`tools/preflight_us_release_gates.py`) and certification still run on the output as before. From c3ebde13b723a290b0cc7dfd000a710863a7c548 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 23 Sep 2026 05:09:33 -0400 Subject: [PATCH 14/19] Charge hashing and mirroring to a Modal attempt and lock a run to one attempt The heartbeat stopped before the state was hashed and mirrored, "so that no commit lands mid-mirror"; Modal commits volumes in the background every few seconds anyway, so that bought nothing, and a preemption during hashing or mirroring (many minutes for a large state) went uncharged. The comment also bounded the undercount at one heartbeat, which left out the cold start before _run_stage. The heartbeat now runs until the attempt's final record, and the comment lists what is actually uncounted: cold start and image load, up to one interval after the last record (more after a failed write), and the preemption grace period. Each attempt now ends once with an outcome: "receipt" (finished), "refused" by the lock or the budget (finished, not charged), or "error" for any exception after its first record (unfinished, so charged, like a preemption). A lock serializes the writes, commits and reloads the heartbeat and the main thread make, and a late heartbeat can no longer overwrite the final record. Records carry Modal's input and call ids. The ledger is also the run's lock. An attempt refuses to start while an earlier attempt of the same run, any stage or plan, is still writing its record. A recent record is ambiguous (Modal restarts a preempted input within moments), so the attempt waits two heartbeats and reads again: a record that moved is running, one that did not is dead. A later attempt is left to refuse itself. Volumes have no atomic lock, so this is best effort. No override: a dead attempt stops blocking after one wait, and a running one must not be raced. The check reports recent records without waiting. Receipts are written under a temporary name and renamed, and record the attempt id. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_us_modal_stage_plan_tool.py | 211 ++++++++++++++++++ tools/modal_us_stage.py | 199 +++++++++++++---- tools/modal_us_stage_plan.py | 130 ++++++++++- 3 files changed, 495 insertions(+), 45 deletions(-) diff --git a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py index 37313ac2a..a0bef846f 100644 --- a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py +++ b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py @@ -942,3 +942,214 @@ def test_tool_environment_strips_credentials_and_stays_offline() -> None: def test_tool_environment_refuses_an_unvalidated_plan_env(key: str) -> None: with pytest.raises(plan_lib.PlanError, match="may not be passed"): plan_lib.tool_environment({}, {key: "x"}) + + +# --------------------------------------------------------------------------- # +# Attempt ledger: outcomes, charging and the run's lock # +# --------------------------------------------------------------------------- # + + +def _attempt( + attempt_id: str, + *, + seen: float, + started: float = 0.0, + plan_data: dict | None = None, + **kw, +) -> dict: + data = plan_data or _plan_data(max_wall_seconds=20_000) + return plan_lib.attempt_record( + plan_lib.parse_plan(data), + plan_lib.plan_digest(data), + attempt_id=attempt_id, + started_epoch=started, + last_seen_epoch=seen, + **kw, + ) + + +def test_attempt_records_carry_outcome_note_and_modal_ids() -> None: + record = _attempt( + "20260923T050000000000Z", + seen=60.0, + outcome=plan_lib.OUTCOME_ERROR, + note="PlanError: input 'feed' ...", + modal={"input_id": "in-1", "function_call_id": "fc-1"}, + ) + assert record["outcome"] == "error" + assert record["finished"] is False + assert record["modal"] == {"input_id": "in-1", "function_call_id": "fc-1"} + assert _attempt("a", seen=1.0)["outcome"] is None + + +def test_charging_counts_every_attempt_that_did_not_finish() -> None: + data = _plan_data(max_wall_seconds=20_000) + plan, sha = plan_lib.parse_plan(data), plan_lib.plan_digest(data) + preempted = _attempt("a1", seen=3_000.0) + errored = _attempt("a2", seen=500.0, outcome=plan_lib.OUTCOME_ERROR) + refused = _attempt("a3", seen=270.0, finished=True, outcome="refused") + stopped = _attempt( + "a4", seen=9_000.0, finished=True, outcome="receipt", receipt="r.json" + ) + own = _attempt("a5", seen=10.0) + charged = plan_lib.unfinished_attempts( + [preempted, errored, refused, stopped, own], plan, sha, exclude="a5" + ) + assert [item["attempt_id"] for item in charged] == ["a1", "a2"] + # A budget stop wrote a FAILED receipt, so the same plan starts afresh. + assert plan_lib.remaining_wall_seconds(plan, charged) == 20_000 - 3_500 + + +def test_recent_unfinished_attempts_are_the_run_s_possible_lock_holders() -> None: + now = 10_000.0 + window = plan_lib.ATTEMPT_LIVE_WINDOW_SECONDS + other_run = _attempt( + "b", seen=now, plan_data=_plan_data(run_id="another-run", stage="calibrate") + ) + records = [ + _attempt("fresh", seen=now - 10), + _attempt("edge", seen=now - window + 1), + _attempt("stale", seen=now - window - 1), + _attempt("ended", seen=now - 10, outcome=plan_lib.OUTCOME_ERROR), + _attempt("done", seen=now - 10, finished=True, outcome="receipt"), + _attempt("mine", seen=now), + other_run, + # A record from the runner before outcomes existed. + { + key: value + for key, value in _attempt("legacy", seen=now - 5).items() + if key not in {"outcome", "note", "modal"} + }, + ] + recent = plan_lib.recent_unfinished_attempts( + records, "acs-local-20260923", now=now, exclude="mine" + ) + assert [item["attempt_id"] for item in recent] == ["fresh", "edge", "legacy"] + + +class _Ledger: + """Attempt records as successive reads of the runs volume return them.""" + + def __init__(self, *reads: list[dict]) -> None: + self.reads = list(reads) + self.sleeps: list[float] = [] + + def read(self) -> list[dict]: + return self.reads.pop(0) if len(self.reads) > 1 else self.reads[0] + + def sleep(self, seconds: float) -> None: + self.sleeps.append(seconds) + + +OWN = "20260923T060000000000Z" +EARLIER = "20260923T050000000000Z" +LATER = "20260923T070000000000Z" + + +def _lock(ledger: _Ledger, now: float = 10_000.0) -> list[dict]: + return plan_lib.live_attempts( + ledger.read, + "acs-local-20260923", + own_attempt_id=OWN, + now=lambda: now, + sleep=ledger.sleep, + ) + + +def test_lock_is_free_without_recent_earlier_attempts() -> None: + ledger = _Ledger([_attempt(EARLIER, seen=1_000.0), _attempt(OWN, seen=10_000.0)]) + assert _lock(ledger) == [] + assert ledger.sleeps == [] + + +def test_lock_waits_out_a_preempted_predecessor_and_starts() -> None: + # Modal restarted this input moments after preempting it: the earlier + # attempt's record is recent but never moves again. + dead = _attempt(EARLIER, seen=9_950.0) + ledger = _Ledger([dead], [dead]) + assert _lock(ledger) == [] + assert ledger.sleeps == [plan_lib.ATTEMPT_RECHECK_SECONDS] + + +def test_lock_refuses_while_an_earlier_attempt_keeps_writing() -> None: + ledger = _Ledger( + [_attempt(EARLIER, seen=9_950.0)], [_attempt(EARLIER, seen=10_200.0)] + ) + running = _lock(ledger) + assert [item["attempt_id"] for item in running] == [EARLIER] + + +def test_lock_frees_when_the_earlier_attempt_finishes_during_the_wait() -> None: + ledger = _Ledger( + [_attempt(EARLIER, seen=9_950.0)], + [_attempt(EARLIER, seen=10_200.0, finished=True, outcome="receipt")], + ) + assert _lock(ledger) == [] + + +def test_lock_leaves_a_later_attempt_to_refuse_itself() -> None: + ledger = _Ledger([_attempt(LATER, seen=10_000.0)]) + assert _lock(ledger) == [] + assert ledger.sleeps == [] + + +def test_lock_covers_every_stage_and_plan_of_the_run() -> None: + calibrate = _attempt( + EARLIER, seen=9_950.0, plan_data=_plan_data("calibrate", max_wall_seconds=600) + ) + moved = {**calibrate, "last_seen_epoch": 10_100.0} + assert len(_lock(_Ledger([calibrate], [moved]))) == 1 + + +def test_app_attempt_record_cannot_be_overwritten_after_it_ends( + app, tmp_path: Path, monkeypatch +) -> None: + monkeypatch.setattr(plan_lib, "RUNS_MOUNT", str(tmp_path)) + data = _plan_data(max_wall_seconds=600) + plan = plan_lib.parse_plan(data) + attempt = app._Attempt(plan, plan_lib.plan_digest(data), started=1_000.0) + attempt.write() + record = json.loads(attempt.path.read_text()) + assert (record["finished"], record["outcome"]) == (False, None) + assert record["modal"] == {"input_id": "in-test", "function_call_id": "fc-test"} + attempt.end("receipt", finished=True, receipt="receipts/x.json") + attempt.write() # a heartbeat that lost the race + record = json.loads(attempt.path.read_text()) + assert (record["finished"], record["outcome"]) == (True, "receipt") + assert record["receipt"] == "receipts/x.json" + assert not list(attempt.path.parent.glob("*.tmp")) + assert app.runs_volume.commit.call_count == 2 + assert [item["attempt_id"] for item in attempt.records()] == [attempt.attempt_id] + app.runs_volume.reload.assert_called() + + +@pytest.mark.parametrize( + ("raised", "outcome", "finished"), + [("refusal", "refused", True), ("error", "error", False)], +) +def test_app_run_stage_ends_the_attempt_with_its_outcome( + app, tmp_path: Path, monkeypatch, raised: str, outcome: str, finished: bool +) -> None: + monkeypatch.setattr(plan_lib, "RUNS_MOUNT", str(tmp_path)) + monkeypatch.setattr( + app, + "_git_state", + lambda plan: { + "head_matches_plan": True, + "tree_clean": True, + "branch_verified": True, + "tool_present": True, + }, + ) + error = app._Refusal("lock held") if raised == "refusal" else OSError("disk full") + + def attempt_stage(*_args, **_kwargs): + raise error + + monkeypatch.setattr(app, "_attempt_stage", attempt_stage) + with pytest.raises(type(error)): + app._run_stage(_plan_data(max_wall_seconds=600)) + (path,) = (tmp_path / "runs" / "acs-local-20260923" / "attempts").glob("*.json") + record = json.loads(path.read_text()) + assert (record["outcome"], record["finished"]) == (outcome, finished) + assert str(error) in record["note"] diff --git a/tools/modal_us_stage.py b/tools/modal_us_stage.py index dbe1a6f80..e64064a79 100644 --- a/tools/modal_us_stage.py +++ b/tools/modal_us_stage.py @@ -367,13 +367,31 @@ def run(self) -> None: self.proc.kill() +def _modal_ids() -> dict[str, str | None]: + """This container's Modal input and call ids, recorded for the audit.""" + + ids: dict[str, str | None] = {} + for key, name in ( + ("input_id", "current_input_id"), + ("function_call_id", "current_function_call_id"), + ): + try: + value = getattr(modal, name)() + except Exception: # noqa: BLE001 - absent outside a container + value = None + ids[key] = value if isinstance(value, str) else None + return ids + + class _Attempt(threading.Thread): - """This container's entry in ``runs//attempts/``. + """This container's record in ``runs//attempts/``. Modal restarts a preempted function from scratch on the same input, and - ``retries=0`` does not stop that. The entry is rewritten and committed - every ``ATTEMPT_HEARTBEAT_SECONDS`` so that a restart can charge the - attempts preemption cut short to the plan's ``max_wall_seconds``. + ``retries=0`` does not stop that. The record is written when the attempt + starts, rewritten and committed every ``ATTEMPT_HEARTBEAT_SECONDS``, and + ended once with an outcome (``end``). The records charge the attempts + preemption cut short to the plan's ``max_wall_seconds`` and serve as the + run's lock (``plan_lib.live_attempts``). """ def __init__(self, plan: plan_lib.Plan, plan_sha256: str, started: float) -> None: @@ -385,24 +403,60 @@ def __init__(self, plan: plan_lib.Plan, plan_sha256: str, started: float) -> Non "%Y%m%dT%H%M%S%fZ" ) self.path = _run_dir(plan) / "attempts" / f"{plan.stage}-{self.attempt_id}.json" + self.modal = _modal_ids() self.stop = threading.Event() + # Serializes writes, commits and reloads of the runs volume between the + # heartbeat and the main thread; a reload fails while a file is open. + self.lock = threading.Lock() + self.ended = False - def write(self, *, finished: bool = False, receipt: str | None = None) -> None: + def _write(self, **final: object) -> None: record = plan_lib.attempt_record( self.plan, self.plan_sha256, attempt_id=self.attempt_id, started_epoch=self.started, last_seen_epoch=time.time(), - finished=finished, - receipt=receipt, + modal=self.modal, + **final, # type: ignore[arg-type] ) self.path.parent.mkdir(parents=True, exist_ok=True) - tmp = self.path.with_suffix(".tmp") + tmp = self.path.with_name(self.path.name + ".tmp") tmp.write_text(json.dumps(record, indent=2) + "\n") os.replace(tmp, self.path) runs_volume.commit() + def write(self) -> None: + """Write or refresh the running record; a no-op once ended.""" + + with self.lock: + if not self.ended: + self._write() + + def end( + self, + outcome: str, + *, + finished: bool, + receipt: str | None = None, + note: str | None = None, + ) -> None: + """Write the final record once; no heartbeat can overwrite it after.""" + + self.stop.set() + with self.lock: + if self.ended: + return + self.ended = True + self._write(finished=finished, receipt=receipt, outcome=outcome, note=note) + + def records(self) -> list[dict]: + """Every attempt record of the run, as committed now.""" + + with self.lock: + runs_volume.reload() + return _attempt_records(self.plan) + def run(self) -> None: while not self.stop.wait(plan_lib.ATTEMPT_HEARTBEAT_SECONDS): try: @@ -411,16 +465,12 @@ def run(self) -> None: print(f"attempt heartbeat failed: {error}", flush=True) -def _prior_attempts(plan: plan_lib.Plan, plan_sha256: str) -> list[dict[str, object]]: +def _attempt_records(plan: plan_lib.Plan) -> list[dict]: attempts_dir = _run_dir(plan) / "attempts" - records = ( - [_load_json(path) for path in sorted(attempts_dir.glob("*.json"))] - if attempts_dir.exists() - else [] - ) - return plan_lib.unfinished_attempts( - [record for record in records if record], plan, plan_sha256 - ) + if not attempts_dir.exists(): + return [] + records = [_load_json(path) for path in sorted(attempts_dir.glob("*.json"))] + return [record for record in records if record] def _load_json(path: Path) -> dict | None: @@ -510,7 +560,10 @@ def check_stage(plan_data: dict) -> dict: ) report["prior_run_identity"] = identity problems += plan_lib.prior_state_problems(plan, identity) - prior_attempts = _prior_attempts(plan, plan_lib.plan_digest(plan_data)) + records = _attempt_records(plan) + prior_attempts = plan_lib.unfinished_attempts( + records, plan, plan_lib.plan_digest(plan_data) + ) report["prior_unfinished_attempts"] = prior_attempts remaining = plan_lib.remaining_wall_seconds(plan, prior_attempts) report["remaining_wall_seconds"] = remaining @@ -518,12 +571,26 @@ def check_stage(plan_data: dict) -> dict: problems.append( "earlier unfinished attempts used this plan's max_wall_seconds budget" ) + # The run's lock, without the wait a run would do. + now = time.time() + running = plan_lib.recent_unfinished_attempts(records, plan.run_id, now=now) + report["recent_unfinished_attempts"] = running + problems += [ + f"attempt {item.get('attempt_id')} ({item.get('stage')}) of this run wrote " + f"its record {now - float(item.get('last_seen_epoch') or 0.0):.0f}s ago and " + "may still be running; a run waits to see and refuses if it is" + for item in running + ] report["problems"] = problems report["ok"] = not problems return report +class _Refusal(plan_lib.PlanError): + """The lock or the budget stopped an attempt before it staged anything.""" + + def _run_stage(plan_data: dict) -> dict: container_started = time.time() plan = plan_lib.parse_plan(plan_data) @@ -539,27 +606,82 @@ def _run_stage(plan_data: dict) -> dict: ): raise plan_lib.PlanError(f"image clone does not match the plan: {git}") - work = Path(plan_lib.WORK_ROOT) - state = work / "state" run_dir = _run_dir(plan) runs_volume.reload() - # 1. Refuse a foreign or missing predecessor before paying for inputs. + # 1. Refuse a foreign or missing predecessor before recording an attempt. identity = _load_json(run_dir / "state" / "checkpoints" / "run_identity.json") problems = plan_lib.prior_state_problems(plan, identity) if problems: raise plan_lib.PlanError("; ".join(problems)) - # 1b. Charge attempts that preemption cut short to this plan's budget, - # then record this one. An attempt that never wrote a receipt stays - # charged; to relaunch past it, raise max_wall_seconds or change run_id. + # 2. Record this attempt. The heartbeat runs until the final record, so the + # lock wait, staging, the tool, hashing and mirroring are all charged + # if preemption cuts the attempt short. Any exception from here on + # leaves the attempt unfinished (charged) with outcome "error"; only a + # receipt or a refusal by the lock or the budget finishes it. plan_sha256 = plan_lib.plan_digest(plan_data) - prior_attempts = _prior_attempts(plan, plan_sha256) + attempt = _Attempt(plan, plan_sha256, container_started) + attempt.write() + attempt.start() + try: + return _attempt_stage(plan, plan_data, plan_sha256, attempt, git) + except _Refusal as refusal: + attempt.end(plan_lib.OUTCOME_REFUSED, finished=True, note=str(refusal)[:500]) + raise + except BaseException as error: + try: + attempt.end( + plan_lib.OUTCOME_ERROR, + finished=False, + note=f"{type(error).__name__}: {error}"[:500], + ) + except Exception as record_error: # noqa: BLE001 - keep the first error + print(f"attempt record failed: {record_error}", flush=True) + raise + + +def _attempt_stage( + plan: plan_lib.Plan, + plan_data: dict, + plan_sha256: str, + attempt: _Attempt, + git: dict[str, object], +) -> dict: + work = Path(plan_lib.WORK_ROOT) + state = work / "state" + run_dir = _run_dir(plan) + + # 3. The run's lock: an earlier attempt of this run that is still + # writing its record would race on the state directory. + running = plan_lib.live_attempts( + attempt.records, + plan.run_id, + own_attempt_id=attempt.attempt_id, + now=time.time, + sleep=time.sleep, + log=lambda message: print(message, flush=True), + ) + if running: + names = ", ".join( + f"{item.get('attempt_id')} ({item.get('stage')})" for item in running + ) + raise _Refusal( + f"run {plan.run_id!r} has a running attempt: {names}. Run one stage " + "of a run at a time; wait for it or stop its app" + ) + + # 4. Charge the attempts of this plan that never finished to its budget. + prior_attempts = plan_lib.unfinished_attempts( + attempt.records(), plan, plan_sha256, exclude=attempt.attempt_id + ) budget_seconds = plan_lib.remaining_wall_seconds(plan, prior_attempts) if budget_seconds is not None and budget_seconds < plan_lib.MIN_ATTEMPT_SECONDS: - raise plan_lib.PlanError( + raise _Refusal( f"{len(prior_attempts)} earlier unfinished attempt(s) of this plan used " - f"the max_wall_seconds budget ({plan.max_wall_seconds}s); not starting" + f"the max_wall_seconds budget ({plan.max_wall_seconds}s); not starting. " + "To launch again, raise max_wall_seconds (a new plan digest) or use a " + "new run_id" ) if prior_attempts: print( @@ -567,11 +689,8 @@ def _run_stage(plan_data: dict) -> dict: f"tool budget for this attempt {budget_seconds}s", flush=True, ) - attempt = _Attempt(plan, plan_sha256, container_started) - attempt.write() - attempt.start() - # 2. Inputs to stable local paths, each digest verified; then this run's + # 5. Inputs to stable local paths, each digest verified; then this run's # prior state (checkpoints, calibrated H5) onto local disk. inputs_verified = [_stage_input(ref, work) for ref in plan.inputs.values()] pulled = plan_lib.mirror_tree(run_dir / "state", state) @@ -580,7 +699,7 @@ def _run_stage(plan_data: dict) -> dict: for path in sorted(receipts_dir.glob("*.json")) if receipts_dir.exists() else []: prior.append({"file": path.name, "sha256": plan_lib.sha256_file(path)[0]}) - # 3. The stage itself, in the pinned tree, logged to the state directory. + # 6. The stage itself, in the pinned tree, logged to the state directory. argv = plan_lib.planned_argv(plan) started_at, started = _now(), time.time() log_path = state / "logs" / f"{plan.stage}-{started_at.replace(':', '')}.log" @@ -608,10 +727,8 @@ def _run_stage(plan_data: dict) -> dict: peak_kib = resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss finished_at = _now() - # 4. Outputs: hash the whole state tree, mirror it to the volume, receipt. - # The heartbeat stops first so that no commit lands mid-mirror. - attempt.stop.set() - attempt.join(timeout=60) + # 7. Outputs: hash the whole state tree, mirror it to the volume, receipt. + # The heartbeat keeps running, so this is charged too. outputs = plan_lib.hash_tree(state) pushed = plan_lib.mirror_tree(state, run_dir / "state") receipt = plan_lib.build_receipt( @@ -631,17 +748,21 @@ def _run_stage(plan_data: dict) -> dict: "state_pulled": pulled, "state_pushed": pushed, "tool_env_removed": env_removed, + "modal": attempt.modal, }, prior_receipts=prior, stopped_at_budget=budget.fired, - container_wall_seconds=time.time() - container_started, + container_wall_seconds=time.time() - attempt.started, prior_attempts=prior_attempts, budget_seconds=budget_seconds, + attempt_id=attempt.attempt_id, ) receipts_dir.mkdir(parents=True, exist_ok=True) name = f"{plan.stage}-{started_at.replace(':', '')}.json" - (receipts_dir / name).write_text(json.dumps(receipt, indent=2) + "\n") - attempt.write(finished=True, receipt=f"receipts/{name}") + tmp = receipts_dir / f".{name}.tmp" + tmp.write_text(json.dumps(receipt, indent=2) + "\n") + os.replace(tmp, receipts_dir / name) + attempt.end(plan_lib.OUTCOME_RECEIPT, finished=True, receipt=f"receipts/{name}") receipt["receipt_path"] = f"runs/{plan.run_id}/receipts/{name}" shutil.rmtree(work / "hf", ignore_errors=True) return receipt diff --git a/tools/modal_us_stage_plan.py b/tools/modal_us_stage_plan.py index 14d5b8d12..0ace51498 100644 --- a/tools/modal_us_stage_plan.py +++ b/tools/modal_us_stage_plan.py @@ -951,6 +951,7 @@ def build_receipt( container_wall_seconds: float | None = None, prior_attempts: Sequence[Mapping[str, object]] = (), budget_seconds: int | None = None, + attempt_id: str | None = None, ) -> dict[str, object]: resources = plan.resources # The tool's wall leaves out staging the inputs, hashing the state tree @@ -996,6 +997,7 @@ def build_receipt( }, "argv": list(argv), "returncode": returncode, + "attempt_id": attempt_id, "max_wall_seconds": plan.max_wall_seconds, "budget_seconds_this_attempt": budget_seconds, "stopped_at_budget": stopped_at_budget, @@ -1078,13 +1080,28 @@ def prior_state_problems(plan: Plan, run_identity: Mapping | None) -> list[str]: # --------------------------------------------------------------------------- # ATTEMPT_SCHEMA = "microcosm-modal-us-stage-attempt/1" -# How often a running attempt records that it is still alive; a preempted -# attempt is charged up to its last heartbeat, so it can be undercounted by -# at most this much. +# How often a running attempt rewrites its record. The record runs from the +# start of ``_run_stage`` until the attempt's final record, so the lock wait, +# staging, the tool, hashing and mirroring are all inside it. A preempted +# attempt is charged from its start to its last record, which undercounts +# what Modal bills for it by: the container's cold start and image load +# before ``_run_stage`` begins; up to one interval after the last record +# (longer if a heartbeat write failed; failures are logged, not retried); +# and the preemption grace period. ATTEMPT_HEARTBEAT_SECONDS = 120 +# An unfinished record newer than this may belong to a running attempt. +ATTEMPT_LIVE_WINDOW_SECONDS = 2 * ATTEMPT_HEARTBEAT_SECONDS + 60 +# How long a new attempt waits to see whether such a record moves: two +# heartbeats and slack, so a single failed write does not read as a death. +ATTEMPT_RECHECK_SECONDS = 2 * ATTEMPT_HEARTBEAT_SECONDS + 30 # An attempt left with less tool time than this refuses to start. MIN_ATTEMPT_SECONDS = 60 +# How an attempt ended. A running (or preempted) attempt has no outcome. +OUTCOME_RECEIPT = "receipt" # the tool ran and a receipt was written +OUTCOME_REFUSED = "refused" # the lock or the budget stopped it before staging +OUTCOME_ERROR = "error" # an exception after its first record + def attempt_record( plan: Plan, @@ -1095,6 +1112,9 @@ def attempt_record( last_seen_epoch: float, finished: bool = False, receipt: str | None = None, + outcome: str | None = None, + note: str | None = None, + modal: Mapping[str, object] | None = None, ) -> dict[str, object]: return { "schema": ATTEMPT_SCHEMA, @@ -1107,16 +1127,29 @@ def attempt_record( "elapsed_seconds": round(max(0.0, last_seen_epoch - started_epoch), 1), "finished": finished, "receipt": receipt, + "outcome": outcome, + "note": note, + "modal": dict(modal or {}), } def unfinished_attempts( - records: Iterable[Mapping], plan: Plan, plan_sha256: str + records: Iterable[Mapping], + plan: Plan, + plan_sha256: str, + *, + exclude: str | None = None, ) -> list[dict[str, object]]: - """Earlier attempts of this exact plan and stage that never wrote a receipt. + """Attempts of this exact plan and stage that are charged to its budget. Modal restarts a preempted function on the same input, from scratch and - regardless of ``retries``; each such attempt is billed. + regardless of ``retries``; each such attempt is billed. Every attempt + that is not ``finished`` is charged: one preemption cut short, and one + that raised after its first record (a digest mismatch, a failed pull or + mirror, a runner bug). A finished attempt is not: one that wrote a + receipt, COMPLETED or FAILED (including a stop at the budget, after + which the same plan starts again with its whole budget), and one the + lock or the budget refused. ``exclude`` is the caller's own attempt. """ return [ @@ -1126,10 +1159,95 @@ def unfinished_attempts( and record.get("run_id") == plan.run_id and record.get("stage") == plan.stage and record.get("plan_sha256") == plan_sha256 + and record.get("attempt_id") != exclude + and not record.get("finished") + ] + + +def recent_unfinished_attempts( + records: Iterable[Mapping], + run_id: str, + *, + now: float, + exclude: str | None = None, +) -> list[dict[str, object]]: + """Attempts of the run, any stage or plan, that may still be running. + + Unfinished, without an outcome, and with a record newer than + ATTEMPT_LIVE_WINDOW_SECONDS. A preempted attempt looks like this until + its record ages out, so this alone cannot tell the two apart. + """ + + return [ + dict(record) + for record in records + if record.get("schema") == ATTEMPT_SCHEMA + and record.get("run_id") == run_id + and record.get("attempt_id") != exclude and not record.get("finished") + and record.get("outcome") is None + and now - float(record.get("last_seen_epoch") or 0.0) + < ATTEMPT_LIVE_WINDOW_SECONDS ] +def live_attempts( + read_records: Callable[[], Iterable[Mapping]], + run_id: str, + *, + own_attempt_id: str, + now: Callable[[], float], + sleep: Callable[[float], None], + log: Callable[[str], None] = lambda _message: None, +) -> list[dict[str, object]]: + """The run's lock: earlier attempts that are still running. + + Two attempts of one run would race on its state directory. An earlier + attempt (smaller ``attempt_id``, a start timestamp) whose record is + recent is either running or was preempted moments ago, as when Modal + restarts this very input. Wait ATTEMPT_RECHECK_SECONDS and read again: a + running attempt has rewritten its record, a dead one has not. Returns the + running ones; the caller refuses to start when there are any. A later + attempt is left alone: it sees this one and refuses itself. Modal + volumes have no atomic lock, so two attempts that start within one + commit of each other can both pass; the runbook asks for one stage of a + run at a time. + """ + + before = [ + record + for record in recent_unfinished_attempts( + read_records(), run_id, now=now(), exclude=own_attempt_id + ) + if str(record.get("attempt_id") or "") < own_attempt_id + ] + if not before: + return [] + names = ", ".join(str(record.get("attempt_id")) for record in before) + log( + f"LOCK: earlier attempt(s) {names} of run {run_id!r} wrote a record " + f"recently; waiting {ATTEMPT_RECHECK_SECONDS}s to see whether they run" + ) + sleep(ATTEMPT_RECHECK_SECONDS) + after = { + record.get("attempt_id"): record + for record in read_records() + if record.get("schema") == ATTEMPT_SCHEMA + } + running = [] + for record in before: + latest = after.get(record.get("attempt_id")) + if ( + latest is not None + and not latest.get("finished") + and latest.get("outcome") is None + and float(latest.get("last_seen_epoch") or 0.0) + > float(record.get("last_seen_epoch") or 0.0) + ): + running.append(dict(latest)) + return running + + def remaining_wall_seconds( plan: Plan, prior_unfinished: Sequence[Mapping] ) -> int | None: From 034c9845d061f1542b702557f2680ceb4ed18856 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 23 Sep 2026 05:11:37 -0400 Subject: [PATCH 15/19] Mirror Modal stage state atomically and verify it against the latest receipt mirror_tree copied each file straight over its target with copy2, so a preemption mid-copy could leave a truncated file on the runs volume, and a preemption between files a mix of new and old files. Nothing checked the state a restarted or later stage pulled; the runbook said a cut-short attempt's work "is lost", which was not true of one cut short while mirroring. Each file is now copied to a hidden "..mirror-partial" in its own directory and renamed over the target. Listings skip partials and the next mirror removes them, so no half-written file is ever part of a state tree. After pulling, a stage verifies the state against the run's latest receipt (by finished_at, since receipt names begin with the stage) with strict verify_receipt, and refuses on any difference, before staging inputs; with no receipt the state must be empty. The check hashes the volume's state in place and reports the same problems. The receipt names the receipt the pulled state matched (prior_state_verified_against). Co-Authored-By: Claude Opus 5.5 --- .../tests/test_us_modal_stage_plan_tool.py | 148 +++++++++++++++++- tools/modal_us_stage.py | 43 ++++- tools/modal_us_stage_plan.py | 87 +++++++++- 3 files changed, 265 insertions(+), 13 deletions(-) diff --git a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py index a0bef846f..4f49d701c 100644 --- a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py +++ b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py @@ -443,10 +443,13 @@ def test_mirror_tree_round_trip_is_a_no_op(tmp_path: Path) -> None: (src / "checkpoints").mkdir(parents=True) (src / "checkpoints" / "targets.json").write_text("[]") (src / "artifact.h5").write_bytes(b"h5") - assert plan_lib.mirror_tree(src, dst) == {"copied": 2, "deleted": 0} - assert plan_lib.mirror_tree(src, dst) == {"copied": 0, "deleted": 0} + counts = plan_lib.mirror_tree(src, dst) + assert counts == {"copied": 2, "deleted": 0, "partials_removed": 0} + counts = plan_lib.mirror_tree(src, dst) + assert counts == {"copied": 0, "deleted": 0, "partials_removed": 0} (src / "artifact.h5").unlink() - assert plan_lib.mirror_tree(src, dst) == {"copied": 0, "deleted": 1} + counts = plan_lib.mirror_tree(src, dst) + assert counts == {"copied": 0, "deleted": 1, "partials_removed": 0} assert plan_lib.tree_listing(dst).keys() == {"checkpoints/targets.json"} @@ -1153,3 +1156,142 @@ def attempt_stage(*_args, **_kwargs): record = json.loads(path.read_text()) assert (record["outcome"], record["finished"]) == (outcome, finished) assert str(error) in record["note"] + + +# --------------------------------------------------------------------------- # +# Atomic mirroring and the pulled-state check # +# --------------------------------------------------------------------------- # + + +def test_a_mirror_cut_short_never_leaves_a_half_written_file( + tmp_path: Path, monkeypatch +) -> None: + src, dst = tmp_path / "src", tmp_path / "dst" + src.mkdir() + (src / "a.npz").write_bytes(b"old-a") + (src / "b.h5").write_bytes(b"old-b") + plan_lib.mirror_tree(src, dst) + (src / "a.npz").write_bytes(b"new-a-longer") + (src / "b.h5").write_bytes(b"new-b-longer") + real_copy = plan_lib.shutil.copy2 + + def copy_then_die(source, target): + if Path(source).name == "b.h5": + Path(target).write_bytes(b"new-") # the partial copy, then preemption + raise KeyboardInterrupt + return real_copy(source, target) + + monkeypatch.setattr(plan_lib.shutil, "copy2", copy_then_die) + with pytest.raises(KeyboardInterrupt): + plan_lib.mirror_tree(src, dst) + # a.npz was replaced whole, b.h5 still holds its old bytes, and the + # partial copy is invisible to listings. + assert (dst / "a.npz").read_bytes() == b"new-a-longer" + assert (dst / "b.h5").read_bytes() == b"old-b" + assert (dst / f".b.h5{plan_lib.MIRROR_PARTIAL_SUFFIX}").exists() + assert set(plan_lib.tree_listing(dst)) == {"a.npz", "b.h5"} + monkeypatch.setattr(plan_lib.shutil, "copy2", real_copy) + counts = plan_lib.mirror_tree(src, dst) + assert counts["partials_removed"] == 1 + assert (dst / "b.h5").read_bytes() == b"new-b-longer" + assert not list(dst.rglob(f"*{plan_lib.MIRROR_PARTIAL_SUFFIX}")) + + +def _state_receipt( + state: Path, finished_at: str, stage: str = "materialize", run_id: str = "" +) -> dict: + data = _plan_data(stage, run_id=run_id or "acs-local-20260923") + plan = plan_lib.parse_plan(data) + return plan_lib.build_receipt( + plan, + data, + argv=["python"], + returncode=0, + started_at=finished_at, + finished_at=finished_at, + wall_seconds=1.0, + peak_rss_bytes=None, + inputs_verified=[], + outputs=plan_lib.hash_tree(state), + git={}, + runner={}, + ) + + +def test_latest_receipt_is_by_finish_time_not_by_name(tmp_path: Path) -> None: + state = tmp_path / "state" + state.mkdir() + materialize = _state_receipt(state, "2026-09-23T09:00:00Z") + calibrate = _state_receipt(state, "2026-09-23T10:00:00Z", stage="calibrate") + foreign = _state_receipt(state, "2026-09-23T11:00:00Z", run_id="other-run") + receipts = [ + ("calibrate-2026-09-23T100000Z.json", calibrate), + ("materialize-2026-09-23T090000Z.json", materialize), + ("materialize-2026-09-23T110000Z.json", foreign), + ("truncated.json", {}), + ] + name, receipt = plan_lib.latest_receipt(receipts, "acs-local-20260923") + assert name == "calibrate-2026-09-23T100000Z.json" + assert receipt["stage"] == "calibrate" + assert plan_lib.latest_receipt([("x.json", {})], "acs-local-20260923") is None + + +def test_pulled_state_must_be_what_the_latest_receipt_lists(tmp_path: Path) -> None: + state = tmp_path / "state" + assert plan_lib.pulled_state_problems(state, None) == [] + (state / "checkpoints").mkdir(parents=True) + (state / "checkpoints" / "targets.json").write_text("[]") + # State with no receipt at all: a first stage cut short after mirroring. + (problem,) = plan_lib.pulled_state_problems(state, None) + assert "has no receipt" in problem + receipt = _state_receipt(state, "2026-09-23T09:00:00Z") + latest = ("materialize-2026-09-23T090000Z.json", receipt) + assert plan_lib.pulled_state_problems(state, latest) == [] + # A later stage's mirror cut short: one file replaced, one added. + (state / "checkpoints" / "targets.json").write_text("[1]") + (state / "weights_latest.npz").write_bytes(b"w") + problems = plan_lib.pulled_state_problems(state, latest) + assert problems == [ + "state differs from receipt materialize-2026-09-23T090000Z.json: " + "size mismatch: checkpoints/targets.json is 3 bytes, receipt 2", + "state differs from receipt materialize-2026-09-23T090000Z.json: " + "not in receipt: weights_latest.npz", + ] + + +def test_receipt_names_the_receipt_the_pulled_state_matched(tmp_path: Path) -> None: + data = _plan_data("calibrate") + plan = plan_lib.parse_plan(data) + receipt = plan_lib.build_receipt( + plan, + data, + argv=["python"], + returncode=0, + started_at="t0", + finished_at="t1", + wall_seconds=1.0, + peak_rss_bytes=None, + inputs_verified=[], + outputs=[], + git={}, + runner={}, + attempt_id=OWN, + prior_state_verified_against="materialize-2026-09-23T090000Z.json", + ) + assert receipt["attempt_id"] == OWN + assert receipt["prior_state_verified_against"] == ( + "materialize-2026-09-23T090000Z.json" + ) + + +def test_app_reads_receipts_even_when_one_is_truncated(app, tmp_path: Path) -> None: + receipts = tmp_path / "receipts" + receipts.mkdir() + (receipts / "materialize-2026-09-23T090000Z.json").write_text('{"schema": "x"}') + (receipts / "calibrate-2026-09-23T100000Z.json").write_text('{"sche') + (receipts / ".calibrate-2026-09-23T110000Z.json.tmp").write_text("{}") + rows = app._receipts(tmp_path) + assert [(name, data) for name, data, _ in rows] == [ + ("calibrate-2026-09-23T100000Z.json", {}), + ("materialize-2026-09-23T090000Z.json", {"schema": "x"}), + ] diff --git a/tools/modal_us_stage.py b/tools/modal_us_stage.py index e64064a79..e6f9eab40 100644 --- a/tools/modal_us_stage.py +++ b/tools/modal_us_stage.py @@ -465,6 +465,18 @@ def run(self) -> None: print(f"attempt heartbeat failed: {error}", flush=True) +def _receipts(run_dir: Path) -> list[tuple[str, dict, str]]: + """``(file name, parsed receipt or {}, sha256)`` for each of a run's receipts.""" + + receipts_dir = run_dir / "receipts" + if not receipts_dir.exists(): + return [] + return [ + (path.name, _load_json(path) or {}, plan_lib.sha256_file(path)[0]) + for path in sorted(receipts_dir.glob("*.json")) + ] + + def _attempt_records(plan: plan_lib.Plan) -> list[dict]: attempts_dir = _run_dir(plan) / "attempts" if not attempts_dir.exists(): @@ -560,6 +572,12 @@ def check_stage(plan_data: dict) -> dict: ) report["prior_run_identity"] = identity problems += plan_lib.prior_state_problems(plan, identity) + # The state the run would pull, hashed in place against its latest receipt. + latest = plan_lib.latest_receipt( + [(name, data) for name, data, _ in _receipts(_run_dir(plan))], plan.run_id + ) + report["prior_state_verified_against"] = latest[0] if latest else None + problems += plan_lib.pulled_state_problems(_run_dir(plan) / "state", latest) records = _attempt_records(plan) prior_attempts = plan_lib.unfinished_attempts( records, plan, plan_lib.plan_digest(plan_data) @@ -690,14 +708,23 @@ def _attempt_stage( flush=True, ) - # 5. Inputs to stable local paths, each digest verified; then this run's - # prior state (checkpoints, calibrated H5) onto local disk. - inputs_verified = [_stage_input(ref, work) for ref in plan.inputs.values()] + # 5. This run's prior state (checkpoints, calibrated H5) onto local disk, + # verified against the run's latest receipt before anything uses it; + # then the inputs to stable local paths, each digest verified. pulled = plan_lib.mirror_tree(run_dir / "state", state) - receipts_dir = run_dir / "receipts" - prior = [] - for path in sorted(receipts_dir.glob("*.json")) if receipts_dir.exists() else []: - prior.append({"file": path.name, "sha256": plan_lib.sha256_file(path)[0]}) + receipts = _receipts(run_dir) + latest = plan_lib.latest_receipt( + [(name, data) for name, data, _ in receipts], plan.run_id + ) + problems = plan_lib.pulled_state_problems(state, latest) + if problems: + raise plan_lib.PlanError( + f"run {plan.run_id!r}: its state on {plan_lib.RUNS_VOLUME} is not what " + "its latest receipt lists, so a stage was cut short after changing it; " + "use a new run_id. " + "; ".join(problems[:10]) + ) + prior = [{"file": name, "sha256": sha} for name, _, sha in receipts] + inputs_verified = [_stage_input(ref, work) for ref in plan.inputs.values()] # 6. The stage itself, in the pinned tree, logged to the state directory. argv = plan_lib.planned_argv(plan) @@ -756,7 +783,9 @@ def _attempt_stage( prior_attempts=prior_attempts, budget_seconds=budget_seconds, attempt_id=attempt.attempt_id, + prior_state_verified_against=latest[0] if latest else None, ) + receipts_dir = run_dir / "receipts" receipts_dir.mkdir(parents=True, exist_ok=True) name = f"{plan.stage}-{started_at.replace(':', '')}.json" tmp = receipts_dir / f".{name}.tmp" diff --git a/tools/modal_us_stage_plan.py b/tools/modal_us_stage_plan.py index 0ace51498..606c9a121 100644 --- a/tools/modal_us_stage_plan.py +++ b/tools/modal_us_stage_plan.py @@ -877,14 +877,24 @@ def verify_digest(ref: InputRef, actual_sha256: str) -> None: ) +# A file mid-copy in mirror_tree; never part of a state tree. +MIRROR_PARTIAL_SUFFIX = ".mirror-partial" + + def tree_listing(root: Path | str) -> dict[str, tuple[int, int]]: - """``{relative posix path: (bytes, mtime_ns)}`` for every regular file.""" + """``{relative posix path: (bytes, mtime_ns)}`` for every regular file. + + A ``MIRROR_PARTIAL_SUFFIX`` file (a copy a preemption cut short) is left + out; mirror_tree removes it on the next push. + """ root = Path(root) listing: dict[str, tuple[int, int]] = {} if not root.exists(): return listing for path in sorted(root.rglob("*")): + if path.name.endswith(MIRROR_PARTIAL_SUFFIX): + continue if path.is_file() and not path.is_symlink(): stat = path.stat() listing[path.relative_to(root).as_posix()] = ( @@ -912,15 +922,37 @@ def mirror_actions( def mirror_tree(source: Path | str, destination: Path | str) -> dict[str, int]: + """Make ``destination`` mirror ``source``, one file at a time, atomically. + + Each file is copied to a temporary name in its own directory and renamed + over the target, so no file on the destination is ever half-written. A + preemption between files still leaves a mix of new and old files (and + files not yet deleted); the next attempt's pulled-state check against + the latest receipt refuses that mix. + """ + source, destination = Path(source), Path(destination) + partials = ( + sorted(destination.rglob(f"*{MIRROR_PARTIAL_SUFFIX}")) + if destination.exists() + else [] + ) + for stale in partials: + stale.unlink() copy, delete = mirror_actions(tree_listing(source), tree_listing(destination)) for rel in copy: target = destination / rel target.parent.mkdir(parents=True, exist_ok=True) - shutil.copy2(source / rel, target) + tmp = target.with_name(f".{target.name}{MIRROR_PARTIAL_SUFFIX}") + shutil.copy2(source / rel, tmp) + os.replace(tmp, target) for rel in delete: (destination / rel).unlink() - return {"copied": len(copy), "deleted": len(delete)} + return { + "copied": len(copy), + "deleted": len(delete), + "partials_removed": len(partials), + } def hash_tree(root: Path | str) -> list[dict[str, object]]: @@ -952,6 +984,7 @@ def build_receipt( prior_attempts: Sequence[Mapping[str, object]] = (), budget_seconds: int | None = None, attempt_id: str | None = None, + prior_state_verified_against: str | None = None, ) -> dict[str, object]: resources = plan.resources # The tool's wall leaves out staging the inputs, hashing the state tree @@ -1012,6 +1045,7 @@ def build_receipt( **container, "inputs": [dict(item) for item in inputs_verified], "prior_receipts": [dict(item) for item in prior_receipts], + "prior_state_verified_against": prior_state_verified_against, "outputs": [dict(item) for item in outputs], } @@ -1046,6 +1080,53 @@ def verify_receipt( return problems +def latest_receipt( + receipts: Iterable[tuple[str, Mapping]], run_id: str +) -> tuple[str, Mapping] | None: + """The run's most recent receipt, by ``finished_at`` (then file name). + + Every receipt is written after its state was mirrored, and one stage of + a run runs at a time, so this is the receipt that describes the state + on the runs volume. Receipt names begin with the stage, so sorting them + by name would not do. + """ + + mine = [ + (str(receipt.get("finished_at") or ""), name, receipt) + for name, receipt in receipts + if receipt.get("schema") == RECEIPT_SCHEMA and receipt.get("run_id") == run_id + ] + if not mine: + return None + _, name, receipt = max(mine, key=lambda item: (item[0], item[1])) + return name, receipt + + +def pulled_state_problems( + state_root: Path | str, latest: tuple[str, Mapping] | None +) -> list[str]: + """Refusals for a run's state that its latest receipt does not describe. + + A preemption or error after the state was mirrored and before the + receipt was written, or during a mirror, leaves state no receipt lists. + With no receipt the state must be empty. + """ + + if latest is None: + files = tree_listing(state_root) + if not files: + return [] + return [ + f"the run's state has {len(files)} file(s) and the run has no receipt " + f"(e.g. {sorted(files)[:3]})" + ] + name, receipt = latest + return [ + f"state differs from receipt {name}: {problem}" + for problem in verify_receipt(receipt, state_root, strict=True) + ] + + def prior_state_problems(plan: Plan, run_identity: Mapping | None) -> list[str]: """Refusals for a stage whose predecessor state is missing or foreign. From 5d5c6e96d77eaab42455ea521b113101c9142e35 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 23 Sep 2026 05:15:38 -0400 Subject: [PATCH 16/19] Test the nonpreemptible option's refusal, runners and pricing Review of the opt-in "nonpreemptible" plan field, whose code landed in e70b23114 and whose first test landed in d0bab3201 (both swept in from a concurrent session's uncommitted edit; neither message says so). The 3x multiplier matches modal.com/docs/guide/preemption and modal.com/pricing (both read 23 September), the installed Modal client (1.3.0) accepts nonpreemptible= on app.function, and the arithmetic in that test holds. Nothing needed fixing; three gaps are now tested: - the check class refuses nonpreemptible (it was untested); - every class and placement a valid plan can produce has a runner in RUNNERS whose Modal options carry that class's cpu, memory and timeout, retries=0 and the placement (the stub modal now keeps decorated functions and records their options); - every attempt of a non-preemptible plan, including earlier ones, is priced at 3x, and switching the flag is a new plan digest. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_us_modal_stage_plan_tool.py | 71 +++++++++++++++++++ 1 file changed, 71 insertions(+) diff --git a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py index 4f49d701c..fff86fe5a 100644 --- a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py +++ b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py @@ -50,6 +50,11 @@ def app(monkeypatch): stub = MagicMock(name="modal") stub.is_local.return_value = False + # @app.function(...) keeps the function and records its Modal options. + stub.App.return_value.function.side_effect = lambda **options: ( + lambda function: setattr(function, "modal_options", options) or function + ) + stub.App.return_value.local_entrypoint.side_effect = lambda **_: lambda f: f stub.current_input_id.return_value = "in-test" stub.current_function_call_id.return_value = "fc-test" monkeypatch.setitem(sys.modules, "modal", stub) @@ -1295,3 +1300,69 @@ def test_app_reads_receipts_even_when_one_is_truncated(app, tmp_path: Path) -> N ("calibrate-2026-09-23T100000Z.json", {}), ("materialize-2026-09-23T090000Z.json", {"schema": "x"}), ] + + +# --------------------------------------------------------------------------- # +# Non-preemptible placement # +# --------------------------------------------------------------------------- # + + +def test_nonpreemptible_is_refused_for_the_check_class() -> None: + data = _smoke_plan_data() + data["nonpreemptible"] = True + with pytest.raises(plan_lib.PlanError, match="always runs preemptible"): + plan_lib.parse_plan(data) + + +def test_every_valid_plan_has_a_runner_with_its_class_and_placement(app) -> None: + keys = { + (plan.resources.name, plan.nonpreemptible) + for plan in [ + plan_lib.parse_plan(_plan_data(stage, nonpreemptible=flag)) + for stage in plan_lib.US_ACS_LOCAL_RELEASE.stages + for flag in (False, True) + ] + + [plan_lib.parse_plan(_smoke_plan_data())] + } + assert keys <= set(app.RUNNERS) + for (name, nonpreemptible), runner in app.RUNNERS.items(): + resources = plan_lib.RESOURCE_CLASSES[name] + options = runner.modal_options + assert (options["cpu"], options["memory"], options["timeout"]) == ( + resources.cpu, + resources.memory_mib, + resources.timeout_s, + ) + assert options.get("nonpreemptible", False) is nonpreemptible + assert options["retries"] == 0 + # The check never asks for non-preemptible placement. + assert "nonpreemptible" not in app.check_stage.modal_options + + +def test_nonpreemptible_prices_every_attempt_at_three_times_list() -> None: + data = _plan_data(max_wall_seconds=20_000, nonpreemptible=True) + plan = plan_lib.parse_plan(data) + preempted = _attempt("a1", seen=3_600.0, plan_data=data) + receipt = plan_lib.build_receipt( + plan, + data, + argv=["python"], + returncode=0, + started_at="t0", + finished_at="t1", + wall_seconds=3_600.0, + peak_rss_bytes=None, + inputs_verified=[], + outputs=[], + git={}, + runner={}, + container_wall_seconds=3_600.0, + prior_attempts=[preempted], + ) + # One heavy container-hour is $1.21 at list; x3, and two hours in all. + assert receipt["estimated_usd_container_at_list_price"] == pytest.approx(3.63) + assert receipt["estimated_usd_all_attempts_at_list_price"] == pytest.approx(7.27) + # The flag is part of the plan, so switching it is a new plan digest. + assert plan_lib.plan_digest(data) != plan_lib.plan_digest( + _plan_data(max_wall_seconds=20_000) + ) From 50af66c51d467d903cc0c057b11a392299b43ab3 Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 23 Sep 2026 05:18:45 -0400 Subject: [PATCH 17/19] Fail a Modal stage launch on any receipt that is not COMPLETED The local entrypoint exited nonzero only on a nonzero return code, so a stage stopped at its max_wall_seconds budget by a tool that exits 0 on the stop signal reported success although its receipt says FAILED. It now fails on any status other than COMPLETED, says why, and prints stopped_at_budget and the all-attempts cost with the rest of the brief. Also: the pricing comment had the floor backwards (an estimate from the request is exact inside the request and a floor above it); the plan module says a Hub revision may be a branch name, the bytes being pinned by sha256; and the app's docstring lists what the check now covers. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_us_modal_stage_plan_tool.py | 36 +++++++++++++++++++ tools/modal_us_stage.py | 25 +++++++++---- tools/modal_us_stage_plan.py | 8 +++-- 3 files changed, 60 insertions(+), 9 deletions(-) diff --git a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py index fff86fe5a..ccd241c6c 100644 --- a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py +++ b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py @@ -1366,3 +1366,39 @@ def test_nonpreemptible_prices_every_attempt_at_three_times_list() -> None: assert plan_lib.plan_digest(data) != plan_lib.plan_digest( _plan_data(max_wall_seconds=20_000) ) + + +@pytest.mark.parametrize( + ("status", "returncode", "stopped", "fails"), + [ + ("COMPLETED", 0, False, False), + ("FAILED", 0, True, True), # stopped at the budget; the tool exited 0 + ("FAILED", 1, False, True), + ], +) +def test_app_main_exits_nonzero_unless_the_stage_completed( + app, monkeypatch, capsys, status, returncode, stopped, fails +) -> None: + data = _plan_data(max_wall_seconds=600) + monkeypatch.setattr(app, "_LOADED", (data, plan_lib.parse_plan(data))) + receipt = { + "status": status, + "stage": "materialize", + "run_id": "acs-local-20260923", + "wall_seconds": 600.0, + "peak_rss_bytes": 1, + "returncode": returncode, + "stopped_at_budget": stopped, + "estimated_usd_at_list_price": 0.2, + "receipt_path": "runs/acs-local-20260923/receipts/x.json", + "outputs": [], + } + runner = MagicMock() + runner.remote.return_value = receipt + monkeypatch.setitem(app.RUNNERS, ("heavy", False), runner) + if fails: + with pytest.raises(SystemExit, match=f"STAGE {status}"): + app.main(run=True) + else: + app.main(run=True) + assert json.loads(capsys.readouterr().out)["stopped_at_budget"] is stopped diff --git a/tools/modal_us_stage.py b/tools/modal_us_stage.py index e6f9eab40..e5aad10f8 100644 --- a/tools/modal_us_stage.py +++ b/tools/modal_us_stage.py @@ -10,9 +10,11 @@ H5, release directory) is mirrored to the runs volume, and every file in it is listed with its sha256 in a receipt written next to it. -Check (cheap, the default; builds the image, verifies the clone, the -environment, the tool's own argument parser on the built argv, every input -digest and the run's prior state, without running the stage):: +Check (cheap, the default; builds the image, verifies the clone and that +the plan's branch contains its commit, the environment, the tool's own +argument parser on the built argv, every input digest, the run's prior +state against its latest receipt, its attempts and budget, without running +the stage):: MICROCOSM_MODAL_PLAN=plan.json modal run tools/modal_us_stage.py @@ -901,9 +903,20 @@ def main(run: bool = False) -> None: "receipt_path", ) } - for key in ("container_wall_seconds", "estimated_usd_container_at_list_price"): + for key in ( + "stopped_at_budget", + "container_wall_seconds", + "estimated_usd_container_at_list_price", + "estimated_usd_all_attempts_at_list_price", + ): brief[key] = receipt.get(key) brief["outputs"] = len(receipt["outputs"]) print(json.dumps(brief, indent=2)) - if receipt["returncode"] != 0: - raise SystemExit(f"STAGE FAILED (receipt {receipt['receipt_path']})") + # FAILED covers a nonzero exit and a stop at the max_wall_seconds budget, + # which a tool that handles the stop signal can still end with exit 0. + if receipt["status"] != "COMPLETED": + raise SystemExit( + f"STAGE {receipt['status']} (returncode {receipt['returncode']}, " + f"stopped_at_budget {receipt['stopped_at_budget']}; " + f"receipt {receipt['receipt_path']})" + ) diff --git a/tools/modal_us_stage_plan.py b/tools/modal_us_stage_plan.py index 606c9a121..6ddd7887f 100644 --- a/tools/modal_us_stage_plan.py +++ b/tools/modal_us_stage_plan.py @@ -63,8 +63,8 @@ # Modal list prices for standard (non-sandbox) compute, read from # https://modal.com/pricing on 2026-09-22. Modal bills the higher of the -# request and actual use, so an estimate from the request is a floor only -# when the stage stays inside it. +# request and actual use, so an estimate from the request is exact while the +# stage stays inside its request and a floor when it uses more. CPU_USD_PER_CORE_SECOND = 0.0000131 MEMORY_USD_PER_GIB_SECOND = 0.00000222 # Modal applies this to the CPU and memory list price of a function set @@ -399,7 +399,9 @@ def parse_input(name: str, spec: object) -> InputRef: ``cas/sha256//`` must carry the input's own digest. ``hf:////@/`` names a file in a Hugging Face repo at an explicit revision (a public repo needs no - token). The sha256 is verified after staging either way. + token). The revision may be a tag, a commit or a branch name; the bytes + are pinned by the sha256 either way, which is verified after staging, so + a branch that has moved fails verification instead of running. """ if not isinstance(spec, Mapping) or set(spec) != {"uri", "sha256"}: From f8e85385716d65506dbf6ddd1084202d33eef97f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 23 Sep 2026 05:18:45 -0400 Subject: [PATCH 18/19] Correct the Modal runbook and drop inert env from the acceptance plan Runbook, against the evidence in _recovered/scratch-backup/893/overnight-20260923/ (status file, local materialize run.log, Modal client log): - Per-chunk figures: Modal 144.0 to 354.5 s against local 36.8 to 80.1 s (mean 66.3 over 80 chunks), 3.3 to 7.6 times per chunk and 3.93 times over chunks 1 to 29 (6,832 s against 1,737 s); loading 280.6 and 254.6 s against 79.3 s. It said 149 to 355 s against 66 to 80 s and "two to five times". The local wall is 5,604.5 s (5,569 s in the tool). The attempt's app id, launch time, runner commit and stop time are recorded, and that it left no receipt and no state. - The projection for an uninterrupted Modal materialize is now derived from the observed pace: 5.3 to 6.9 hours, $6.40 to $8.40 preemptible or $19 to $25 non-preemptible (it said 3.5 to 7 hours, $4 to $9 or $13 to $26). - A new "Preemption and restarts" section describes the ledger's outcomes and what is charged, the budget (a budget stop writes a FAILED receipt, so a relaunch of the same plan gets its whole budget), the lock and its wait, atomic mirroring and the pulled-state check, and recommends "nonpreemptible": true for heavy stages expected to run longer than about an hour, with the arithmetic: at the 0.67 preemptions an hour the run saw, a preemptible attempt is more likely than not to be cut short past 1.03 hours, and non-preemptible pays for itself in expected cost from 2.8 hours. The sample is two events; the table is labelled an order of magnitude. - How a run works: credential env refusal, input hashing (volume inputs while copied, Hub inputs after download), Hub branch revisions, the pulled-state check and the new receipt fields. The acceptance plan set four peak-limit variables that do not reach materialize at 767312d6: the two MICROCOSM_ names are read nowhere, and the two POPULACE_ names only set defaults for with_optional_acs_spine and _preflight_staging_export, which only the staging builder calls. They are removed, which changes the plan sha256 from 5ec595c7 to d1341ec0; the runbook records that the 23 September attempt ran the earlier version (b2e34fe4f). Co-Authored-By: Claude Opus 5.5 --- ...-modal-stage-acceptance-20260923-plan.json | 6 - docs/us-modal-stage-runbook.md | 251 +++++++++++++----- 2 files changed, 188 insertions(+), 69 deletions(-) diff --git a/docs/us-modal-stage-acceptance-20260923-plan.json b/docs/us-modal-stage-acceptance-20260923-plan.json index 807f5f602..54344c176 100644 --- a/docs/us-modal-stage-acceptance-20260923-plan.json +++ b/docs/us-modal-stage-acceptance-20260923-plan.json @@ -29,11 +29,5 @@ "soi_mode": "state", "hh_chunk": 20000 }, - "env": { - "MICROCOSM_ACS_POOL_PEAK_LIMIT_BYTES": "100000000000", - "MICROCOSM_STAGING_EXPORT_PEAK_LIMIT_BYTES": "100000000000", - "POPULACE_ACS_POOL_PEAK_LIMIT_BYTES": "100000000000", - "POPULACE_STAGING_EXPORT_PEAK_LIMIT_BYTES": "100000000000" - }, "max_wall_seconds": 26400 } diff --git a/docs/us-modal-stage-runbook.md b/docs/us-modal-stage-runbook.md index 84c90fe43..fb3f225f3 100644 --- a/docs/us-modal-stage-runbook.md +++ b/docs/us-modal-stage-runbook.md @@ -29,9 +29,13 @@ release directory onto the runs volume. Publication stays the human step in `run_id`, a full 40-hex commit and the branch it was pushed on, and every input as `{uri, sha256}`. Options (`soi_mode`, `hh_chunk`, `epochs` and so on) and environment overrides (`MICROCOSM_*`, `POPULACE_*`, - `*_NUM_THREADS`) come from allowlists. The runner sets the flags for - input paths, checkpoints and outputs, and a plan cannot pass them. It - also cannot pass `--allow-dirty`: every run builds from a clean clone. + `*_NUM_THREADS`) come from allowlists; an environment variable whose name + looks like a credential (`KEY`, `TOKEN`, `SECRET`, `PASSW`, `SIGNING`, + `CREDENTIAL`) is refused even under an allowed prefix. The runner sets the + flags for input paths, checkpoints and outputs, and a plan cannot pass + them. It also cannot pass `--allow-dirty`: every run builds from a clean + clone. `"nonpreemptible": true` asks for Modal's non-preemptible + placement at three times the list price (see "Preemption and restarts"). 2. **The image is the commit.** It starts from `debian_slim` with Python 3.14, the minor version the local US builds record (`runtime.python` 3.14.4 in the #974 build manifest; Modal served 3.14.2). It adds git and uv 0.11.7, makes a shallow @@ -58,15 +62,24 @@ release directory onto the runs volume. Publication stays the human step in file at an explicit revision: `hf://datasets//@/`. Public repos need no token. Each input is copied or downloaded to a stable local path, - `/work/inputs//`, hashed in the same pass, and refused on - mismatch. + `/work/inputs//`, and refused on a digest mismatch before the + tool starts: a volume input is hashed while it is copied, a Hub input + after it downloads. A Hub revision may be a branch name; the sha256 still + pins the bytes, so a branch that has moved fails verification. 4. **State lives on the runs volume.** Each `run_id` keeps `runs//state/` on `microcosm-us-stage-runs`. That directory holds the checkpoints, the calibrated H5, the release root for `package` and the stage logs. A stage pulls the state to local disk and runs the tool - there, since materialize writes a memory-mapped matrix. It then mirrors - the state back: changed files are copied and files the tool deleted are - removed. Later stages of the same run pick up from that state. Before + there, since materialize writes a memory-mapped matrix. Before anything + uses the pulled state, it is verified against the run's latest receipt + (by `finished_at`): every file the receipt lists, with its bytes and + sha256, and nothing else. With no receipt the state must be empty. A + mismatch means a stage was cut short after changing the state, and the + stage refuses; start a new `run_id`. When the tool exits the state is + mirrored back: each changed file is copied under a temporary name in its + own directory and renamed over the old one, and files the tool deleted + are removed, so a mirror cut short never leaves a half-written file. + Later stages of the same run pick up from that state. Before paying for inputs, a stage after materialize checks that the run's `checkpoints/run_identity.json` exists and pins the same staging and ladder digests. The tool re-verifies the staging digest itself @@ -74,12 +87,19 @@ release directory onto the runs volume. Publication stays the human step in 5. **Every output has a receipt.** `runs//receipts/-.json` (`microcosm-modal-us-stage-receipt/1`) records: - the plan and its sha256; - - the commit, the clone's `HEAD`, clean state and branch; + - the commit, the clone's `HEAD`, clean state, and the branch with the + remote tip it was verified against; - the sha256 of the runner files and of `uv.lock`; - - the resource class, the exact argv, the return code and wall seconds; - - peak child RSS and the list-price cost estimate; + - the resource class and placement, the exact argv, the return code and + wall seconds; + - peak child RSS and the list-price cost estimate, for this container and + for every earlier unfinished attempt of the plan; + - the attempt id and this attempt's share of `max_wall_seconds`; - every verified input; - - the sha256 of the earlier receipts in the run; + - the names, never the values, of environment variables withheld from + the tool; + - the sha256 of the earlier receipts in the run, and the receipt the + pulled state matched; - every file in the state tree with its bytes and sha256. A failed stage still mirrors its state, so `calibrate` can resume from @@ -92,9 +112,9 @@ The heavy and light classes are sized from the #974 measured peaks totals SOI surface). On the state SOI surface, the overnight build of 23 September measured materialize locally at a 77.9 GB peak, 5,569 s of tool wall and 5,472 CPU-s (4,459 targets; see the acceptance attempt below). On -Modal the same stage ran two to five times slower per chunk and held 15 to -24 GB more RSS at the same point, so the local wall times in this table -understate Modal's wall and cost. +Modal the same stage ran 3.3 to 7.6 times slower per chunk (3.93 times over +chunks 1 to 29) and held 15 to 24 GB more RSS at the same chunk, so the +local wall times in this table understate Modal's wall and cost. | Stage | Measured locally | Class | Request | List price at measured wall | | --- | --- | --- | --- | --- | @@ -109,9 +129,11 @@ The prices are Modal's list prices for standard compute, read from modal.com/pricing on 22 September 2026: $0.0000131 per core-second and $0.00000222 per GiB-second. Modal bills the higher of the request and actual use. The heavy class costs about $1.21 an hour, so an 8-hour timeout costs at -most about $9.70. The engine pass in materialize is single-threaded (CPU -seconds roughly equal wall seconds), so extra cores would not speed it up. -The table leaves out volume storage and image builds. +most about $9.70. With `"nonpreemptible": true` every figure is three times +higher: about $3.63 an hour for the heavy class. The engine pass in +materialize is single-threaded (CPU seconds roughly equal wall seconds), so +extra cores would not speed it up. The table leaves out volume storage and +image builds. ## Commands @@ -138,18 +160,23 @@ modal volume put microcosm-us-stage-inputs cas/sha256// python3 tools/modal_us_stage_plan.py validate plan.json # 3. Check on Modal (default mode; 2 cores / 8 GiB). It builds the image, -# verifies the clone, imports the environment, runs the pinned tool's own -# _parse_args on the built argv, checks every input digest (Hub inputs -# by their LFS sha256, without downloading) and the run's prior state. +# verifies the clone and that the plan's branch contains its commit, +# imports the environment, runs the pinned tool's own _parse_args on the +# built argv, checks every input digest (Hub inputs by their LFS sha256, +# without downloading), hashes the run's state against its latest +# receipt, and reports the attempts charged to the budget and any +# attempt of the run that may still be running. MICROCOSM_MODAL_PLAN=plan.json modal run tools/modal_us_stage.py # 4. Run the stage. --detach keeps it running if this terminal goes away; # the receipt lands on the runs volume either way. Set "max_wall_seconds" # in the plan to cap the cost below the class's hard timeout: the runner # stops the tool then, and the receipt says FAILED, stopped_at_budget. -# The budget covers every attempt: when Modal restarts a preempted -# container, the time the cut-short attempts ran comes off it (see -# "Preemption" below). +# The budget covers every attempt that never finished: when Modal +# restarts a preempted container, the time the cut-short attempts ran +# comes off it (see "Preemption and restarts" below). A stop at the +# budget writes a receipt, so launching the same plan again gets the +# whole budget again. MICROCOSM_MODAL_PLAN=plan.json modal run --detach tools/modal_us_stage.py --run # 5. Fetch the state and verify it against the receipt. @@ -170,7 +197,8 @@ and receipt code as a real stage. Next stage: copy the plan, change `stage`, keep the `run_id`, and drop `feed` if you like, since only materialize reads it. Run steps 3 and 4 again. Run one stage of a run at a time, because two concurrent stages would race on the -same state directory. +same state directory. The attempt ledger enforces this, best effort (see +"Preemption and restarts"). For a gated or private Hub input, set `MICROCOSM_MODAL_HF_SECRET` to the name of a Modal secret that holds `HF_TOKEN`, for example `huggingface-token` in @@ -229,8 +257,12 @@ overnight build and the Hub were not touched. recorded 3.14.4. The heavy materialize replay has not been run. It needs the staging H5 on -the inputs volume, which is about 10.7 GB to upload, plus one heavy run at -about $1.71 list price for the measured wall time. That run is the next +the inputs volume, which is about 10.7 GB to upload, plus one heavy run. +That is about $1.71 at the local wall time, but at the Modal pace measured +on 23 September (3.93 times local, below) it is closer to 5.5 hours: about +$6.70 if nothing preempts it, or about $20 non-preemptible, which a run that +long needs (see "Preemption and restarts"). Add `"nonpreemptible": true` and +a `max_wall_seconds` to a copy of the plan first. That run is the next step: ```bash @@ -249,10 +281,20 @@ The plan `docs/us-modal-stage-acceptance-20260923-plan.json` ran the materialize stage of the overnight build of 23 September on Modal (`run_id` `overnight-20260923-materialize-state`). It pinned build commit `767312d6` on branch `overnight-acs-local-20260923`, `soi_mode` `state`, -`hh_chunk` 20,000, the local run's peak-limit environment and a -26,400-second wall budget. The same stage ran locally at the same time from -the same inputs. The Modal run did not finish, so its outputs could not be -compared. +`hh_chunk` 20,000 and a 26,400-second wall budget. The same stage ran +locally at the same time from the same inputs. The Modal run did not finish, +so its outputs could not be compared. + +The plan as launched (`b2e34fe4f`, plan sha256 `5ec595c7…`) also set four +peak-limit variables copied from the local run's environment: +`MICROCOSM_ACS_POOL_PEAK_LIMIT_BYTES`, +`MICROCOSM_STAGING_EXPORT_PEAK_LIMIT_BYTES` and their `POPULACE_` twins. At +`767312d6` none of them reaches materialize. The two `MICROCOSM_` names are +read nowhere. The two `POPULACE_` names only set defaults for +`with_optional_acs_spine` and `_preflight_staging_export`, which only the +staging builder calls. They were removed from the committed plan, whose +sha256 is now `d1341ec0…`, so a relaunch of this file is a new plan to the +attempt ledger. - **Inputs.** The staging H5 (`ed2e6308…`, 10,685,765,051 bytes) went to the inputs volume in 255 seconds, about 42 MB/s, and the staging summary @@ -262,17 +304,24 @@ compared. policyengine-us 2.2.1 and policyengine-core 3.32.5, the same versions as the local environment. The tool's parser accepted the argv, and all four inputs matched their digests on the volume. -- **Run.** The container started at 06:03:06 UTC and the tool's first - log line came at 06:06:22, so staging and verifying the inputs took under - 3.5 minutes. The tool's own timings, Modal against local: +- **Run.** It was launched detached at 06:02:58 UTC (02:02:58 EDT) as app + `ap-mzp14wEyVLIFlYX5qQQxFZ`, on runner commit `1d80287af`, before the + attempt ledger existed. The container started at 06:03:06 UTC and the + tool's first log line came at 06:06:22, so staging and verifying the + inputs took under 3.5 minutes. The tool's own timings, Modal against + local: | Step | Modal | Local | | --- | --- | --- | - | Start to staging frame loaded (hashing, specs, load) | 255 to 281 s | 79 s | - | One chunk of 20,000 households | 149 to 355 s | 66 to 80 s | + | Start to staging frame loaded (hashing, specs, load) | 280.6 and 254.6 s | 79.3 s | + | One chunk of 20,000 households | 144.0 to 354.5 s | 36.8 to 80.1 s (mean 66.3 s over 80) | + | Chunks 1 to 29 in all (second attempt) | 6,832 s | 1,737 s | | peak RSS after chunk 1 | 66.3 GB | 51.0 GB | | peak RSS after chunk 29 | 74.5 GB | 51.0 GB | + Per chunk Modal was 3.3 to 7.6 times slower, and 3.93 times over chunks 1 + to 29. + The container ran gVisor (`Linux-4.19.0-gvisor-x86_64`) with Python 3.14.2. In the check-class container Modal set `OMP_NUM_THREADS`, `OPENBLAS_NUM_THREADS`, `MKL_NUM_THREADS` and `BLIS_NUM_THREADS` to that @@ -286,16 +335,20 @@ compared. minutes (after chunk 12 of 80) and the second after 2 hours 4 minutes (after chunk 29). Each time it restarted the function from zero. The run was launched before the attempt ledger existed, so the budget restarted - too. It was stopped by hand (`modal app stop`) 34 seconds into a third - attempt, because at the observed pace a third attempt needed 6.5 to 7 - more hours and would have taken the total past $10. + too. It was stopped by hand (`modal app stop`, 09:02:22 UTC) 34 seconds + into a third attempt, because at the observed pace a third attempt needed + 6.5 to 7 more hours and would have taken the total past $10. It left no + receipt and no state: that runner mirrors state and writes the receipt + only when the tool exits, and the runs volume has no + `runs/overnight-20260923-materialize-state`. - **Cost.** Modal's workspace billing report (`modal.billing.workspace_billing_report`, hourly) shows $3.57 for the run through 09:00 UTC. The last 2.5 minutes add about $0.05, for about $3.62 in total. The check cost $0.02. Billing was at the request, about $1.21 an hour. - **Local result for the next comparison.** The local run finished in - 5,569 s with a 77.9 GB peak. `run_identity`: staging `ed2e6308…`, ladder + 5,604.5 s of wall (5,569 s in the tool) with a 77.9 GB peak (the tool's + log; 76.7 GB by the supervisor). `run_identity`: staging `ed2e6308…`, ladder `39a2ab2a…`, 1,588,854 households, 4,459 targets (3,972 admin specs declared and compiled, plus 487 state and CD population cells), no dropped cells, `targets_sha256` @@ -309,14 +362,99 @@ compared. local file next to the overnight run. At the pace observed, one uninterrupted Modal materialize on the state -surface takes about 3.5 to 7 hours. That is about $4 to $9 on preemptible -placement, or $13 to $26 with `"nonpreemptible": true`, against 1.5 hours -locally. The Modal path works, but it does not yet make this stage cheaper -or faster. The run did show that the stage runs as one long single-core +surface takes about 5.3 to 6.9 hours: 80 chunks at the second attempt's +mean of 235.6 s, or at its last five chunks' 308.2 s, plus about 4.5 +minutes of loading. That is about $6.40 to $8.40 at the preemptible list +price if nothing preempts it, and $19 to $25 with `"nonpreemptible": true`, +against 1.6 hours locally. At the preemption rate this run saw, a 5- to +7-hour preemptible run finishes uninterrupted 1 to 3% of the time, so +non-preemptible is the practical placement for it. The Modal path works, +but it does not yet make this stage cheaper or faster. The run did show that the stage runs as one long single-core loop over 80 chunks of households. If the chunks are independent, which this runbook has not checked in `materialize_chunked`, fanning them out across containers would let preemption lose one chunk instead of the run. +## Preemption and restarts + +Functions run on Modal's preemptible placement unless the plan sets +`"nonpreemptible": true`. When Modal preempts a container it restarts the +function on the same input, from zero, whatever `retries` says +(modal.com/docs/guide/preemption). The 23 September acceptance run was +preempted twice. The state is mirrored and the receipt written only when +the tool exits, so a preempted attempt's tool work is lost, but it is +billed. + +**The attempt ledger.** Each attempt writes +`runs//attempts/-.json` when it starts and +rewrites it every 120 seconds until it ends: through the lock wait, input +staging, the tool, hashing and mirroring. The record ends with one outcome. + +| Outcome | When | Charged to the budget | +| --- | --- | --- | +| `receipt` | the tool ran and a receipt was written, COMPLETED or FAILED | no | +| `refused` | the lock or the budget stopped it before it staged anything | no | +| `error` | any exception after its first record: a digest mismatch, a failed pull or mirror, a pulled-state mismatch, a runner bug | yes | +| none | preempted (or still running) | yes | + +**The budget.** A new attempt charges every unfinished attempt of the same +plan (same `run_id`, stage and plan sha256) to `max_wall_seconds`. Its tool +gets what is left, and it refuses to start with less than 60 seconds. A +preempted attempt is charged from the start of `_run_stage` to its last +record, so the charge leaves out the container's cold start and image load, +up to 120 seconds after the last record (more if a heartbeat write failed), +and the preemption grace period. A stop at the budget is not a preemption: +the tool is stopped, the state mirrored and a FAILED receipt with +`stopped_at_budget` written, so that attempt is finished and launching the +same plan again gets the whole budget again. To launch past a spent budget, +raise `max_wall_seconds` or use a new `run_id`. Any change to the plan, +`nonpreemptible` included, is a new plan sha256 and a fresh budget. + +**The lock.** Two attempts of one run would race on its state directory, +so the ledger is also the run's lock. An attempt refuses to start while an +earlier attempt of the same run, of any stage or plan, is still writing its +record. A record less than 300 seconds old is ambiguous, because Modal +restarts a preempted input within moments. The new attempt waits 270 +seconds (two heartbeats) and reads again. A record that moved belongs to a +running attempt, and the new one refuses; a record that did not move +belongs to a dead one, and the new attempt starts. A restart after +preemption can therefore spend up to 4.5 more minutes. There is no +override: a dead attempt stops blocking after one wait, and a running one +must not be raced. Modal volumes have no atomic lock, so two attempts that +start within one volume commit of each other can both pass; run one stage +of a run at a time. The check reports recent records without waiting. + +**The state after a cut.** The mirror writes each file under a temporary +name and renames it, so it never leaves a half-written file. A preemption +during mirroring can still leave some files new and some old. The next +attempt verifies the pulled state against the run's latest receipt and +refuses that mix; start a new `run_id`. + +**When to set `"nonpreemptible": true`.** Non-preemptible placement costs +three times the list price for CPU and memory (modal.com/docs/guide/preemption +and modal.com/pricing, read 23 September 2026): about $3.63 an hour for the +heavy class instead of $1.21. Set it for heavy stages expected to run +longer than about an hour. The acceptance run was preempted twice in about +2.97 hours of running (after 54 minutes and after 2 hours 4 minutes), a +rate λ of about 0.67 an hour. Each preemption restarts the stage from zero, +so at that rate a stage of T hours on preemptible placement takes +(e^(λT) − 1)/λ hours in expectation: + +| Stage length | Uninterrupted on preemptible | Expected preemptible hours and cost | Non-preemptible cost | +| --- | --- | --- | --- | +| 1 h | 51% | 1.4 h, $1.73 | $3.63 | +| 2 h | 26% | 4.2 h, $5.12 | $7.27 | +| 3 h | 13% | 9.7 h, $11.78 | $10.90 | +| 6 h | 2% | 83 h, $101 | $21.81 | + +Past about an hour a preemptible attempt is more likely than not to be cut +short. Non-preemptible placement pays for itself in expected cost from +about 2.8 hours. Before that it costs up to three times as much for a +result that arrives on time. A `max_wall_seconds` budget also usually stops +a long preemptible run before it finishes. Two preemptions in one run is a +small sample: a 95% interval for the rate runs from about 0.08 to 2.4 an +hour, so read the table as an order of magnitude. Light stages and short +heavy stages (calibrate took 7 minutes locally) are cheaper preemptible. + ## Data placement Inputs go to the `policyengine` Modal workspace. Upload only files that are @@ -343,25 +481,12 @@ hold digests, sizes and paths, never file contents. registered either. It takes a directory input (`--inputs-dir`, the ACS PUMS archive cache), which the plan format does not support. Supporting it would take an archive digest plus extraction. -- **Preemption and out-of-memory kills.** Functions run on Modal's default - (preemptible) placement. When Modal preempts a container it restarts the - function on the same input, from scratch, whatever `retries` says; this - happened to the 23 September acceptance run after 54 minutes. State is - only mirrored and the receipt only written when the tool exits, so the - cut-short attempt leaves no receipt and its work is lost, but it is - billed. Each attempt therefore writes - `runs//attempts/-.json` when it starts and rewrites - it every 120 seconds. A restart charges the time of every earlier attempt - of the same plan that never wrote a receipt to `max_wall_seconds`, and - refuses to start when less than a minute is left. The check reports those - attempts and the time left. To launch again past a spent budget, raise - `max_wall_seconds` (a new plan digest) or use a new `run_id`. A plan can - set `"nonpreemptible": true` to run on Modal's non-preemptible placement, - which avoids restarts at three times the list price for CPU and memory - (modal.com/docs/guide/preemption, read 23 September 2026); `validate` and - the receipt price it that way. The heavy class sets a memory request but - no hard limit. What Modal does with a container killed for memory has not - been observed here. +- **Work lost to preemption, and memory kills.** A preempted attempt's + tool work is lost: the state is mirrored only when the tool exits, and + materialize has no resume point inside its chunk loop. Non-preemptible + placement avoids preemption (see "Preemption and restarts"). The heavy + class sets a memory request but no hard limit. What Modal does with a + container killed for memory has not been observed here. - **Certification.** A receipt proves which bytes a stage produced. It does not certify a release. Preflight (`tools/preflight_us_release_gates.py`) and certification still run on the output as before. From e40c93be1f32a332346f541b8b137ae294c7a94f Mon Sep 17 00:00:00 2001 From: Max Ghenis Date: Wed, 23 Sep 2026 08:37:13 -0400 Subject: [PATCH 19/19] Name the thread-count overrides a Modal stage plan may set The runbook and the refusal message said any `*_NUM_THREADS` override was allowed, but `_ENV_KEY` admitted only OMP, MKL, OPENBLAS and NUMEXPR, so a plan overriding BLIS_NUM_THREADS (which Modal itself sets in the container) was refused with a message saying it was allowed. The allowlist is now the explicit `THREAD_ENV_KEYS` tuple, BLIS included; the refusal message and the runbook list the five names, and tests pin every listed name as allowed and an unlisted *_NUM_THREADS name as refused. Co-Authored-By: Claude Opus 5.5 --- docs/us-modal-stage-runbook.md | 6 ++++-- .../tests/test_us_modal_stage_plan_tool.py | 10 ++++++++++ tools/modal_us_stage_plan.py | 14 ++++++++++++-- 3 files changed, 26 insertions(+), 4 deletions(-) diff --git a/docs/us-modal-stage-runbook.md b/docs/us-modal-stage-runbook.md index fb3f225f3..0f198d13b 100644 --- a/docs/us-modal-stage-runbook.md +++ b/docs/us-modal-stage-runbook.md @@ -28,8 +28,10 @@ release directory onto the runs volume. Publication stays the human step in `docs/us-modal-stage-example-plan.json`). It names the tool, the stage, a `run_id`, a full 40-hex commit and the branch it was pushed on, and every input as `{uri, sha256}`. Options (`soi_mode`, `hh_chunk`, `epochs` and - so on) and environment overrides (`MICROCOSM_*`, `POPULACE_*`, - `*_NUM_THREADS`) come from allowlists; an environment variable whose name + so on) and environment overrides (`MICROCOSM_*`, `POPULACE_*`, and the + thread counts `OMP_NUM_THREADS`, `MKL_NUM_THREADS`, + `OPENBLAS_NUM_THREADS`, `NUMEXPR_NUM_THREADS` and `BLIS_NUM_THREADS`) + come from allowlists; an environment variable whose name looks like a credential (`KEY`, `TOKEN`, `SECRET`, `PASSW`, `SIGNING`, `CREDENTIAL`) is refused even under an allowed prefix. The runner sets the flags for input paths, checkpoints and outputs, and a plan cannot pass diff --git a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py index ccd241c6c..7e77ea8ae 100644 --- a/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py +++ b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py @@ -1402,3 +1402,13 @@ def test_app_main_exits_nonzero_unless_the_stage_completed( else: app.main(run=True) assert json.loads(capsys.readouterr().out)["stopped_at_budget"] is stopped + + +@pytest.mark.parametrize("key", plan_lib.THREAD_ENV_KEYS) +def test_every_documented_thread_count_is_allowlisted(key) -> None: + assert plan_lib._ENV_KEY.fullmatch(key) + + +@pytest.mark.parametrize("key", ["VECLIB_MAXIMUM_THREADS", "FOO_NUM_THREADS"]) +def test_an_undocumented_thread_variable_is_refused_by_name(key) -> None: + assert not plan_lib._ENV_KEY.fullmatch(key) diff --git a/tools/modal_us_stage_plan.py b/tools/modal_us_stage_plan.py index 6ddd7887f..559f9d7f0 100644 --- a/tools/modal_us_stage_plan.py +++ b/tools/modal_us_stage_plan.py @@ -80,8 +80,17 @@ _FILENAME = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,199}") _HF_REPO_ID = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9][A-Za-z0-9._-]*") _HF_REVISION = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,199}") +#: The thread-count variables a plan may override. BLIS is here because Modal +#: sets BLIS_NUM_THREADS in the container alongside the other three. +THREAD_ENV_KEYS = ( + "OMP_NUM_THREADS", + "MKL_NUM_THREADS", + "OPENBLAS_NUM_THREADS", + "NUMEXPR_NUM_THREADS", + "BLIS_NUM_THREADS", +) _ENV_KEY = re.compile( - r"(?:MICROCOSM|POPULACE)_[A-Z0-9_]+|(?:OMP|MKL|OPENBLAS|NUMEXPR)_NUM_THREADS" + r"(?:MICROCOSM|POPULACE)_[A-Z0-9_]+|(?:OMP|MKL|OPENBLAS|NUMEXPR|BLIS)_NUM_THREADS" ) # Names that look like credentials. A plan may not set one, even under an # allowlisted prefix (its value would be copied into every receipt), and the @@ -614,7 +623,8 @@ def parse_plan(data: object) -> Plan: for key in sorted(raw_env): if not isinstance(key, str) or not _ENV_KEY.fullmatch(key): raise PlanError( - f"env {key!r} is not allowlisted (MICROCOSM_*, POPULACE_*, *_NUM_THREADS)" + f"env {key!r} is not allowlisted (MICROCOSM_*, POPULACE_*, or one " + f"of {', '.join(THREAD_ENV_KEYS)})" ) if is_credential_env_key(key): raise PlanError(