Fix build pipeline: add missing script, remove geography.npz, input-scoped checkpoints

- Add create_source_imputed_cps.py to data_build.py Phase 5 (was skipped in CI)
- Remove geography.npz dependency from Modal pipeline; workers regenerate
  geography deterministically from (n_records, n_clones, seed)
- Add input-scoped checkpoints to publish_local_area.py: hash weights+dataset
  to auto-clear stale checkpoints when inputs change
- Remove stale artifacts from push-to-modal (stacked_blocks, stacked_takeup,
  geo_labels)
- Stop uploading source_imputed H5 as intermediate; promote-dataset uploads
  at promotion time instead
- Default skip_download=True in Modal local_area (reads from volume)
- Remove _upload_source_imputed from remote_calibration_runner
- Clean up huggingface.py: remove geography/blocks/geo_labels from
  download and upload functions
- ruff format

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: baogorek

Makefile

@@ -87,9 +87,11 @@ promote-database:
 	@echo "Copied DB and raw_inputs to HF clone. Now cd to HF repo, commit, and push."
 
 promote-dataset:
-	cp policyengine_us_data/storage/source_imputed_stratified_extended_cps_2024.h5 \
-		$(HF_CLONE_DIR)/calibration/source_imputed_stratified_extended_cps.h5
-	@echo "Copied dataset to HF clone. Now cd to HF repo, commit, and push."
+	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
 	python policyengine_us_data/utils/uprating.py
@@ -141,11 +143,9 @@ upload-calibration:
 		upload_calibration_artifacts()"
 
 upload-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 uploaded to HF."
+	@echo "NOTE: source_imputed H5 is an intermediate artifact."
+	@echo "Use 'make push-to-modal' to push to Modal volume,"
+	@echo "or 'make promote-dataset' to publish to HF at promotion time."
 
 upload-database:
 	python -c "from policyengine_us_data.utils.huggingface import upload; \
@@ -158,18 +158,9 @@ push-to-modal:
 	modal volume put local-area-staging \
 		policyengine_us_data/storage/calibration/calibration_weights.npy \
 		calibration_inputs/calibration/calibration_weights.npy --force
-	modal volume put local-area-staging \
-		policyengine_us_data/storage/calibration/stacked_blocks.npy \
-		calibration_inputs/calibration/stacked_blocks.npy --force
-	modal volume put local-area-staging \
-		policyengine_us_data/storage/calibration/stacked_takeup.npz \
-		calibration_inputs/calibration/stacked_takeup.npz --force
 	modal volume put local-area-staging \
 		policyengine_us_data/storage/calibration/policy_data.db \
 		calibration_inputs/calibration/policy_data.db --force
-	modal volume put local-area-staging \
-		policyengine_us_data/storage/calibration/geo_labels.json \
-		calibration_inputs/calibration/geo_labels.json --force
 	modal volume put local-area-staging \
 		policyengine_us_data/storage/source_imputed_stratified_extended_cps_2024.h5 \
 		calibration_inputs/calibration/source_imputed_stratified_extended_cps.h5 --force
@@ -195,8 +186,7 @@ calibrate-both:
 
 stage-h5s:
 	modal run modal_app/local_area.py::main \
-		--branch $(BRANCH) --num-workers $(NUM_WORKERS) \
-		$(if $(SKIP_DOWNLOAD),--skip-download)
+		--branch $(BRANCH) --num-workers $(NUM_WORKERS)
 
 stage-national-h5:
 	modal run modal_app/local_area.py::main_national \
@@ -231,7 +221,7 @@ check-sanity:
 	python -m policyengine_us_data.calibration.validate_staging \
 		--sanity-only --area-type states --areas NC
 
-pipeline: data upload-dataset build-matrices calibrate-both stage-all-h5s
+pipeline: data push-to-modal build-matrices calibrate-both stage-all-h5s
 	@echo ""
 	@echo "========================================"
 	@echo "Pipeline complete. H5s are in HF staging."

modal_app/README.md

@@ -76,7 +76,6 @@ Every run produces these local files (whichever the calibration script emits):
 - **unified_diagnostics.csv** — Final per-target diagnostics
 - **calibration_log.csv** — Per-target metrics across epochs (requires `--log-freq`)
 - **unified_run_config.json** — Run configuration and summary stats
-- **stacked_blocks.npy** — Census block assignments for stacked records
 
 ## Artifact Upload to HuggingFace
 
@@ -86,7 +85,6 @@ atomic commit after writing them locally:
 | Local file | HF path |
 |------------|---------|
 | `calibration_weights.npy` | `calibration/calibration_weights.npy` |
-| `stacked_blocks.npy` | `calibration/stacked_blocks.npy` |
 | `calibration_log.csv` | `calibration/logs/calibration_log.csv` |
 | `unified_diagnostics.csv` | `calibration/logs/unified_diagnostics.csv` |
 | `unified_run_config.json` | `calibration/logs/unified_run_config.json` |
@@ -201,7 +199,6 @@ Artifacts uploaded to HF by `--push-results`:
 | Local file | HF path |
 |------------|---------|
 | `calibration_weights.npy` | `calibration/calibration_weights.npy` |
-| `stacked_blocks.npy` | `calibration/stacked_blocks.npy` |
 | `calibration_log.csv` | `calibration/logs/calibration_log.csv` |
 | `unified_diagnostics.csv` | `calibration/logs/unified_diagnostics.csv` |
 | `unified_run_config.json` | `calibration/logs/unified_run_config.json` |

modal_app/data_build.py

@@ -433,15 +433,38 @@ def build_datasets(
             for future in as_completed(futures):
                 future.result()
 
-        # SEQUENTIAL: Small enhanced CPS (needs enhanced_cps)
-        print("=== Phase 5: Building small enhanced CPS ===")
-        run_script_with_checkpoint(
-            "policyengine_us_data/datasets/cps/small_enhanced_cps.py",
-            SCRIPT_OUTPUTS["policyengine_us_data/datasets/cps/small_enhanced_cps.py"],
-            branch,
-            checkpoint_volume,
-            env=env,
+        # GROUP 4: After Phase 4 - run in parallel
+        # create_source_imputed_cps needs stratified_cps
+        # small_enhanced_cps needs enhanced_cps
+        print(
+            "=== Phase 5: Building source imputed CPS "
+            "and small enhanced CPS (parallel) ==="
         )
+        with ThreadPoolExecutor(max_workers=2) as executor:
+            futures = [
+                executor.submit(
+                    run_script_with_checkpoint,
+                    "policyengine_us_data/calibration/create_source_imputed_cps.py",
+                    SCRIPT_OUTPUTS[
+                        "policyengine_us_data/calibration/create_source_imputed_cps.py"
+                    ],
+                    branch,
+                    checkpoint_volume,
+                    env=env,
+                ),
+                executor.submit(
+                    run_script_with_checkpoint,
+                    "policyengine_us_data/datasets/cps/small_enhanced_cps.py",
+                    SCRIPT_OUTPUTS[
+                        "policyengine_us_data/datasets/cps/small_enhanced_cps.py"
+                    ],
+                    branch,
+                    checkpoint_volume,
+                    env=env,
+                ),
+            ]
+            for future in as_completed(futures):
+                future.result()
 
     # Run tests with checkpointing
     print("=== Running tests with checkpointing ===")

modal_app/local_area.py

@@ -313,17 +313,10 @@ def build_areas_worker(
         "--output-dir",
         str(output_dir),
     ]
-    if "geography" not in calibration_inputs:
-        raise RuntimeError(
-            "geography.npz path missing from calibration_inputs. "
-            "Re-run calibration to generate this artifact."
-        )
-    worker_cmd.extend(
-        [
-            "--geography-path",
-            calibration_inputs["geography"],
-        ]
-    )
+    if "n_clones" in calibration_inputs:
+        worker_cmd.extend(["--n-clones", str(calibration_inputs["n_clones"])])
+    if "seed" in calibration_inputs:
+        worker_cmd.extend(["--seed", str(calibration_inputs["seed"])])
     result = subprocess.run(
         worker_cmd,
         capture_output=True,
@@ -575,7 +568,7 @@ def coordinate_publish(
     branch: str = "main",
     num_workers: int = 8,
     skip_upload: bool = False,
-    skip_download: bool = False,
+    skip_download: bool = True,
 ) -> str:
     """Coordinate the full publishing workflow."""
     setup_gcp_credentials()
@@ -612,12 +605,12 @@ def coordinate_publish(
             "weights": weights_path,
             "dataset": dataset_path,
             "database": db_path,
-            "geography": (calibration_dir / "calibration" / "geography.npz"),
-            "run_config": (calibration_dir / "calibration" / "unified_run_config.json"),
         }
         for label, p in required.items():
             if not p.exists():
-                raise RuntimeError(f"Missing required calibration input ({label}): {p}")
+                raise RuntimeError(
+                    f"Missing required calibration input ({label}): {p}"
+                )
         print("All required calibration inputs found on volume.")
     else:
         if calibration_dir.exists():
@@ -649,20 +642,14 @@ def coordinate_publish(
         calibration_dir / "calibration" / "source_imputed_stratified_extended_cps.h5"
     )
 
-    geo_npz_path = calibration_dir / "calibration" / "geography.npz"
     config_json_path = calibration_dir / "calibration" / "unified_run_config.json"
     calibration_inputs = {
         "weights": str(weights_path),
         "dataset": str(dataset_path),
         "database": str(db_path),
+        "n_clones": 430,
+        "seed": 42,
     }
-    if not geo_npz_path.exists():
-        raise RuntimeError(
-            f"geography.npz not found at {geo_npz_path}. "
-            f"Re-run calibration to generate this artifact."
-        )
-    calibration_inputs["geography"] = str(geo_npz_path)
-    print(f"Geography artifact found: {geo_npz_path}")
     validate_artifacts(
         config_json_path,
         calibration_dir / "calibration",
@@ -793,7 +780,7 @@ def main(
     branch: str = "main",
     num_workers: int = 8,
     skip_upload: bool = False,
-    skip_download: bool = False,
+    skip_download: bool = True,
 ):
     """Local entrypoint for Modal CLI."""
     result = coordinate_publish.remote(
@@ -859,23 +846,16 @@ def coordinate_national_publish(
         calibration_dir / "calibration" / "source_imputed_stratified_extended_cps.h5"
     )
 
-    geo_npz_path = calibration_dir / "calibration" / "national_geography.npz"
     config_json_path = (
         calibration_dir / "calibration" / "national_unified_run_config.json"
     )
     calibration_inputs = {
         "weights": str(weights_path),
         "dataset": str(dataset_path),
         "database": str(db_path),
+        "n_clones": 430,
+        "seed": 42,
     }
-    if not geo_npz_path.exists():
-        raise RuntimeError(
-            f"national_geography.npz not found at "
-            f"{geo_npz_path}. Re-run national calibration "
-            f"to generate this artifact."
-        )
-    calibration_inputs["geography"] = str(geo_npz_path)
-    print(f"National geography artifact found: {geo_npz_path}")
     validate_artifacts(
         config_json_path,
         calibration_dir / "calibration",