Add Stage 2 package payload reader

Author: anth-volk

changelog.d/1073.changed

@@ -0,0 +1 @@
+Add typed Stage 2 calibration package payload reader and writer helpers.

docs/pipeline_map.yaml

@@ -821,8 +821,12 @@ stages:
     - build_matrix
     - build_matrix_chunked
     - stage2_artifact_specs
+    - stage2_payload_boundary
+    - stage2_payload_writer
+    - stage2_payload_reader
     - stage2_calibration_package_writer
     - out_pkg
+    - out_metadata
     - stage2_calibration_package_contract_writer
     - out_contract
     - stage2_calibration_package_contract_validator
@@ -875,6 +879,10 @@ stages:
     label: calibration_package.pkl
     node_type: artifact
     description: X_sparse CSR matrix, targets_df, initial_weights, metadata
+  - id: out_metadata
+    label: calibration_package_meta.json
+    node_type: artifact
+    description: Metadata sidecar generated from the typed package payload and Stage 2 contract
   - id: out_contract
     label: calibration_package_contract.json
     node_type: artifact
@@ -983,21 +991,40 @@ stages:
     edge_type: uses_library
     label: chunked path
   - source: build_matrix
-    target: stage2_calibration_package_writer
+    target: stage2_payload_boundary
     edge_type: data_flow
   - source: build_matrix_chunked
-    target: stage2_calibration_package_writer
+    target: stage2_payload_boundary
     edge_type: data_flow
+  - source: stage2_payload_boundary
+    target: stage2_payload_writer
+    edge_type: data_flow
+    label: typed pickle payload
+  - source: stage2_payload_writer
+    target: stage2_calibration_package_writer
+    edge_type: uses_library
+    label: pickle write
   - source: stage2_artifact_specs
     target: stage2_calibration_package_writer
     edge_type: uses_utility
     label: package path
   - source: stage2_calibration_package_writer
     target: out_pkg
     edge_type: produces_artifact
+  - source: stage2_payload_writer
+    target: out_metadata
+    edge_type: produces_artifact
+    label: sidecar metadata
+  - source: out_pkg
+    target: stage2_payload_reader
+    edge_type: data_flow
   - source: out_pkg
     target: stage2_calibration_package_contract_writer
     edge_type: data_flow
+  - source: stage2_payload_reader
+    target: stage2_calibration_package_contract_writer
+    edge_type: uses_library
+    label: summary and checksum
   - source: stage2_artifact_specs
     target: stage2_calibration_package_contract_writer
     edge_type: uses_utility

modal_app/remote_calibration_runner.py

@@ -321,23 +321,30 @@ def _print_provenance_from_meta(meta: dict, current_branch: str = None) -> None:
 
 
 def _write_package_sidecar(pkg_path: str) -> bool:
-    """Extract metadata from a pickle package and write a JSON sidecar.
+    """Write package metadata from the typed payload and contract sidecar.
 
     Returns:
         True if sidecar was written successfully, False otherwise.
     """
-    import json
     import logging
-    import pickle
 
-    sidecar_path = pkg_path.replace(".pkl", "_meta.json")
     try:
-        with open(pkg_path, "rb") as f:
-            package = pickle.load(f)
-        meta = package.get("metadata", {})
-        del package
-        with open(sidecar_path, "w") as f:
-            json.dump(meta, f, indent=2)
+        from policyengine_us_data.calibration_package.payload import (
+            CalibrationPackageReader,
+            CalibrationPackageWriter,
+        )
+        from policyengine_us_data.calibration_package.specs import (
+            CALIBRATION_PACKAGE_CONTRACT_FILENAME,
+        )
+        from policyengine_us_data.stage_contracts.io import read_contract
+
+        package_path = Path(pkg_path)
+        payload = CalibrationPackageReader(package_path=package_path).read()
+        contract_path = package_path.with_name(CALIBRATION_PACKAGE_CONTRACT_FILENAME)
+        contract = read_contract(contract_path) if contract_path.exists() else None
+        sidecar_path = CalibrationPackageWriter(
+            package_path=package_path,
+        ).write_metadata_sidecar(payload, contract=contract)
         print(
             f"Sidecar metadata written to {sidecar_path}",
             flush=True,

policyengine_us_data/calibration/unified_calibration.py

@@ -44,6 +44,11 @@
 from policyengine_us_data.calibration.calibration_utils import (
     create_target_groups,
 )
+from policyengine_us_data.calibration_package.payload import (
+    CalibrationPackagePayload,
+    CalibrationPackageReader,
+    CalibrationPackageWriter,
+)
 from policyengine_us_data.calibration_package.specs import (
     DEFAULT_TARGET_CONFIG_PATH as DEFAULT_TARGET_CONFIG_RELATIVE_PATH,
     TargetConfigIdentity,
@@ -677,20 +682,16 @@ def save_calibration_package(
         cd_geoid: CD GEOID array from geography assignment.
         block_geoid: Block GEOID array from geography assignment.
     """
-    import pickle
-
-    package = {
-        "X_sparse": X_sparse,
-        "targets_df": targets_df,
-        "target_names": target_names,
-        "metadata": metadata,
-        "initial_weights": initial_weights,
-        "cd_geoid": cd_geoid,
-        "block_geoid": block_geoid,
-    }
-    Path(path).parent.mkdir(parents=True, exist_ok=True)
-    with open(path, "wb") as f:
-        pickle.dump(package, f, protocol=pickle.HIGHEST_PROTOCOL)
+    payload = CalibrationPackagePayload(
+        X_sparse=X_sparse,
+        targets_df=targets_df,
+        target_names=target_names,
+        metadata=metadata,
+        initial_weights=initial_weights,
+        cd_geoid=cd_geoid,
+        block_geoid=block_geoid,
+    )
+    CalibrationPackageWriter(package_path=Path(path)).write(payload)
     logger.info("Calibration package saved to %s", path)
 
 
@@ -703,16 +704,14 @@ def load_calibration_package(path: str) -> dict:
     Returns:
         Dict with X_sparse, targets_df, target_names, metadata.
     """
-    import pickle
-
-    with open(path, "rb") as f:
-        package = pickle.load(f)
+    payload = CalibrationPackageReader(package_path=Path(path)).read()
+    package = payload.to_mapping()
     logger.info(
         "Loaded package: %d targets, %d records",
-        package["X_sparse"].shape[0],
-        package["X_sparse"].shape[1],
+        payload.X_sparse.shape[0],
+        payload.X_sparse.shape[1],
     )
-    meta = package.get("metadata", {})
+    meta = payload.metadata
     print_package_provenance(meta)
     check_package_staleness(meta)
     return package
@@ -1727,15 +1726,15 @@ def run_calibration(
 
     initial_weights = compute_initial_weights(X_sparse, targets_df)
     if package_output_path:
-        package_payload = {
-            "X_sparse": X_sparse,
-            "targets_df": targets_df,
-            "target_names": target_names,
-            "metadata": metadata,
-            "initial_weights": initial_weights,
-            "cd_geoid": geography.cd_geoid,
-            "block_geoid": geography.block_geoid,
-        }
+        package_payload = CalibrationPackagePayload(
+            X_sparse=X_sparse,
+            targets_df=targets_df,
+            target_names=target_names,
+            metadata=metadata,
+            initial_weights=initial_weights,
+            cd_geoid=geography.cd_geoid,
+            block_geoid=geography.block_geoid,
+        )
         save_calibration_package(
             package_output_path,
             X_sparse,

policyengine_us_data/calibration_package/__init__.py

@@ -26,6 +26,13 @@
     stage2_input_bundle_from_stage1_contract,
     stage2_input_bundle_from_stage1_contract_path,
 )
+from .payload import (
+    LEGACY_MISSING_GEOGRAPHY_WARNING,
+    REQUIRED_PACKAGE_KEYS,
+    CalibrationPackagePayload,
+    CalibrationPackageReader,
+    CalibrationPackageWriter,
+)
 
 __all__ = [
     "CALIBRATION_PACKAGE_CONTRACT_FILENAME",
@@ -41,10 +48,15 @@
     "TARGET_DATABASE_FILENAME",
     "CalibrationPackageArtifactPaths",
     "CalibrationPackageOutputBundle",
+    "CalibrationPackagePayload",
+    "CalibrationPackageReader",
+    "CalibrationPackageWriter",
+    "LEGACY_MISSING_GEOGRAPHY_WARNING",
     "Stage2BuildContext",
     "Stage2InputBundle",
     "Stage2InputBundleError",
     "Stage2InputSource",
+    "REQUIRED_PACKAGE_KEYS",
     "TargetConfigIdentity",
     "calibration_package_artifact_paths",
     "resolve_target_config_identity",