Merge pull request #700 from PolicyEngine/fix/build-database-from-source

Always build policy_data.db from source.

Both Maria and I worked on this and Anthony gave me his blessing via Slack. Given this is my last week and things need to move, I am bypassing the reviewer check mark and merging.

Author: baogorek

Makefile

@@ -1,8 +1,10 @@
-.PHONY: all format test test-unit test-integration install download upload docker documentation data validate-data calibrate calibrate-build publish-local-area upload-calibration upload-dataset upload-database push-to-modal build-data-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 refresh-soi-targets push-pr-branch
+.PHONY: all format test test-unit test-integration install download upload docker documentation data validate-data calibrate calibrate-build publish-local-area upload-calibration upload-dataset push-to-modal build-data-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-dataset promote build-h5s validate-local refresh-soi-targets push-pr-branch
 
 SOI_SOURCE_YEAR ?= 2021
 SOI_TARGET_YEAR ?= 2023
 
+YEAR ?= 2024
+
 GPU ?= T4
 EPOCHS ?= 1000
 NATIONAL_GPU ?= T4
@@ -75,38 +77,29 @@ documentation-dev:
 database:
 	rm -f policyengine_us_data/storage/calibration/policy_data.db
 	python policyengine_us_data/db/create_database_tables.py
-	python policyengine_us_data/db/create_initial_strata.py
-	python policyengine_us_data/db/etl_national_targets.py
-	python policyengine_us_data/db/etl_age.py
-	python policyengine_us_data/db/etl_medicaid.py
-	python policyengine_us_data/db/etl_snap.py
-	python policyengine_us_data/db/etl_state_income_tax.py
-	python policyengine_us_data/db/etl_irs_soi.py
-	python policyengine_us_data/db/etl_pregnancy.py
+	python policyengine_us_data/db/create_initial_strata.py --year $(YEAR)
+	python policyengine_us_data/db/etl_national_targets.py --year $(YEAR)
+	python policyengine_us_data/db/etl_age.py --year $(YEAR)
+	python policyengine_us_data/db/etl_medicaid.py --year $(YEAR)
+	python policyengine_us_data/db/etl_snap.py --year $(YEAR)
+	python policyengine_us_data/db/etl_state_income_tax.py --year $(YEAR)
+	python policyengine_us_data/db/etl_irs_soi.py --year $(YEAR)
+	python policyengine_us_data/db/etl_pregnancy.py --year $(YEAR)
 	python policyengine_us_data/db/validate_database.py
 
 database-refresh:
 	rm -f policyengine_us_data/storage/calibration/policy_data.db
 	rm -rf policyengine_us_data/storage/calibration/raw_inputs/
 	$(MAKE) database
 
-promote-database:
-	sqlite3 policyengine_us_data/storage/calibration/policy_data.db "PRAGMA wal_checkpoint(TRUNCATE);"
-	cp policyengine_us_data/storage/calibration/policy_data.db \
-		$(HF_CLONE_DIR)/calibration/policy_data.db
-	rm -rf $(HF_CLONE_DIR)/calibration/raw_inputs
-	cp -r policyengine_us_data/storage/calibration/raw_inputs \
-		$(HF_CLONE_DIR)/calibration/raw_inputs
-	@echo "Copied DB and raw_inputs to HF clone. Now cd to HF repo, commit, and push."
-
 promote-dataset:
 	python -c "from policyengine_us_data.utils.huggingface import upload; \
 		upload('policyengine_us_data/storage/source_imputed_stratified_extended_cps_2024.h5', \
 		'policyengine/policyengine-us-data', \
 		'calibration/source_imputed_stratified_extended_cps.h5')"
 	@echo "Dataset promoted to HF."
 
-data: download
+data: download database
 	python policyengine_us_data/utils/uprating.py
 	python policyengine_us_data/datasets/acs/acs.py
 	python policyengine_us_data/datasets/cps/cps.py
@@ -174,13 +167,6 @@ upload-dataset:
 		'calibration/source_imputed_stratified_extended_cps.h5')"
 	@echo "Dataset uploaded to HF."
 
-upload-database:
-	python -c "from policyengine_us_data.utils.huggingface import upload; \
-		upload('policyengine_us_data/storage/calibration/policy_data.db', \
-		'policyengine/policyengine-us-data', \
-		'calibration/policy_data.db')"
-	@echo "Database uploaded to HF."
-
 push-to-modal:
 	modal volume put pipeline-artifacts \
 		policyengine_us_data/storage/calibration/calibration_weights.npy \

changelog.d/fix-build-database-from-source.changed.md

@@ -0,0 +1 @@
+Build policy_data.db from source instead of downloading from HuggingFace, replace H5 dataset dependency with a --year CLI flag for all database ETL scripts, fix Modal data build ordering (CPS before PUF), and add missing heapq import in local area builder.

modal_app/data_build.py

@@ -435,7 +435,14 @@ def build_datasets(
         env=env,
         log_file=log_file,
     )
-    # Checkpoint policy_data.db immediately after download so it survives
+    # Build policy_data.db from source
+    subprocess.run(
+        ["uv", "run", "make", "database"],
+        check=True,
+        cwd="/root/policyengine-us-data",
+        env=env,
+    )
+    # Checkpoint policy_data.db immediately after build so it survives
     # test failures and can be restored on retries.
     save_checkpoint(
         branch,

modal_app/images.py

@@ -52,7 +52,7 @@ def _base_image(extras: list[str] | None = None):
     extra_flags = " ".join(f"--extra {e}" for e in (extras or []))
     return (
         modal.Image.debian_slim(python_version="3.14")
-        .apt_install("git")
+        .apt_install("git", "make")
         .pip_install("uv>=0.8")
         .add_local_dir(
             str(REPO_ROOT),

modal_app/local_area.py

@@ -11,6 +11,7 @@
     modal run modal_app/local_area.py --branch=main --num-workers=8
 """
 
+import heapq
 import json
 import os
 import subprocess

modal_app/pipeline.py

@@ -880,11 +880,12 @@ def run_pipeline(
             print("\n[Step 3/5] Fit weights (skipped - completed)")
 
         # ── Step 4: Build H5s + stage + diagnostics (parallel) ──
-        # Per plan: all four tasks run in parallel:
         #   4a. coordinate_publish (regional H5s)
         #   4b. coordinate_national_publish (national H5)
         #   4c. stage_base_datasets (datasets → HF staging)
-        #   4d. upload_run_diagnostics (diagnostics → HF)
+        #   4d. upload_run_diagnostics (calibration diagnostics → HF)
+        #   4e. _write_validation_diagnostics (after H5 builds)
+        #   4f. upload_run_diagnostics (validation diagnostics → HF)
         if not _step_completed(meta, "publish_and_stage"):
             print(
                 "\n[Step 4/5] Building H5s, staging datasets, "
@@ -918,16 +919,12 @@ def run_pipeline(
                     f"    → coordinate_national_publish fc: {national_h5_handle.object_id}"
                 )
 
-            # While H5 builds run, stage base datasets
-            # and upload diagnostics in this container
+            # While H5 builds run, stage base datasets in this container
             pipeline_volume.reload()
 
             print("  Staging base datasets to HF...")
             stage_base_datasets(run_id, version, branch)
 
-            print("  Uploading run diagnostics...")
-            upload_run_diagnostics(run_id, branch)
-
             # Now wait for H5 builds to finish
             print("  Waiting for regional H5 build...")
             regional_h5_result = regional_h5_handle.get()
@@ -964,6 +961,10 @@ def run_pipeline(
                 vol=pipeline_volume,
             )
 
+            # Upload validation diagnostics (written after H5 builds)
+            print("  Uploading validation diagnostics...")
+            upload_run_diagnostics(run_id, branch)
+
             _record_step(
                 meta,
                 "publish_and_stage",

policyengine_us_data/db/DATABASE_GUIDE.md

@@ -15,7 +15,6 @@ cd ~/devl/sep/policyengine-us-data
 
 make database           # Build (uses cached downloads if available)
 make database-refresh   # Force re-download all sources and rebuild
-make promote-database   # Copy DB + raw inputs to HuggingFace clone
 ```
 
 ### Pipeline Stages
@@ -44,19 +43,6 @@ Set `PE_REFRESH_RAW=1` to force re-download:
 PE_REFRESH_RAW=1 make database
 ```
 
-### Promotion to HuggingFace
-
-After building and validating:
-```bash
-make promote-database
-cd ~/devl/huggingface/policyengine-us-data
-git add calibration/policy_data.db calibration/raw_inputs/
-git commit -m "Update policy_data.db - <description>"
-git push
-```
-
-This copies both the database and the raw inputs that built it, preserving provenance in the HF repo's git history.
-
 ### Recovery
 
 If a step fails mid-pipeline, delete the database and re-run. With cached downloads this takes ~10-15 minutes:
@@ -286,4 +272,4 @@ ORDER BY geographic_id;
 
 `policyengine_us_data/storage/calibration/policy_data.db`
 
-Downloaded from HuggingFace by `download_private_prerequisites.py` and `download_calibration_inputs()` in `utils/huggingface.py`.
+Built from source via `make database`. See [Building the Database](#building-the-database) above.

policyengine_us_data/db/etl_national_targets.py

@@ -13,20 +13,19 @@
     RETIREMENT_CONTRIBUTION_TARGETS,
 )
 from policyengine_us_data.utils.db import (
-    DEFAULT_DATASET,
+    DEFAULT_YEAR,
     etl_argparser,
 )
 
 
-def extract_national_targets(dataset: str = DEFAULT_DATASET):
+def extract_national_targets(year: int = DEFAULT_YEAR):
     """
     Extract national calibration targets from various sources.
 
     Parameters
     ----------
-    dataset : str
-        Path to the calibration dataset (local path or HuggingFace URL).
-        The time period is derived from the dataset's default_calculation_period.
+    year : int
+        Target year for calibration data.
 
     Returns
     -------
@@ -38,15 +37,14 @@ def extract_national_targets(dataset: str = DEFAULT_DATASET):
         - conditional_count_targets: Enrollment counts requiring constraints
         - cbo_targets: List of CBO projection targets
         - treasury_targets: List of Treasury/JCT targets
-        - time_period: The year derived from the dataset
+        - time_period: The target year
     """
-    from policyengine_us import Microsimulation
+    from policyengine_us import CountryTaxBenefitSystem
 
-    print(f"Loading dataset: {dataset}")
-    sim = Microsimulation(dataset=dataset)
+    time_period = year
+    print(f"Using time_period: {time_period}")
 
-    time_period = int(sim.default_calculation_period)
-    print(f"Derived time_period from dataset: {time_period}")
+    tax_benefit_system = CountryTaxBenefitSystem()
 
     # Hardcoded dollar targets are specific to 2024 and should be
     # labeled as such.  Only CBO/Treasury parameter lookups use the
@@ -400,7 +398,7 @@ def extract_national_targets(dataset: str = DEFAULT_DATASET):
     for variable_name in cbo_vars:
         param_name = cbo_param_name_map.get(variable_name, variable_name)
         try:
-            value = sim.tax_benefit_system.parameters(
+            value = tax_benefit_system.parameters(
                 time_period
             ).calibration.gov.cbo._children[param_name]
             cbo_targets.append(
@@ -420,7 +418,7 @@ def extract_national_targets(dataset: str = DEFAULT_DATASET):
 
     # Treasury/JCT targets (EITC) - use time_period derived from dataset
     try:
-        eitc_value = sim.tax_benefit_system.parameters.calibration.gov.treasury.tax_expenditures.eitc(
+        eitc_value = tax_benefit_system.parameters.calibration.gov.treasury.tax_expenditures.eitc(
             time_period
         )
         treasury_targets = [
@@ -883,11 +881,11 @@ def load_national_targets(
 
 def main():
     """Main ETL pipeline for national targets."""
-    args, _ = etl_argparser("ETL for national calibration targets")
+    _, year = etl_argparser("ETL for national calibration targets")
 
     # Extract
     print("Extracting national targets...")
-    raw_targets = extract_national_targets(dataset=args.dataset)
+    raw_targets = extract_national_targets(year=year)
     time_period = raw_targets["time_period"]
     print(f"Using time_period={time_period} for CBO/Treasury targets")
 

policyengine_us_data/storage/download_private_prerequisites.py

@@ -1,9 +1,5 @@
-import os
-
 from pathlib import Path
-from policyengine_us_data.db.create_database_tables import (
-    refresh_views_for_db_path,
-)
+
 from policyengine_us_data.utils.huggingface import download
 
 FOLDER = Path(__file__).parent
@@ -26,16 +22,3 @@
     local_folder=FOLDER,
     version=None,
 )
-if os.environ.get("SKIP_POLICY_DB_DOWNLOAD"):
-    print(
-        "SKIP_POLICY_DB_DOWNLOAD set — skipping "
-        "policy_data.db download from HuggingFace"
-    )
-else:
-    download(
-        repo="policyengine/policyengine-us-data",
-        repo_filename="calibration/policy_data.db",
-        local_folder=FOLDER,
-        version=None,
-    )
-    refresh_views_for_db_path(FOLDER / "calibration" / "policy_data.db")

policyengine_us_data/utils/db.py

@@ -1,5 +1,4 @@
 import argparse
-from pathlib import Path
 from typing import Dict, List, Optional, Tuple
 
 from sqlmodel import Session, select
@@ -9,54 +8,37 @@
     Stratum,
     StratumConstraint,
 )
-from policyengine_us_data.storage import STORAGE_FOLDER
 
-DEFAULT_DATASET = str(STORAGE_FOLDER / "source_imputed_stratified_extended_cps_2024.h5")
+DEFAULT_YEAR = 2024
 
 
 def etl_argparser(
     description: str,
     extra_args_fn=None,
 ) -> Tuple[argparse.Namespace, int]:
-    """Shared argument parsing and dataset-year derivation for ETL scripts.
+    """Shared argument parsing for ETL scripts.
 
     Args:
         description: Description for the argparse help text.
         extra_args_fn: Optional callable that receives the parser to add
             extra arguments before parsing.
 
     Returns:
-        (args, year) where *year* is derived from the dataset's
-        ``default_calculation_period``.
+        (args, year) tuple.
     """
     parser = argparse.ArgumentParser(description=description)
     parser.add_argument(
-        "--dataset",
-        default=DEFAULT_DATASET,
-        help=(
-            "Source dataset (local path or HuggingFace URL). "
-            "The year is derived from the dataset's "
-            "default_calculation_period. Default: %(default)s"
-        ),
+        "--year",
+        type=int,
+        default=DEFAULT_YEAR,
+        help="Target year for calibration data. Default: %(default)s",
     )
     if extra_args_fn is not None:
         extra_args_fn(parser)
 
     args = parser.parse_args()
-
-    if not args.dataset.startswith("hf://") and not Path(args.dataset).exists():
-        raise FileNotFoundError(
-            f"Dataset not found: {args.dataset}\n"
-            f"Either build it locally (`make data`) or pass a "
-            f"HuggingFace URL via --dataset hf://policyengine/..."
-        )
-
-    from policyengine_us import Microsimulation
-
-    print(f"Loading dataset: {args.dataset}")
-    sim = Microsimulation(dataset=args.dataset)
-    year = int(sim.default_calculation_period)
-    print(f"Derived year from dataset: {year}")
+    year = args.year
+    print(f"Using year: {year}")
 
     return args, year