Harden CPS basic ORG loading and caching (#727)

* Harden CPS basic ORG loading

* Validate cached ORG donor data

* Format ORG cache validation helper

* Fix CPS basic ORG column alignment

Author: MaxGhenis

changelog.d/fix-org-loader.fixed.md

@@ -0,0 +1 @@
+Harden CPS basic ORG donor loading against transient fetch failures and concurrent cache builds.

policyengine_us_data/datasets/org/org.py

@@ -6,11 +6,16 @@
 imputation onto CPS records.
 """
 
+from contextlib import contextmanager
 from functools import lru_cache
+from io import BytesIO
+from pathlib import Path
+import fcntl
 
 from microimpute.models.qrf import QRF
 import numpy as np
 import pandas as pd
+import requests
 
 from policyengine_us_data.storage import STORAGE_FOLDER
 
@@ -181,11 +186,13 @@ def _cps_basic_org_month_url(year: int, month: str) -> str:
     )
 
 
-def _select_cps_basic_org_columns(month_df: pd.DataFrame) -> pd.DataFrame:
-    """Normalize CPS basic-month columns onto the ORG schema."""
+def _resolve_cps_basic_org_column_names(
+    columns: pd.Index | list[str],
+) -> list[str]:
+    """Resolve CPS basic-month columns onto the expected ORG schema order."""
     column_lookup = {
-        str(column).lower(): column
-        for column in month_df.columns
+        str(column).lower(): str(column)
+        for column in columns
         if isinstance(column, str)
     }
     missing = [
@@ -196,9 +203,12 @@ def _select_cps_basic_org_columns(month_df: pd.DataFrame) -> pd.DataFrame:
     if missing:
         raise ValueError(f"CPS basic ORG month is missing required columns: {missing}")
 
-    selected = month_df[
-        [column_lookup[column.lower()] for column in CPS_BASIC_MONTHLY_ORG_COLUMNS]
-    ].copy()
+    return [column_lookup[column.lower()] for column in CPS_BASIC_MONTHLY_ORG_COLUMNS]
+
+
+def _select_cps_basic_org_columns(month_df: pd.DataFrame) -> pd.DataFrame:
+    """Normalize CPS basic-month columns onto the ORG schema."""
+    selected = month_df[_resolve_cps_basic_org_column_names(month_df.columns)].copy()
     selected.columns = CPS_BASIC_MONTHLY_ORG_COLUMNS
     return selected
 
@@ -211,16 +221,18 @@ def _load_cps_basic_org_month(
 ) -> pd.DataFrame:
     """Load one CPS basic-month file with light retry around transient fetch/parser issues."""
     url = _cps_basic_org_month_url(year, month)
-    required_columns = {column.lower() for column in CPS_BASIC_MONTHLY_ORG_COLUMNS}
     last_error: Exception | None = None
 
     for _ in range(max_attempts):
         try:
+            response = requests.get(url, timeout=60)
+            response.raise_for_status()
+            content = response.content
+            header = pd.read_csv(BytesIO(content), nrows=0)
+            selected_columns = _resolve_cps_basic_org_column_names(header.columns)
             month_df = pd.read_csv(
-                url,
-                usecols=lambda column: (
-                    isinstance(column, str) and column.lower() in required_columns
-                ),
+                BytesIO(content),
+                usecols=selected_columns,
                 low_memory=False,
             )
             return _select_cps_basic_org_columns(month_df)
@@ -233,6 +245,36 @@ def _load_cps_basic_org_month(
     ) from last_error
 
 
+@contextmanager
+def _org_cache_build_lock(lock_path: Path):
+    lock_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(lock_path, "w") as lock_file:
+        fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
+        try:
+            yield
+        finally:
+            fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
+
+
+def _load_valid_cached_org_training_data(cache_path: Path) -> pd.DataFrame | None:
+    """Return a cached ORG training frame when it is present and structurally valid."""
+    required_columns = set(
+        ORG_PREDICTORS + ORG_QRF_IMPUTED_VARIABLES + ["sample_weight"]
+    )
+    try:
+        cached = pd.read_csv(cache_path)
+    except (FileNotFoundError, OSError, pd.errors.EmptyDataError):
+        return None
+
+    if cached.empty:
+        return None
+
+    if not required_columns.issubset(cached.columns):
+        return None
+
+    return cached
+
+
 def _transform_cps_basic_org_month(month_df: pd.DataFrame) -> pd.DataFrame:
     """Convert one monthly CPS basic file into ORG donor rows.
 
@@ -451,17 +493,31 @@ def _predict_union_coverage_from_bls_tables(
 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)
-
-    months = []
-    for month in ORG_MONTHS:
-        month_df = _load_cps_basic_org_month(ORG_YEAR, month)
-        months.append(_transform_cps_basic_org_month(month_df))
-
-    org = pd.concat(months, ignore_index=True)
-    org.to_csv(cache_path, index=False, compression="gzip")
-    return org
+    lock_path = cache_path.parent / f"{cache_path.name}.lock"
+    cached = _load_valid_cached_org_training_data(cache_path)
+    if cached is not None:
+        return cached
+
+    with _org_cache_build_lock(lock_path):
+        cached = _load_valid_cached_org_training_data(cache_path)
+        if cached is not None:
+            return cached
+        if cache_path.exists():
+            cache_path.unlink()
+
+        months = []
+        for month in ORG_MONTHS:
+            month_df = _load_cps_basic_org_month(ORG_YEAR, month)
+            months.append(_transform_cps_basic_org_month(month_df))
+
+        org = pd.concat(months, ignore_index=True)
+        temp_path = cache_path.parent / f"{cache_path.name}.tmp.gz"
+        org.to_csv(temp_path, index=False, compression="gzip")
+        temp_path.replace(cache_path)
+        cached = _load_valid_cached_org_training_data(cache_path)
+        if cached is None:
+            raise ValueError("Failed to build a valid cached ORG donor file")
+        return cached
 
 
 @lru_cache(maxsize=1)

tests/unit/datasets/test_org.py

@@ -1,5 +1,7 @@
 import numpy as np
 import pandas as pd
+from concurrent.futures import ThreadPoolExecutor
+import time
 
 from policyengine_us_data.datasets.cps import cps as cps_module
 from policyengine_us_data.datasets.org import (
@@ -10,6 +12,7 @@
     CPS_BASIC_MONTHLY_ORG_COLUMNS,
     _build_union_priority_weights,
     _load_cps_basic_org_month,
+    load_org_training_data,
     _predict_union_coverage_from_bls_tables,
     _select_cps_basic_org_columns,
     _transform_cps_basic_org_month,
@@ -154,42 +157,199 @@ def test_load_cps_basic_org_month_retries_after_transient_parser_failure(
     monkeypatch,
 ):
     calls = []
-    month_df = pd.DataFrame(
-        {
-            "hrmis": [4],
-            "GESTFIPS": [6],
-            "PRTAGE": [30],
-            "PESEX": [2],
-            "PTDTRACE": [1],
-            "PEHSPNON": [2],
-            "PWORWGT": [100.0],
-            "PTERNWA": [100000.0],
-            "PTERNHLY": [2500.0],
-            "PEERNHRY": [1],
-            "PEHRUSLT": [40.0],
-            "PRERELG": [1],
-            "PEMLR": [1],
-            "PEIO1COW": [1],
-        }
+    csv_text = (
+        "hrmis,GESTFIPS,PRTAGE,PESEX,PTDTRACE,PEHSPNON,PWORWGT,"
+        "PTERNWA,PTERNHLY,PEERNHRY,PEHRUSLT,PRERELG,PEMLR,PEIO1COW\n"
+        "4,6,30,2,1,2,100.0,100000.0,2500.0,1,40.0,1,1,1\n"
     )
 
-    def fake_read_csv(*args, **kwargs):
+    class FakeResponse:
+        def __init__(self, text: str, status_code: int = 200):
+            self.content = text.encode("utf-8")
+            self.status_code = status_code
+
+        def raise_for_status(self):
+            if self.status_code >= 400:
+                raise ValueError("bad status")
+
+    responses = [
+        FakeResponse("<html>temporary error</html>"),
+        FakeResponse(csv_text),
+    ]
+
+    def fake_get(*args, **kwargs):
         calls.append(kwargs)
-        if len(calls) == 1:
-            raise ValueError("Usecols do not match columns")
-        return month_df
+        return responses.pop(0)
 
-    monkeypatch.setattr(
-        "policyengine_us_data.datasets.org.org.pd.read_csv", fake_read_csv
-    )
+    monkeypatch.setattr("policyengine_us_data.datasets.org.org.requests.get", fake_get)
 
     loaded = _load_cps_basic_org_month(2024, "may", max_attempts=2)
 
     assert len(calls) == 2
-    assert callable(calls[0]["usecols"])
     assert loaded.columns.tolist() == CPS_BASIC_MONTHLY_ORG_COLUMNS
 
 
+def test_load_cps_basic_org_month_reorders_file_order_columns(monkeypatch):
+    csv_text = (
+        "PTERNWA,PEHRUSLT,hrmis,PEMLR,PEERNHRY,PEHSPNON,PRTAGE,"
+        "PTDTRACE,pworwgt,peio1cow,GESTFIPS,PESEX,PTERNHLY,PRERELG\n"
+        "100000.0,40.0,4,1,1,2,30,1,100.0,1,6,2,2500.0,1\n"
+    )
+
+    class FakeResponse:
+        def __init__(self, text: str):
+            self.content = text.encode("utf-8")
+
+        def raise_for_status(self):
+            return None
+
+    monkeypatch.setattr(
+        "policyengine_us_data.datasets.org.org.requests.get",
+        lambda *args, **kwargs: FakeResponse(csv_text),
+    )
+
+    loaded = _load_cps_basic_org_month(2024, "may", max_attempts=1)
+
+    assert loaded.columns.tolist() == CPS_BASIC_MONTHLY_ORG_COLUMNS
+    assert loaded.iloc[0].to_dict() == {
+        "HRMIS": 4,
+        "gestfips": 6,
+        "prtage": 30,
+        "pesex": 2,
+        "ptdtrace": 1,
+        "pehspnon": 2,
+        "pworwgt": 100.0,
+        "pternwa": 100000.0,
+        "pternhly": 2500.0,
+        "peernhry": 1,
+        "pehruslt": 40.0,
+        "prerelg": 1,
+        "pemlr": 1,
+        "peio1cow": 1,
+    }
+
+
+def test_load_org_training_data_serializes_first_cache_build(monkeypatch, tmp_path):
+    raw_month = pd.DataFrame(
+        {
+            "HRMIS": [4],
+            "gestfips": [6],
+            "prtage": [30],
+            "pesex": [2],
+            "ptdtrace": [1],
+            "pehspnon": [2],
+            "pworwgt": [100.0],
+            "pternwa": [100000.0],
+            "pternhly": [2500.0],
+            "peernhry": [1],
+            "pehruslt": [40.0],
+            "prerelg": [1],
+            "pemlr": [1],
+            "peio1cow": [1],
+        }
+    )
+    call_count = {"value": 0}
+
+    monkeypatch.setattr(
+        "policyengine_us_data.datasets.org.org.STORAGE_FOLDER", tmp_path
+    )
+    monkeypatch.setattr(
+        "policyengine_us_data.datasets.org.org.ORG_MONTHS",
+        ("may",),
+    )
+
+    def fake_load_month(year, month):
+        call_count["value"] += 1
+        time.sleep(0.2)
+        return raw_month.copy()
+
+    monkeypatch.setattr(
+        "policyengine_us_data.datasets.org.org._load_cps_basic_org_month",
+        fake_load_month,
+    )
+
+    load_org_training_data.cache_clear()
+    try:
+        with ThreadPoolExecutor(max_workers=2) as executor:
+            left = executor.submit(load_org_training_data)
+            right = executor.submit(load_org_training_data)
+            left_result = left.result()
+            right_result = right.result()
+    finally:
+        load_org_training_data.cache_clear()
+
+    assert call_count["value"] == 1
+    pd.testing.assert_frame_equal(left_result, right_result)
+
+
+def test_load_org_training_data_rebuilds_invalid_cached_file(monkeypatch, tmp_path):
+    raw_month = pd.DataFrame(
+        {
+            "HRMIS": [4],
+            "gestfips": [6],
+            "prtage": [30],
+            "pesex": [2],
+            "ptdtrace": [1],
+            "pehspnon": [2],
+            "pworwgt": [100.0],
+            "pternwa": [100000.0],
+            "pternhly": [2500.0],
+            "peernhry": [1],
+            "pehruslt": [40.0],
+            "prerelg": [1],
+            "pemlr": [1],
+            "peio1cow": [1],
+        }
+    )
+    cache_path = tmp_path / "census_cps_org_2024_wages.csv.gz"
+    pd.DataFrame(columns=["employment_income", "weekly_hours_worked"]).to_csv(
+        cache_path,
+        index=False,
+        compression="gzip",
+    )
+    call_count = {"value": 0}
+
+    monkeypatch.setattr(
+        "policyengine_us_data.datasets.org.org.STORAGE_FOLDER", tmp_path
+    )
+    monkeypatch.setattr(
+        "policyengine_us_data.datasets.org.org.ORG_MONTHS",
+        ("may",),
+    )
+
+    def fake_load_month(year, month):
+        call_count["value"] += 1
+        return raw_month.copy()
+
+    monkeypatch.setattr(
+        "policyengine_us_data.datasets.org.org._load_cps_basic_org_month",
+        fake_load_month,
+    )
+
+    load_org_training_data.cache_clear()
+    try:
+        rebuilt = load_org_training_data()
+    finally:
+        load_org_training_data.cache_clear()
+
+    assert call_count["value"] == 1
+    assert not rebuilt.empty
+    assert set(
+        [
+            "employment_income",
+            "weekly_hours_worked",
+            "age",
+            "is_female",
+            "is_hispanic",
+            "race_wbho",
+            "state_fips",
+            "hourly_wage",
+            "is_paid_hourly",
+            "sample_weight",
+        ]
+    ).issubset(rebuilt.columns)
+
+
 def test_build_union_priority_weights_reflect_bls_demographics():
     receiver = pd.DataFrame(
         {