Fix release manifest and ORG CI regressions

Author: MaxGhenis

policyengine_us_data/datasets/org/org.py

@@ -203,6 +203,30 @@ def _select_cps_basic_org_columns(month_df: pd.DataFrame) -> pd.DataFrame:
     return selected
 
 
+def _resolve_cps_basic_org_usecols(url: str) -> list[str]:
+    """Resolve the exact remote column names before reading the full CPS month.
+
+    Pandas' callable `usecols` path against remote CSVs can intermittently
+    mis-handle the header row and return an empty selection. Resolving the
+    concrete header first avoids that parser path while keeping the full read
+    column-limited.
+    """
+    header_df = pd.read_csv(url, nrows=0)
+    column_lookup = {
+        str(column).lower(): column
+        for column in header_df.columns
+        if isinstance(column, str)
+    }
+    missing = [
+        column
+        for column in CPS_BASIC_MONTHLY_ORG_COLUMNS
+        if column.lower() not in column_lookup
+    ]
+    if missing:
+        raise ValueError(f"CPS basic ORG month is missing required columns: {missing}")
+    return [column_lookup[column.lower()] for column in CPS_BASIC_MONTHLY_ORG_COLUMNS]
+
+
 def _load_cps_basic_org_month(
     year: int,
     month: str,
@@ -211,16 +235,14 @@ 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:
+            usecols = _resolve_cps_basic_org_usecols(url)
             month_df = pd.read_csv(
                 url,
-                usecols=lambda column: (
-                    isinstance(column, str) and column.lower() in required_columns
-                ),
+                usecols=usecols,
                 low_memory=False,
             )
             return _select_cps_basic_org_columns(month_df)

policyengine_us_data/utils/data_upload.py

@@ -380,6 +380,14 @@ def get_matching_finalized_release_manifest(
         ],
         existing_manifest=finalized_manifest,
     )
+    if "created_at" in finalized_manifest:
+        candidate_manifest["created_at"] = finalized_manifest["created_at"]
+    finalized_build = finalized_manifest.get("build")
+    if isinstance(finalized_build, dict):
+        candidate_build = candidate_manifest.setdefault("build", {})
+        for field in ("build_id", "built_at"):
+            if field in finalized_build:
+                candidate_build[field] = finalized_build[field]
     if candidate_manifest != finalized_manifest:
         raise RuntimeError(
             f"Release {version} is already finalized on {hf_repo_name}. "

tests/unit/datasets/test_org.py

@@ -172,11 +172,14 @@ def test_load_cps_basic_org_month_retries_after_transient_parser_failure(
             "PEIO1COW": [1],
         }
     )
+    header_df = pd.DataFrame(columns=month_df.columns)
 
     def fake_read_csv(*args, **kwargs):
         calls.append(kwargs)
         if len(calls) == 1:
-            raise ValueError("Usecols do not match columns")
+            raise ValueError("Temporary header parse failure")
+        if kwargs.get("nrows") == 0:
+            return header_df
         return month_df
 
     monkeypatch.setattr(
@@ -185,8 +188,10 @@ def fake_read_csv(*args, **kwargs):
 
     loaded = _load_cps_basic_org_month(2024, "may", max_attempts=2)
 
-    assert len(calls) == 2
-    assert callable(calls[0]["usecols"])
+    assert len(calls) == 3
+    assert calls[0]["nrows"] == 0
+    assert calls[1]["nrows"] == 0
+    assert calls[2]["usecols"] == month_df.columns.tolist()
     assert loaded.columns.tolist() == CPS_BASIC_MONTHLY_ORG_COLUMNS
 
 

tests/unit/test_release_manifest.py

@@ -28,7 +28,9 @@ def _sha256(content: bytes) -> str:
     return hashlib.sha256(content).hexdigest()
 
 
-EXPECTED_COMPATIBLE_MODEL_PACKAGES = [{"name": "policyengine-us", "version": "1.634.4"}]
+EXPECTED_COMPATIBLE_MODEL_PACKAGES = [
+    {"name": "policyengine-us", "specifier": "==1.634.4"}
+]
 
 
 def _build_local_area_manifest(
@@ -64,6 +66,8 @@ def test_build_release_manifest_tracks_uploaded_artifacts(tmp_path):
         version="1.73.0",
         repo_id="policyengine/policyengine-us-data",
         model_package_version="1.634.4",
+        model_package_git_sha="deadbeef",
+        model_package_data_build_fingerprint="sha256:fingerprint",
         created_at="2026-04-10T12:00:00Z",
     )
 
@@ -73,6 +77,16 @@ def test_build_release_manifest_tracks_uploaded_artifacts(tmp_path):
     }
     assert manifest["schema_version"] == RELEASE_MANIFEST_SCHEMA_VERSION
     assert manifest["compatible_model_packages"] == EXPECTED_COMPATIBLE_MODEL_PACKAGES
+    assert manifest["build"] == {
+        "build_id": "policyengine-us-data-1.73.0",
+        "built_at": "2026-04-10T12:00:00Z",
+        "built_with_model_package": {
+            "name": "policyengine-us",
+            "version": "1.634.4",
+            "git_sha": "deadbeef",
+            "data_build_fingerprint": "sha256:fingerprint",
+        },
+    }
     assert manifest["default_datasets"] == {"national": "enhanced_cps_2024"}
 
     assert manifest["artifacts"]["enhanced_cps_2024"] == {
@@ -122,13 +136,25 @@ def test_build_release_manifest_merges_existing_release_same_version(tmp_path):
         version="1.73.0",
         repo_id="policyengine/policyengine-us-data",
         model_package_version="1.634.4",
+        model_package_git_sha="deadbeef",
+        model_package_data_build_fingerprint="sha256:fingerprint",
         existing_manifest=existing_manifest,
         created_at="2026-04-10T12:00:00Z",
     )
 
     assert set(manifest["artifacts"]) == {"enhanced_cps_2024", "districts/NC-01"}
     assert manifest["default_datasets"] == {"national": "enhanced_cps_2024"}
-    assert manifest["created_at"] == "2026-04-09T12:00:00Z"
+    assert manifest["created_at"] == "2026-04-10T12:00:00Z"
+    assert manifest["build"] == {
+        "build_id": "policyengine-us-data-1.73.0",
+        "built_at": "2026-04-10T12:00:00Z",
+        "built_with_model_package": {
+            "name": "policyengine-us",
+            "version": "1.634.4",
+            "git_sha": "deadbeef",
+            "data_build_fingerprint": "sha256:fingerprint",
+        },
+    }
     assert manifest["artifacts"]["districts/NC-01"]["sha256"] == _sha256(district_bytes)
 
 
@@ -167,8 +193,12 @@ def test_upload_files_to_hf_adds_release_manifest_operations(tmp_path):
             return_value=None,
         ),
         patch(
-            "policyengine_us_data.utils.data_upload.metadata.version",
-            return_value="1.634.4",
+            "policyengine_us_data.utils.data_upload._get_model_package_build_metadata",
+            return_value={
+                "version": "1.634.4",
+                "git_sha": "deadbeef",
+                "data_build_fingerprint": "sha256:fingerprint",
+            },
         ),
         patch.dict(
             "policyengine_us_data.utils.data_upload.os.environ",
@@ -215,8 +245,12 @@ def test_upload_files_to_hf_does_not_tag_until_finalize(tmp_path):
             return_value=None,
         ),
         patch(
-            "policyengine_us_data.utils.data_upload.metadata.version",
-            return_value="1.634.4",
+            "policyengine_us_data.utils.data_upload._get_model_package_build_metadata",
+            return_value={
+                "version": "1.634.4",
+                "git_sha": "deadbeef",
+                "data_build_fingerprint": "sha256:fingerprint",
+            },
         ),
         patch.dict(
             "policyengine_us_data.utils.data_upload.os.environ",
@@ -250,6 +284,16 @@ def test_publish_release_manifest_to_hf_can_finalize_and_tag(tmp_path):
         "compatible_model_packages": EXPECTED_COMPATIBLE_MODEL_PACKAGES,
         "default_datasets": {"national": "enhanced_cps_2024"},
         "created_at": "2026-04-10T12:00:00Z",
+        "build": {
+            "build_id": "policyengine-us-data-1.73.0",
+            "built_at": "2026-04-10T12:00:00Z",
+            "built_with_model_package": {
+                "name": "policyengine-us",
+                "version": "1.634.4",
+                "git_sha": "deadbeef",
+                "data_build_fingerprint": "sha256:fingerprint",
+            },
+        },
         "artifacts": {
             "enhanced_cps_2024": {
                 "kind": "microdata",
@@ -271,22 +315,36 @@ def test_publish_release_manifest_to_hf_can_finalize_and_tag(tmp_path):
             ),
         ),
         patch(
-            "policyengine_us_data.utils.data_upload.metadata.version",
-            return_value="1.634.4",
+            "policyengine_us_data.utils.data_upload._get_model_package_build_metadata",
+            return_value={
+                "version": "1.634.4",
+                "git_sha": "deadbeef",
+                "data_build_fingerprint": "sha256:fingerprint",
+            },
         ),
         patch.dict(
             "policyengine_us_data.utils.data_upload.os.environ",
             {"HUGGING_FACE_TOKEN": "token"},
             clear=False,
         ),
     ):
-        publish_release_manifest_to_hf(
+        manifest = publish_release_manifest_to_hf(
             [(state_path, "states/AL.h5")],
             version="1.73.0",
             create_tag=True,
         )
 
     mock_api.create_tag.assert_called_once()
+    assert manifest["build"] == {
+        "build_id": "policyengine-us-data-1.73.0",
+        "built_at": "2026-04-10T12:00:00Z",
+        "built_with_model_package": {
+            "name": "policyengine-us",
+            "version": "1.634.4",
+            "git_sha": "deadbeef",
+            "data_build_fingerprint": "sha256:fingerprint",
+        },
+    }
 
 
 def test_missing_release_prefixes_requires_full_local_area_bundle():
@@ -373,6 +431,16 @@ def test_publish_release_manifest_to_hf_rejects_finalized_release(tmp_path):
         "compatible_model_packages": EXPECTED_COMPATIBLE_MODEL_PACKAGES,
         "default_datasets": {"national": "enhanced_cps_2024"},
         "created_at": "2026-04-10T12:00:00Z",
+        "build": {
+            "build_id": "policyengine-us-data-1.73.0",
+            "built_at": "2026-04-10T12:00:00Z",
+            "built_with_model_package": {
+                "name": "policyengine-us",
+                "version": "1.634.4",
+                "git_sha": "deadbeef",
+                "data_build_fingerprint": "sha256:fingerprint",
+            },
+        },
         "artifacts": {
             "states/AL": {
                 "kind": "microdata",
@@ -396,6 +464,14 @@ def test_publish_release_manifest_to_hf_rejects_finalized_release(tmp_path):
             "policyengine_us_data.utils.data_upload._get_model_package_version",
             return_value="1.634.4",
         ),
+        patch(
+            "policyengine_us_data.utils.data_upload._get_model_package_build_metadata",
+            return_value={
+                "version": "1.634.4",
+                "git_sha": "deadbeef",
+                "data_build_fingerprint": "sha256:fingerprint",
+            },
+        ),
     ):
         manifest = publish_release_manifest_to_hf(
             [(state_path, "states/AL.h5")],
@@ -420,6 +496,16 @@ def test_publish_release_manifest_to_hf_rejects_mutating_finalized_release(tmp_p
         "compatible_model_packages": EXPECTED_COMPATIBLE_MODEL_PACKAGES,
         "default_datasets": {"national": "enhanced_cps_2024"},
         "created_at": "2026-04-10T12:00:00Z",
+        "build": {
+            "build_id": "policyengine-us-data-1.73.0",
+            "built_at": "2026-04-10T12:00:00Z",
+            "built_with_model_package": {
+                "name": "policyengine-us",
+                "version": "1.634.4",
+                "git_sha": "deadbeef",
+                "data_build_fingerprint": "sha256:fingerprint",
+            },
+        },
         "artifacts": {
             "states/AL": {
                 "kind": "microdata",
@@ -443,6 +529,14 @@ def test_publish_release_manifest_to_hf_rejects_mutating_finalized_release(tmp_p
             "policyengine_us_data.utils.data_upload._get_model_package_version",
             return_value="1.634.4",
         ),
+        patch(
+            "policyengine_us_data.utils.data_upload._get_model_package_build_metadata",
+            return_value={
+                "version": "1.634.4",
+                "git_sha": "deadbeef",
+                "data_build_fingerprint": "sha256:fingerprint",
+            },
+        ),
     ):
         try:
             publish_release_manifest_to_hf(