fixing data build order from #674

Author: juaristi22

modal_app/data_build.py

@@ -493,49 +493,27 @@ def build_datasets(
             for future in as_completed(futures):
                 future.result()  # Raises if script failed
 
-        # GROUP 2: Depends on Group 1 - run in parallel
-        # cps.py needs acs, puf.py needs irs_puf + uprating
-        print("=== Phase 2: Building CPS and PUF (parallel) ===")
-        group2 = [
-            (
-                "policyengine_us_data/datasets/cps/cps.py",
-                SCRIPT_OUTPUTS["policyengine_us_data/datasets/cps/cps.py"],
-            ),
-            (
-                "policyengine_us_data/datasets/puf/puf.py",
-                SCRIPT_OUTPUTS["policyengine_us_data/datasets/puf/puf.py"],
-            ),
-        ]
-        with ThreadPoolExecutor(max_workers=2) as executor:
-            futures = {
-                executor.submit(
-                    run_script_with_checkpoint,
-                    script,
-                    output,
-                    branch,
-                    checkpoint_volume,
-                    env=env,
-                    log_file=log_file,
-                ): script
-                for script, output in group2
-            }
-            for future in as_completed(futures):
-                future.result()
-
-        # SEQUENTIAL: Extended CPS (needs both cps and puf)
-        print("=== Phase 3: Building extended CPS ===")
-        run_script_with_checkpoint(
+        # GROUP 2: Sequential chain — each step depends on the previous.
+        # cps.py needs acs; puf.py needs irs_puf + uprating + cps
+        # (pension imputation); extended_cps.py needs both cps and puf.
+        print("=== Phase 2: Building CPS → PUF → extended CPS ===")
+        for script in (
+            "policyengine_us_data/datasets/cps/cps.py",
+            "policyengine_us_data/datasets/puf/puf.py",
             "policyengine_us_data/datasets/cps/extended_cps.py",
-            SCRIPT_OUTPUTS["policyengine_us_data/datasets/cps/extended_cps.py"],
-            branch,
-            checkpoint_volume,
-            env=env,
-            log_file=log_file,
-        )
+        ):
+            run_script_with_checkpoint(
+                script,
+                SCRIPT_OUTPUTS[script],
+                branch,
+                checkpoint_volume,
+                env=env,
+                log_file=log_file,
+            )
 
         # GROUP 3: After extended_cps - run in parallel
         # enhanced_cps and stratified_cps both depend on extended_cps
-        print("=== Phase 4: Building enhanced and stratified CPS (parallel) ===")
+        print("=== Phase 3: Building enhanced and stratified CPS (parallel) ===")
         phase4_futures = []
         with ThreadPoolExecutor(max_workers=2) as executor:
             if not skip_enhanced_cps:
@@ -570,11 +548,11 @@ def build_datasets(
             for future in as_completed(phase4_futures):
                 future.result()
 
-        # GROUP 4: After Phase 4 - run in parallel
+        # GROUP 4: After Phase 3 - run in parallel
         # create_source_imputed_cps needs stratified_cps
         # small_enhanced_cps needs enhanced_cps
         print(
-            "=== Phase 5: Building source imputed CPS "
+            "=== Phase 4: Building source imputed CPS "
             "and small enhanced CPS (parallel) ==="
         )
         phase5_futures = []

policyengine_us_data/datasets/org/org.py

@@ -1,3 +1,4 @@
+import logging
 from functools import lru_cache
 
 from microimpute.models.qrf import QRF
@@ -392,7 +393,13 @@ def load_org_training_data() -> pd.DataFrame:
     """Load ORG donor rows built from official CPS basic monthly files."""
     cache_path = STORAGE_FOLDER / ORG_FILENAME
     if cache_path.exists():
-        return pd.read_csv(cache_path)
+        try:
+            return pd.read_csv(cache_path)
+        except (EOFError, pd.errors.ParserError):
+            logging.warning(
+                "Corrupt ORG cache %s — deleting and rebuilding", cache_path
+            )
+            cache_path.unlink()
 
     months = []
     for month in ORG_MONTHS:

policyengine_us_data/datasets/puf/puf.py

@@ -167,14 +167,7 @@ def impute_pension_contributions_to_puf(puf_df):
     from policyengine_us import Microsimulation
     from policyengine_us_data.datasets.cps import CPS_2024
 
-    # CPS_2024 may not exist yet during parallel CI builds.
-    # Fall back to CPS_2021 release artifact if needed.
-    try:
-        cps = Microsimulation(dataset=CPS_2024)
-    except Exception:
-        from policyengine_us_data.datasets.cps import CPS_2021
-
-        cps = Microsimulation(dataset=CPS_2021)
+    cps = Microsimulation(dataset=CPS_2024)
     cps.subsample(10_000)
 
     predictors = [