Tighten Stage 2 target config identity validation

Author: anth-volk

policyengine_us_data/calibration/unified_calibration.py

@@ -1452,6 +1452,11 @@ def run_calibration(
 
     started_at = _utc_now_isoformat()
     t0 = time.time()
+    resolved_target_identity = _target_config_identity_for_metadata(
+        target_config=target_config,
+        target_config_path=target_config_path,
+        target_config_identity=target_config_identity,
+    )
 
     # Early exit: load pre-built package
     if package_path is not None:
@@ -1697,11 +1702,6 @@ def run_calibration(
     # Step 6b: Save the calibration package. By default this is the
     # minimal package selected by target_config.yaml; use
     # --all-active-targets to build a broad diagnostic package.
-    resolved_target_identity = _target_config_identity_for_metadata(
-        target_config=target_config,
-        target_config_path=target_config_path,
-        target_config_identity=target_config_identity,
-    )
     metadata = {
         "dataset_path": dataset_path,
         "db_path": db_path,

policyengine_us_data/stage_contracts/calibration_package.py

@@ -568,7 +568,7 @@ def _parameters_with_package_identity(
         payload["target_config_mode"] = (
             "all_active_targets" if payload.get("target_config") is None else "explicit"
         )
-    return payload
+    return CalibrationPackageParameters.from_dict(payload).to_dict()
 
 
 def _require_existing_file(path: Path, label: str) -> None:

policyengine_us_data/stage_contracts/calibration_package_schema.py

@@ -238,11 +238,10 @@ def __post_init__(self) -> None:
         _validate_bool(self.parallel_matrix, "parallel_matrix")
         if self.target_config is not None and not isinstance(self.target_config, str):
             raise ValueError("target_config must be a string or None")
-        if self.target_config_sha256 is not None and not isinstance(
+        _validate_optional_sha256(
             self.target_config_sha256,
-            str,
-        ):
-            raise ValueError("target_config_sha256 must be a string or None")
+            "target_config_sha256",
+        )
         if self.target_config_mode is not None:
             if not isinstance(self.target_config_mode, str):
                 raise ValueError("target_config_mode must be a string or None")
@@ -649,5 +648,8 @@ def _validate_optional_sha256(value: Any, key: str) -> None:
         return
     if not isinstance(value, str) or not value.startswith("sha256:"):
         raise ValueError(f"Calibration package field {key!r} must be a SHA-256 digest")
-    if len(value) != len("sha256:") + 64:
+    digest = value.removeprefix("sha256:")
+    if len(digest) != 64:
+        raise ValueError(f"Calibration package field {key!r} must be a SHA-256 digest")
+    if any(character not in "0123456789abcdef" for character in digest.lower()):
         raise ValueError(f"Calibration package field {key!r} must be a SHA-256 digest")

tests/unit/calibration/test_unified_calibration.py

@@ -44,18 +44,21 @@
     _calibration_package_contract_parameters,
     _target_config_identity_for_metadata,
     check_package_staleness,
+    run_calibration,
 )
 from policyengine_us_data.stage_contracts.calibration_package import (
     CalibrationPackageParameters,
 )
 
+TARGET_CONFIG_SHA256 = "sha256:" + "a" * 64
+
 
 def test_calibration_package_contract_parameters_track_effective_matrix_mode():
     params = _calibration_package_contract_parameters(
         workers=8,
         n_clones=430,
         target_config_path="policyengine_us_data/calibration/target_config.yaml",
-        target_config_sha256="abc123",
+        target_config_sha256=TARGET_CONFIG_SHA256,
         target_config_mode="default",
         skip_county=True,
         skip_source_impute=True,
@@ -71,7 +74,7 @@ def test_calibration_package_contract_parameters_track_effective_matrix_mode():
         "workers": None,
         "n_clones": 430,
         "target_config": "policyengine_us_data/calibration/target_config.yaml",
-        "target_config_sha256": "abc123",
+        "target_config_sha256": TARGET_CONFIG_SHA256,
         "target_config_mode": "default",
         "skip_county": True,
         "skip_source_impute": True,
@@ -116,6 +119,22 @@ def test_target_config_identity_for_metadata_requires_identity_for_parsed_config
         )
 
 
+def test_run_calibration_validates_target_identity_before_dataset_loading():
+    with (
+        patch.dict(sys.modules, {"policyengine_us": None}),
+        pytest.raises(
+            ValueError,
+            match="target_config_path or target_config_identity",
+        ),
+    ):
+        run_calibration(
+            dataset_path="/missing/source.h5",
+            db_path="/missing/policy_data.db",
+            target_config={"include": []},
+            target_config_path=None,
+        )
+
+
 def test_check_package_staleness_warns_for_old_utc_timestamp(
     capsys,
     monkeypatch,

tests/unit/fixtures/calibration_package_stage_contract.py

@@ -22,6 +22,7 @@
 CALIBRATION_COMPLETED_AT = "2026-05-08T12:02:00Z"
 CALIBRATION_DURATION_S = 120.0
 TARGET_CONFIG_PATH = "policyengine_us_data/calibration/target_config.yaml"
+TARGET_CONFIG_SHA256 = "sha256:" + "a" * 64
 
 CALIBRATION_BLOCK_GEOIDS = ("010010001", "010010002", "020010001")
 CALIBRATION_CD_GEOIDS = ("0101", "0102", "0201")
@@ -66,7 +67,7 @@ def calibration_package_payload() -> dict[str, Any]:
             "dataset_sha256": "sha256:dataset",
             "db_sha256": "sha256:db",
             "target_config_path": TARGET_CONFIG_PATH,
-            "target_config_sha256": "sha256:target-config",
+            "target_config_sha256": TARGET_CONFIG_SHA256,
             "target_config_mode": "explicit",
             "n_clones": 3,
             "seed": 42,
@@ -149,7 +150,7 @@ def calibration_package_parameters() -> dict[str, Any]:
         "workers": None,
         "n_clones": 3,
         "target_config": TARGET_CONFIG_PATH,
-        "target_config_sha256": "sha256:target-config",
+        "target_config_sha256": TARGET_CONFIG_SHA256,
         "target_config_mode": "explicit",
         "skip_county": True,
         "skip_source_impute": True,

tests/unit/test_calibration_package_stage_contract.py

@@ -1,4 +1,5 @@
 from tests.unit.fixtures.calibration_package_stage_contract import (
+    TARGET_CONFIG_SHA256,
     TARGET_CONFIG_PATH,
     calibration_package_contract,
     calibration_package_parameters,
@@ -60,7 +61,7 @@ def test_calibration_package_parameters_parse_runtime_args():
         workers=8,
         n_clones=430,
         target_config_path=TARGET_CONFIG_PATH,
-        target_config_sha256="sha256:target-config",
+        target_config_sha256=TARGET_CONFIG_SHA256,
         target_config_mode="explicit",
         skip_county=True,
         skip_source_impute=True,
@@ -82,7 +83,7 @@ def test_calibration_package_parameters_parse_runtime_args():
         "skip_takeup_rerandomize": False,
         "target_config": TARGET_CONFIG_PATH,
         "target_config_mode": "explicit",
-        "target_config_sha256": "sha256:target-config",
+        "target_config_sha256": TARGET_CONFIG_SHA256,
         "workers": None,
     }
 
@@ -109,6 +110,28 @@ def test_calibration_package_parameters_require_identity_for_config_modes():
         raise AssertionError("Explicit target config mode should require checksum")
 
 
+def test_calibration_package_parameters_reject_malformed_target_config_checksum():
+    try:
+        CalibrationPackageParameters.from_runtime_args(
+            workers=8,
+            n_clones=430,
+            target_config_path=TARGET_CONFIG_PATH,
+            target_config_sha256="sha256:target-config",
+            target_config_mode="explicit",
+            skip_county=True,
+            skip_source_impute=True,
+            skip_takeup_rerandomize=False,
+            chunked_matrix=False,
+            chunk_size=25_000,
+            parallel=False,
+            num_matrix_workers=50,
+        )
+    except ValueError as exc:
+        assert "SHA-256 digest" in str(exc)
+    else:
+        raise AssertionError("Malformed target config checksum should fail")
+
+
 def test_calibration_package_parameters_accept_legacy_identity_fields_missing():
     params = CalibrationPackageParameters.from_dict(
         {
@@ -130,6 +153,31 @@ def test_calibration_package_parameters_accept_legacy_identity_fields_missing():
     assert params.target_config_sha256 is None
 
 
+def test_calibration_package_contract_revalidates_backfilled_identity(tmp_path):
+    dataset_path, db_path, package_path = contract_input_paths(tmp_path)
+    package = calibration_package_payload()
+    package["metadata"].pop("target_config_sha256")
+    write_calibration_package_payload(package_path, package)
+    parameters = calibration_package_parameters()
+    parameters.pop("target_config_mode")
+    parameters.pop("target_config_sha256")
+
+    try:
+        build_calibration_package_contract(
+            package_path=package_path,
+            dataset_path=dataset_path,
+            db_path=db_path,
+            package=package,
+            parameters=parameters,
+            run_id="run-a",
+            completed_at="2026-05-08T12:02:00Z",
+        )
+    except ValueError as exc:
+        assert "target_config and target_config_sha256" in str(exc)
+    else:
+        raise AssertionError("Backfilled target config identity should be revalidated")
+
+
 def test_calibration_package_parameters_reject_inconsistent_chunk_shape():
     try:
         CalibrationPackageParameters(
@@ -252,7 +300,7 @@ def test_calibration_package_contract_records_matrix_summary(tmp_path):
     assert summary["matrix_density"] == 0.5
     assert summary["n_targets"] == 2
     assert summary["target_name_count"] == 2
-    assert summary["target_config_sha256"] == "sha256:target-config"
+    assert summary["target_config_sha256"] == TARGET_CONFIG_SHA256
     assert summary["n_clones"] == 3
     assert summary["seed"] == 42
     assert summary["matrix_builder"] == "chunked"