Fix ORG loader fallback and Modal CI env

- Make ORG loading fall back from the CSV endpoint to the Census ZIP/dat source when the CSV response does not expose the required ORG fields.
- Parse the ZIP member with the published fixed-width layout, then normalize it back onto the existing ORG schema before building donor rows.
- Make pytest collection in the optimized Modal seam job lazy-load heavy version-manifest fixtures so the job does not need the full Google dependency chain.
- Pin PR and push Modal runs to the main environment for consistent CI behavior.

Author: juaristi22

.github/workflows/pr.yaml

@@ -100,7 +100,7 @@ jobs:
         with:
           python-version: "3.14"
       - name: Install optimized test deps
-        run: pip install modal pandas pytest
+        run: pip install modal pytest
       - name: Deploy Modal pipeline app
         run: modal deploy --env="${MODAL_ENVIRONMENT}" modal_app/pipeline.py
       - name: Run optimized integration tests
@@ -160,6 +160,7 @@ jobs:
     env:
       MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
       MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+      MODAL_ENVIRONMENT: main
       HUGGING_FACE_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }}
     steps:
       - uses: actions/checkout@v4
@@ -184,6 +185,6 @@ jobs:
             } >> "$GITHUB_STEP_SUMMARY"
           fi
 
-          modal run modal_app/data_build.py \
+          modal run --env="${MODAL_ENVIRONMENT}" modal_app/data_build.py \
             --branch=${{ github.head_ref || github.ref_name }} \
             ${STAGE_ARGS}

.github/workflows/push.yaml

@@ -21,6 +21,7 @@ jobs:
     env:
       MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }}
       MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }}
+      MODAL_ENVIRONMENT: main
       HUGGING_FACE_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }}
     steps:
       - uses: actions/checkout@v4
@@ -30,7 +31,7 @@ jobs:
       - run: pip install modal
       - name: Run linear integration tests on Modal
         run: |
-          modal run modal_app/data_build.py \
+          modal run --env="${MODAL_ENVIRONMENT}" modal_app/data_build.py \
             --upload \
             --branch=${{ github.ref_name }}
 

modal_app/README.md

@@ -145,7 +145,7 @@ Produces `source_imputed_stratified_extended_cps_2024.h5` from raw CPS/PUF/ACS/S
 | Method | Command |
 |--------|---------|
 | **Local** | `make data` |
-| **Modal (CI)** | `modal run modal_app/data_build.py --branch=<branch>` |
+| **Modal (CI)** | `modal run --env=main modal_app/data_build.py --branch=<branch>` |
 | **GitHub Actions** | Automatic on merge to `main` via `code_changes.yaml` → `reusable_test.yaml` (with `full_suite: true`). Also triggered by `pr_code_changes.yaml` on PRs. |
 
 Notes:

policyengine_us_data/datasets/org/README.md

@@ -11,9 +11,10 @@ the CPS ASEC records.
 
 The checked-in code does not vendor the donor file itself. Instead,
 `org.py` builds `census_cps_org_2024_wages.csv.gz` on demand by
-downloading the twelve official CPS basic monthly public-use CSVs for
+downloading the twelve official CPS basic monthly public-use files for
 `ORG_YEAR` (currently 2024) directly from the Census Bureau and filtering
-each file to the ORG rotations.
+each file to the ORG rotations. The loader prefers the CSV endpoint and
+falls back to the ZIP'd fixed-width archive when the CSV is unavailable.
 
 ## Documentation
 
@@ -31,8 +32,8 @@ See also:
 
 ## Data products in this folder
 
-- `org.py` — downloads the twelve monthly CSVs, filters to the MIS-4 and
-  MIS-8 outgoing rotations (`HRMIS`), and caches the combined ORG donor
-  frame. Trains a QRF model to impute `wage_income`, `hourly_wage`, and
-  union-coverage variables onto the CPS ASEC records used by the
-  Enhanced CPS pipeline.
+- `org.py` — downloads the twelve monthly public-use files, filters to
+  the MIS-4 and MIS-8 outgoing rotations (`HRMIS`), and caches the
+  combined ORG donor frame. Trains a QRF model to impute `wage_income`,
+  `hourly_wage`, and union-coverage variables onto the CPS ASEC records
+  used by the Enhanced CPS pipeline.

policyengine_us_data/datasets/org/org.py

@@ -1,16 +1,18 @@
 """Build and cache ORG-style labor-market donor rows from CPS basic-month files.
 
 The checked-in code does not vendor the donor file itself. `ORG_FILENAME` is a
-generated cache built from the 12 official CPS basic monthly public-use CSVs for
-`ORG_YEAR`, then reused as the donor sample for wage, hourly-pay, and union
-imputation onto CPS records.
+generated cache built from the 12 official CPS basic monthly public-use files
+for `ORG_YEAR`, then reused as the donor sample for wage, hourly-pay, and
+union imputation onto CPS records.
 """
 
 from contextlib import contextmanager
 from functools import lru_cache
 from io import BytesIO
+from io import TextIOWrapper
 from pathlib import Path
 import fcntl
+import zipfile
 
 from microimpute.models.qrf import QRF
 import numpy as np
@@ -52,6 +54,44 @@
     "peio1cow",
 ]
 
+# Field order in the 2024 CPS basic-month fixed-width dat file.
+CPS_BASIC_MONTHLY_ORG_FWF_COLUMNS = [
+    "HRMIS",
+    "gestfips",
+    "prtage",
+    "pesex",
+    "ptdtrace",
+    "pehspnon",
+    "pemlr",
+    "pehruslt",
+    "peio1cow",
+    "prerelg",
+    "peernhry",
+    "pternhly",
+    "pternwa",
+    "pworwgt",
+]
+
+# 2024 Basic CPS record layout positions for the ORG columns we need, in file order.
+# The CSV endpoint sometimes returns challenge HTML in Modal, so the ZIP'd
+# fixed-width dat fallback is the reliable source of truth.
+CPS_BASIC_MONTHLY_ORG_COLSPECS = [
+    (62, 64),  # HRMIS 63-64
+    (92, 94),  # GESTFIPS 93-94
+    (121, 123),  # PRTAGE 122-123
+    (128, 130),  # PESEX 129-130
+    (138, 140),  # PTDTRACE 139-140
+    (156, 158),  # PEHSPNON 157-158
+    (179, 181),  # PEMLR 180-181
+    (223, 226),  # PEHRUSLT 224-226
+    (430, 433),  # PEIO1COW 432-433
+    (497, 499),  # PRERELG 498-499
+    (505, 507),  # PEERNHRY 506-507
+    (519, 523),  # PTERNHLY 520-523
+    (526, 534),  # PTERNWA 527-534
+    (602, 612),  # PWORWGT 603-612
+]
+
 ORG_PREDICTORS = [
     "employment_income",
     "weekly_hours_worked",
@@ -186,6 +226,14 @@ def _cps_basic_org_month_url(year: int, month: str) -> str:
     )
 
 
+def _cps_basic_org_month_zip_url(year: int, month: str) -> str:
+    year_suffix = str(year)[-2:]
+    return (
+        f"https://www2.census.gov/programs-surveys/cps/datasets/"
+        f"{year}/basic/{month}{year_suffix}pub.zip"
+    )
+
+
 def _resolve_cps_basic_org_column_names(
     columns: pd.Index | list[str],
 ) -> list[str]:
@@ -213,13 +261,20 @@ def _select_cps_basic_org_columns(month_df: pd.DataFrame) -> pd.DataFrame:
     return selected
 
 
-def _load_cps_basic_org_month(
+def _coerce_cps_basic_org_numeric_columns(
+    month_df: pd.DataFrame,
+) -> pd.DataFrame:
+    """Convert ORG columns into numeric values regardless of source format."""
+    return month_df.apply(lambda column: pd.to_numeric(column, errors="coerce"))
+
+
+def _load_cps_basic_org_month_from_csv(
     year: int,
     month: str,
     *,
     max_attempts: int = 3,
 ) -> pd.DataFrame:
-    """Load one CPS basic-month file with light retry around transient fetch/parser issues."""
+    """Load one CPS basic-month file from the CSV endpoint."""
     url = _cps_basic_org_month_url(year, month)
     last_error: Exception | None = None
 
@@ -235,16 +290,99 @@ def _load_cps_basic_org_month(
                 usecols=selected_columns,
                 low_memory=False,
             )
-            return _select_cps_basic_org_columns(month_df)
+            return _coerce_cps_basic_org_numeric_columns(
+                _select_cps_basic_org_columns(month_df)
+            )
         except Exception as error:
             last_error = error
 
     raise ValueError(
-        f"Failed to load CPS basic ORG month {month} {year} after "
+        f"Failed to load CPS basic ORG month {month} {year} from CSV after "
         f"{max_attempts} attempts"
     ) from last_error
 
 
+def _load_cps_basic_org_month_from_zip(
+    year: int,
+    month: str,
+    *,
+    max_attempts: int = 3,
+) -> pd.DataFrame:
+    """Load one CPS basic-month file from the ZIP'd fixed-width dat archive."""
+    url = _cps_basic_org_month_zip_url(year, month)
+    last_error: Exception | None = None
+
+    for _ in range(max_attempts):
+        try:
+            response = requests.get(url, timeout=60)
+            response.raise_for_status()
+            with zipfile.ZipFile(BytesIO(response.content)) as archive:
+                dat_name = next(
+                    (
+                        name
+                        for name in archive.namelist()
+                        if name.lower().endswith(".dat")
+                    ),
+                    None,
+                )
+                if dat_name is None:
+                    raise ValueError(
+                        "CPS basic ORG month zip did not contain a .dat file"
+                    )
+                with (
+                    archive.open(dat_name) as raw,
+                    TextIOWrapper(
+                        raw,
+                        encoding="utf-8",
+                        newline="",
+                    ) as text,
+                ):
+                    month_df = pd.read_fwf(
+                        text,
+                        colspecs=CPS_BASIC_MONTHLY_ORG_COLSPECS,
+                        names=CPS_BASIC_MONTHLY_ORG_FWF_COLUMNS,
+                        header=None,
+                        dtype=str,
+                    )
+            return _coerce_cps_basic_org_numeric_columns(
+                _select_cps_basic_org_columns(month_df)
+            )
+        except Exception as error:
+            last_error = error
+
+    raise ValueError(
+        f"Failed to load CPS basic ORG month {month} {year} from ZIP after "
+        f"{max_attempts} attempts"
+    ) from last_error
+
+
+def _load_cps_basic_org_month(
+    year: int,
+    month: str,
+    *,
+    max_attempts: int = 3,
+) -> pd.DataFrame:
+    """Load one CPS basic-month file, preferring CSV and falling back to ZIP."""
+    try:
+        return _load_cps_basic_org_month_from_csv(
+            year,
+            month,
+            max_attempts=max_attempts,
+        )
+    except Exception as csv_error:
+        try:
+            return _load_cps_basic_org_month_from_zip(
+                year,
+                month,
+                max_attempts=max_attempts,
+            )
+        except Exception as zip_error:
+            raise ValueError(
+                f"Failed to load CPS basic ORG month {month} {year} from CSV or ZIP: "
+                f"csv={csv_error!r}; zip={zip_error!r}"
+            ) from zip_error
+
+
 @contextmanager
 def _org_cache_build_lock(lock_path: Path):
     lock_path.parent.mkdir(parents=True, exist_ok=True)

tests/conftest.py

@@ -1,18 +1,12 @@
 """Shared fixtures and helpers for version manifest tests."""
 
+from __future__ import annotations
+
 import json
 from unittest.mock import MagicMock
 
 import pytest
 
-from policyengine_us_data.utils.version_manifest import (
-    HFVersionInfo,
-    GCSVersionInfo,
-    VersionManifest,
-    VersionRegistry,
-)
-from policyengine_us_data.utils.policyengine import PolicyEngineUSBuildInfo
-
 # -- Fixtures ------------------------------------------------------
 
 
@@ -27,6 +21,8 @@ def sample_generations() -> dict[str, int]:
 
 @pytest.fixture
 def sample_hf_info() -> HFVersionInfo:
+    from policyengine_us_data.utils.version_manifest import HFVersionInfo
+
     return HFVersionInfo(
         repo="policyengine/policyengine-us-data",
         commit="abc123def456",
@@ -35,6 +31,8 @@ def sample_hf_info() -> HFVersionInfo:
 
 @pytest.fixture
 def sample_policyengine_us_info() -> PolicyEngineUSBuildInfo:
+    from policyengine_us_data.utils.policyengine import PolicyEngineUSBuildInfo
+
     return PolicyEngineUSBuildInfo(
         version="1.587.0",
         locked_version="1.587.0",
@@ -49,6 +47,11 @@ def sample_manifest(
     sample_hf_info: HFVersionInfo,
     sample_policyengine_us_info: PolicyEngineUSBuildInfo,
 ) -> VersionManifest:
+    from policyengine_us_data.utils.version_manifest import (
+        GCSVersionInfo,
+        VersionManifest,
+    )
+
     return VersionManifest(
         version="1.72.3",
         created_at="2026-03-10T14:30:00Z",
@@ -66,6 +69,8 @@ def sample_registry(
     sample_manifest: VersionManifest,
 ) -> VersionRegistry:
     """A registry with one version entry."""
+    from policyengine_us_data.utils.version_manifest import VersionRegistry
+
     return VersionRegistry(
         current="1.72.3",
         versions=[sample_manifest],