Skip to content

Commit 3fae1bb

Browse files
baogorekbaogorek
authored
Merge pull request #538 from PolicyEngine/calibration-pipeline-improvements
Add calibration package checkpointing, target config, and hyperparameter CLI
2 parents ca2000f + 121d603 commit 3fae1bb

65 files changed

Lines changed: 11500 additions & 2872 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.github/workflows/local_area_publish.yaml‎

Lines changed: 59 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ on:
55
# push:
66
# branches: [main]
77
# paths:
8-
# - 'policyengine_us_data/datasets/cps/local_area_calibration/**'
8+
# - 'policyengine_us_data/calibration/**'
99
# - '.github/workflows/local_area_publish.yaml'
1010
# - 'modal_app/**'
1111
# repository_dispatch:
@@ -24,7 +24,7 @@ on:
2424
type: boolean
2525

2626
# Trigger strategy:
27-
# 1. Automatic: Code changes to local_area_calibration/ pushed to main
27+
# 1. Automatic: Code changes to calibration/ pushed to main
2828
# 2. repository_dispatch: Calibration workflow triggers after uploading new weights
2929
# 3. workflow_dispatch: Manual trigger with optional parameters
3030

@@ -72,5 +72,60 @@ jobs:
7272
echo "" >> $GITHUB_STEP_SUMMARY
7373
echo "Files have been uploaded to GCS and staged on HuggingFace." >> $GITHUB_STEP_SUMMARY
7474
echo "" >> $GITHUB_STEP_SUMMARY
75-
echo "### Next step: Promote to production" >> $GITHUB_STEP_SUMMARY
76-
echo "Trigger the **Promote Local Area H5 Files** workflow with the version from the build output." >> $GITHUB_STEP_SUMMARY
75+
echo "### Next step: Validation runs automatically" >> $GITHUB_STEP_SUMMARY
76+
echo "The validate-staging job will now check all staged H5s." >> $GITHUB_STEP_SUMMARY
77+
78+
validate-staging:
79+
needs: publish-local-area
80+
runs-on: ubuntu-latest
81+
env:
82+
HUGGING_FACE_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }}
83+
steps:
84+
- name: Checkout repo
85+
uses: actions/checkout@v4
86+
87+
- name: Set up Python
88+
uses: actions/setup-python@v5
89+
with:
90+
python-version: '3.13'
91+
92+
- name: Set up uv
93+
uses: astral-sh/setup-uv@v5
94+
95+
- name: Install dependencies
96+
run: uv sync
97+
98+
- name: Validate staged H5s
99+
run: |
100+
uv run python -m policyengine_us_data.calibration.validate_staging \
101+
--area-type states --output validation_results.csv
102+
103+
- name: Upload validation results to HF
104+
run: |
105+
uv run python -c "
106+
from policyengine_us_data.utils.huggingface import upload
107+
upload('validation_results.csv',
108+
'policyengine/policyengine-us-data',
109+
'calibration/logs/validation_results.csv')
110+
"
111+
112+
- name: Post validation summary
113+
if: always()
114+
run: |
115+
echo "## Validation Results" >> $GITHUB_STEP_SUMMARY
116+
if [ -f validation_results.csv ]; then
117+
TOTAL=$(tail -n +2 validation_results.csv | wc -l)
118+
FAILS=$(grep -c ',FAIL,' validation_results.csv || true)
119+
echo "- **${TOTAL}** targets validated" >> $GITHUB_STEP_SUMMARY
120+
echo "- **${FAILS}** sanity failures" >> $GITHUB_STEP_SUMMARY
121+
echo "" >> $GITHUB_STEP_SUMMARY
122+
echo "Review in dashboard, then trigger **Promote** workflow." >> $GITHUB_STEP_SUMMARY
123+
else
124+
echo "Validation did not produce output." >> $GITHUB_STEP_SUMMARY
125+
fi
126+
127+
- name: Upload validation artifact
128+
uses: actions/upload-artifact@v4
129+
with:
130+
name: validation-results
131+
path: validation_results.csv

‎.gitignore‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,12 +30,12 @@ docs/.ipynb_checkpoints/
3030
## ACA PTC state-level uprating factors
3131
!policyengine_us_data/storage/aca_ptc_multipliers_2022_2024.csv
3232

33-
## Raw input cache for database pipeline
34-
policyengine_us_data/storage/calibration/raw_inputs/
33+
## Calibration run outputs (weights, diagnostics, packages, config)
34+
policyengine_us_data/storage/calibration/
3535

3636
## Batch processing checkpoints
3737
completed_*.txt
3838

3939
## Test fixtures
40-
!policyengine_us_data/tests/test_local_area_calibration/test_fixture_50hh.h5
40+
!policyengine_us_data/tests/test_calibration/test_fixture_50hh.h5
4141
oregon_ctc_analysis.py

‎Makefile‎

Lines changed: 140 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,12 @@
1-
.PHONY: all format test install download upload docker documentation data validate-data calibrate publish-local-area clean build paper clean-paper presentations database database-refresh promote-database promote-dataset
1+
.PHONY: all format test install download upload docker documentation data validate-data calibrate calibrate-build publish-local-area upload-calibration upload-dataset upload-database push-to-modal build-matrices calibrate-modal calibrate-modal-national calibrate-both stage-h5s stage-national-h5 stage-all-h5s pipeline validate-staging validate-staging-full upload-validation check-staging check-sanity clean build paper clean-paper presentations database database-refresh promote-database promote-dataset promote build-h5s validate-local
2+
3+
GPU ?= A100-80GB
4+
EPOCHS ?= 200
5+
NATIONAL_GPU ?= T4
6+
NATIONAL_EPOCHS ?= 200
7+
BRANCH ?= $(shell git rev-parse --abbrev-ref HEAD)
8+
NUM_WORKERS ?= 8
9+
VERSION ?=
210

311
HF_CLONE_DIR ?= $(HOME)/huggingface/policyengine-us-data
412

@@ -79,8 +87,8 @@ promote-database:
7987
@echo "Copied DB and raw_inputs to HF clone. Now cd to HF repo, commit, and push."
8088

8189
promote-dataset:
82-
cp policyengine_us_data/storage/stratified_extended_cps_2024.h5 \
83-
$(HF_CLONE_DIR)/calibration/stratified_extended_cps.h5
90+
cp policyengine_us_data/storage/source_imputed_stratified_extended_cps_2024.h5 \
91+
$(HF_CLONE_DIR)/calibration/source_imputed_stratified_extended_cps.h5
8492
@echo "Copied dataset to HF clone. Now cd to HF repo, commit, and push."
8593

8694
data: download
@@ -90,20 +98,146 @@ data: download
9098
python policyengine_us_data/datasets/puf/irs_puf.py
9199
python policyengine_us_data/datasets/puf/puf.py
92100
python policyengine_us_data/datasets/cps/extended_cps.py
101+
python policyengine_us_data/calibration/create_stratified_cps.py
102+
python policyengine_us_data/calibration/create_source_imputed_cps.py
103+
104+
data-legacy: data
93105
python policyengine_us_data/datasets/cps/enhanced_cps.py
94106
python policyengine_us_data/datasets/cps/small_enhanced_cps.py
95-
python policyengine_us_data/datasets/cps/local_area_calibration/create_stratified_cps.py
96107

97108
calibrate: data
98109
python -m policyengine_us_data.calibration.unified_calibration \
99-
--puf-dataset policyengine_us_data/storage/puf_2024.h5
110+
--target-config policyengine_us_data/calibration/target_config.yaml
111+
112+
calibrate-build: data
113+
python -m policyengine_us_data.calibration.unified_calibration \
114+
--target-config policyengine_us_data/calibration/target_config.yaml \
115+
--build-only
116+
117+
validate-package:
118+
python -m policyengine_us_data.calibration.validate_package
100119

101120
publish-local-area:
102-
python policyengine_us_data/datasets/cps/local_area_calibration/publish_local_area.py
121+
python policyengine_us_data/calibration/publish_local_area.py --upload
122+
123+
build-h5s:
124+
python -m policyengine_us_data.calibration.publish_local_area \
125+
--weights-path policyengine_us_data/storage/calibration/calibration_weights.npy \
126+
--dataset-path policyengine_us_data/storage/source_imputed_stratified_extended_cps_2024.h5 \
127+
--n-clones 430 \
128+
--seed 42 \
129+
--states-only
130+
131+
validate-local:
132+
python -m policyengine_us_data.calibration.validate_staging \
133+
--hf-prefix local_area_build \
134+
--area-type states --output validation_results.csv
103135

104136
validate-data:
105137
python -c "from policyengine_us_data.storage.upload_completed_datasets import validate_all_datasets; validate_all_datasets()"
106138

139+
upload-calibration:
140+
python -c "from policyengine_us_data.utils.huggingface import upload_calibration_artifacts; \
141+
upload_calibration_artifacts()"
142+
143+
upload-dataset:
144+
python -c "from policyengine_us_data.utils.huggingface import upload; \
145+
upload('policyengine_us_data/storage/source_imputed_stratified_extended_cps_2024.h5', \
146+
'policyengine/policyengine-us-data', \
147+
'calibration/source_imputed_stratified_extended_cps.h5')"
148+
@echo "Dataset uploaded to HF."
149+
150+
upload-database:
151+
python -c "from policyengine_us_data.utils.huggingface import upload; \
152+
upload('policyengine_us_data/storage/calibration/policy_data.db', \
153+
'policyengine/policyengine-us-data', \
154+
'calibration/policy_data.db')"
155+
@echo "Database uploaded to HF."
156+
157+
push-to-modal:
158+
modal volume put local-area-staging \
159+
policyengine_us_data/storage/calibration/calibration_weights.npy \
160+
calibration_inputs/calibration/calibration_weights.npy --force
161+
modal volume put local-area-staging \
162+
policyengine_us_data/storage/calibration/stacked_blocks.npy \
163+
calibration_inputs/calibration/stacked_blocks.npy --force
164+
modal volume put local-area-staging \
165+
policyengine_us_data/storage/calibration/stacked_takeup.npz \
166+
calibration_inputs/calibration/stacked_takeup.npz --force
167+
modal volume put local-area-staging \
168+
policyengine_us_data/storage/calibration/policy_data.db \
169+
calibration_inputs/calibration/policy_data.db --force
170+
modal volume put local-area-staging \
171+
policyengine_us_data/storage/calibration/geo_labels.json \
172+
calibration_inputs/calibration/geo_labels.json --force
173+
modal volume put local-area-staging \
174+
policyengine_us_data/storage/source_imputed_stratified_extended_cps_2024.h5 \
175+
calibration_inputs/calibration/source_imputed_stratified_extended_cps.h5 --force
176+
@echo "All calibration inputs pushed to Modal volume."
177+
178+
build-matrices:
179+
modal run modal_app/remote_calibration_runner.py::build_package \
180+
--branch $(BRANCH)
181+
182+
calibrate-modal:
183+
modal run modal_app/remote_calibration_runner.py::main \
184+
--branch $(BRANCH) --gpu $(GPU) --epochs $(EPOCHS) \
185+
--push-results
186+
187+
calibrate-modal-national:
188+
modal run modal_app/remote_calibration_runner.py::main \
189+
--branch $(BRANCH) --gpu $(NATIONAL_GPU) \
190+
--epochs $(NATIONAL_EPOCHS) \
191+
--push-results --national
192+
193+
calibrate-both:
194+
$(MAKE) calibrate-modal & $(MAKE) calibrate-modal-national & wait
195+
196+
stage-h5s:
197+
modal run modal_app/local_area.py::main \
198+
--branch $(BRANCH) --num-workers $(NUM_WORKERS) \
199+
$(if $(SKIP_DOWNLOAD),--skip-download)
200+
201+
stage-national-h5:
202+
modal run modal_app/local_area.py::main_national \
203+
--branch $(BRANCH)
204+
205+
stage-all-h5s:
206+
$(MAKE) stage-h5s & $(MAKE) stage-national-h5 & wait
207+
208+
promote:
209+
$(eval VERSION := $(or $(VERSION),$(shell python -c "import tomllib; print(tomllib.load(open('pyproject.toml','rb'))['project']['version'])")))
210+
modal run modal_app/local_area.py::main_promote \
211+
--branch $(BRANCH) --version $(VERSION)
212+
213+
validate-staging:
214+
python -m policyengine_us_data.calibration.validate_staging \
215+
--area-type states --output validation_results.csv
216+
217+
validate-staging-full:
218+
python -m policyengine_us_data.calibration.validate_staging \
219+
--area-type states,districts --output validation_results.csv
220+
221+
upload-validation:
222+
python -c "from policyengine_us_data.utils.huggingface import upload; \
223+
upload('validation_results.csv', \
224+
'policyengine/policyengine-us-data', \
225+
'calibration/logs/validation_results.csv')"
226+
227+
check-staging:
228+
python -m policyengine_us_data.calibration.check_staging_sums
229+
230+
check-sanity:
231+
python -m policyengine_us_data.calibration.validate_staging \
232+
--sanity-only --area-type states --areas NC
233+
234+
pipeline: data upload-dataset build-matrices calibrate-both stage-all-h5s
235+
@echo ""
236+
@echo "========================================"
237+
@echo "Pipeline complete. H5s are in HF staging."
238+
@echo "Run 'Promote Local Area H5 Files' workflow in GitHub to publish."
239+
@echo "========================================"
240+
107241
clean:
108242
rm -f policyengine_us_data/storage/*.h5
109243
rm -f policyengine_us_data/storage/*.db

‎changelog.d/add-database-build-test.added.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Add end-to-end test for calibration database build pipeline.

‎changelog.d/calibration-pipeline-improvements.added.md‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
Unified calibration pipeline with GPU-accelerated L1/L0 solver, target config YAML, and CLI package validator.
2+
Per-state and per-county precomputation replacing per-clone Microsimulation (51 sims instead of 436).
3+
Parallel state, county, and clone loop processing via ProcessPoolExecutor.
4+
Block-level takeup re-randomization with deterministic seeded draws.
5+
Hierarchical uprating with ACA PTC state-level CSV factors and CD reconciliation.
6+
Modal remote runner with Volume support, CUDA OOM fixes, and checkpointing.
7+
H5 builder that filters calibrated clone weights by CD subset, uses pre-assigned random census blocks from `geography.npz` to derive full sub-state geography, and produces self-contained local area datasets.
8+
Staging validation script (validate_staging.py) with sim.calculate() comparison and sanity checks.

‎changelog.d/calibration-pipeline-improvements.changed.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Geography assignment now prevents clone-to-CD collisions.
2+
County-dependent vars (aca_ptc) selectively precomputed per county; other vars use state-only path.
3+
Target config switched to finest-grain include mode (~18K targets).

‎changelog.d/calibration-pipeline-improvements.fixed.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
Cross-state cache pollution in matrix builder precomputation.
2+
Takeup draw ordering mismatch between matrix builder and stacked builder.
3+
At-large district geoid mismatch (7 districts had 0 estimates).

‎changelog.d/migrate-to-towncrier.changed.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Migrated from changelog_entry.yaml to towncrier fragments to eliminate merge conflicts.

0 commit comments

Comments
 (0)