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..54344c176 --- /dev/null +++ b/docs/us-modal-stage-acceptance-20260923-plan.json @@ -0,0 +1,33 @@ +{ + "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 + }, + "max_wall_seconds": 26400 +} 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/docs/us-modal-stage-runbook.md b/docs/us-modal-stage-runbook.md new file mode 100644 index 000000000..0f198d13b --- /dev/null +++ b/docs/us-modal-stage-runbook.md @@ -0,0 +1,494 @@ +# 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_*`, 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 + 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 + 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. + 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. +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//`, 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. 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 + (`_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 the branch with the + remote tip it was verified against; + - the sha256 of the runner files and of `uv.lock`; + - 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 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 + `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). 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 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 | +| --- | --- | --- | --- | --- | +| 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. 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 + +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 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 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. +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 +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 +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 + +`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. + +## 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. +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 +# 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 +``` + +## 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 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 + (`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.** 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) | 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 + 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`, 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,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` + `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 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 +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. +- **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. 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 new file mode 100644 index 000000000..7e77ea8ae --- /dev/null +++ b/packages/microcosm-build/tests/test_us_modal_stage_plan_tool.py @@ -0,0 +1,1414 @@ +"""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. 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 + +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() + + +@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 + # @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) + 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 +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) + + +@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") + 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") + 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() + 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"} + + +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, + "nonpreemptible": False, + } + 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_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" + + +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 + + +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_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 = 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" + 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"] + + +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 + ) + + +# --------------------------------------------------------------------------- # +# 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 + + +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"}) + + +# --------------------------------------------------------------------------- # +# 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"] + + +# --------------------------------------------------------------------------- # +# 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"}), + ] + + +# --------------------------------------------------------------------------- # +# 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) + ) + + +@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 + + +@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.py b/tools/modal_us_stage.py new file mode 100644 index 000000000..e5aad10f8 --- /dev/null +++ b/tools/modal_us_stage.py @@ -0,0 +1,922 @@ +"""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 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 + +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 platform +import resource +import shutil +import subprocess +import sys +import tempfile +import threading +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"), + **_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], + # 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__, + }.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], 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: + 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"] + 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, + env=env, + text=True, + capture_output=True, + timeout=900, + ) + return { + "returncode": proc.returncode, + "stdout_tail": proc.stdout[-2000:], + "stderr_tail": proc.stderr[-4000:], + } + + +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 _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 record in ``runs//attempts/``. + + Modal restarts a preempted function from scratch on the same input, and + ``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: + 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.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, **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(), + modal=self.modal, + **final, # type: ignore[arg-type] + ) + self.path.parent.mkdir(parents=True, exist_ok=True) + 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: + self.write() + except Exception as error: # noqa: BLE001 - a missed beat only undercounts + 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(): + 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: + 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["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}") + + # 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", + "-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, + env=tool_env, + 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, 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") + + 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) + # 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) + ) + 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" + ) + # 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) + git = _git_state(plan) + 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}") + + run_dir = _run_dir(plan) + runs_volume.reload() + + # 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)) + + # 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) + 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 _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. " + "To launch again, raise max_wall_seconds (a new plan digest) or use a " + "new run_id" + ) + if prior_attempts: + print( + f"ATTEMPTS: {len(prior_attempts)} earlier unfinished attempt(s); " + f"tool budget for this attempt {budget_seconds}s", + flush=True, + ) + + # 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 = _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) + 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, 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( + argv, + cwd=plan_lib.IMAGE_REPO_ROOT, + env=env, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + budget = _BudgetWatch(proc, budget_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() + + # 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( + 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, + "tool_env_removed": env_removed, + "modal": attempt.modal, + }, + prior_receipts=prior, + stopped_at_budget=budget.fired, + container_wall_seconds=time.time() - attempt.started, + 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" + 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 + + +@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.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, + 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) + + +@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, + 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", 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, +} + + +@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, plan.nonpreemptible)] + 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", + ) + } + 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)) + # 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 new file mode 100644 index 000000000..559f9d7f0 --- /dev/null +++ b/tools/modal_us_stage_plan.py @@ -0,0 +1,1460 @@ +#!/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 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 +# 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. +_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}") +#: 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|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 +# 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 # +# --------------------------------------------------------------------------- # + + +@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, 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 * multiplier, 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 + # 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] + 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", +) + +# 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) +} + + +# --------------------------------------------------------------------------- # +# 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.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.fullmatch(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 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"}: + 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.fullmatch(revision): + raise PlanError( + f"input {name!r}: hf URI needs an explicit @revision " + "(tag, branch without '/', or commit)" + ) + 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( + 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) + # 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: + return self.tool.stages[self.stage] + + @property + 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", + "tool", + "stage", + "run_id", + "source", + "inputs", + "options", + "env", + "max_wall_seconds", + "nonpreemptible", +} +# Time the container keeps for staging inputs, hashing and mirroring state. +_RUNNER_OVERHEAD_SECONDS = 15 * 60 + + +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_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 {tool_name!r}; known: {sorted(TOOLS)}") + stage = data.get("stage") + 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.fullmatch(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.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.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.fullmatch(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 isinstance(key, str) or not _ENV_KEY.fullmatch(key): + raise PlanError( + 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( + 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] + + 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}" + ) + + 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, + run_id=run_id, + repo_url=repo_url, + commit=commit, + branch=branch, + inputs=inputs, + options=options, + env=env, + max_wall_seconds=max_wall, + nonpreemptible=nonpreemptible, + ) + + +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) + 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 {commit}", + f"git -C {root} checkout -q -B {shlex.quote(plan.branch)} FETCH_HEAD", + 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)"', + ] + + +_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()): + 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) + 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]: + 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"), + ) + + +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 # +# --------------------------------------------------------------------------- # + +_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}" + ) + + +# 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. + + 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()] = ( + 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]: + """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) + 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), + "partials_removed": len(partials), + } + + +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]] = (), + stopped_at_budget: bool = False, + container_wall_seconds: float | None = None, + 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 + # 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, plan.price_multiplier + ), + "estimated_usd_all_attempts_at_list_price": resources.estimated_usd( + container_wall_seconds + prior_seconds, plan.price_multiplier + ), + } + return { + "schema": RECEIPT_SCHEMA, + "status": "COMPLETED" + if returncode == 0 and not stopped_at_budget + 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, + "nonpreemptible": plan.nonpreemptible, + }, + "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, + "prior_unfinished_attempts": [dict(item) for item in prior_attempts], + "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, plan.price_multiplier + ), + **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], + } + + +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 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. + + 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 + + +# --------------------------------------------------------------------------- # +# Attempts: a budget that holds across preemption restarts # +# --------------------------------------------------------------------------- # + +ATTEMPT_SCHEMA = "microcosm-modal-us-stage-attempt/1" +# 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, + plan_sha256: str, + *, + attempt_id: str, + started_epoch: float, + 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, + "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, + "outcome": outcome, + "note": note, + "modal": dict(modal or {}), + } + + +def unfinished_attempts( + records: Iterable[Mapping], + plan: Plan, + plan_sha256: str, + *, + exclude: str | None = None, +) -> list[dict[str, object]]: + """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. 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 [ + 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 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: + """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)) + 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, + "nonpreemptible": plan.nonpreemptible, + }, + "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, 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 + + +# --------------------------------------------------------------------------- # +# 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())