Add Stage 2 package identity specs

Author: anth-volk

changelog.d/1041.changed

@@ -0,0 +1 @@
+Stage 2 calibration package manifests now track the explicit target config identity and contract artifact path.

docs/pipeline_map.yaml

@@ -805,15 +805,25 @@ stages:
     label: run_calibration()
     description: 'Build phase: resolve targets and constraints, assemble clone values, and package the sparse calibration matrix'
     node_ids:
+    - stage2_target_config_identity
+    - stage2_target_catalog_load
     - target_resolve
+    - stage2_target_config_apply
     - target_uprate
     - geo_build
     - constraint_resolve
     - state_precomp
     - clone_assembly
     - takeup_rerand
     - sparse_build
+    - build_matrix
+    - build_matrix_chunked
+    - stage2_artifact_specs
+    - stage2_calibration_package_writer
     - out_pkg
+    - stage2_calibration_package_contract_writer
+    - out_contract
+    - stage2_calibration_package_contract_validator
   extra_nodes:
   - id: in_cps_s5
     label: source_imputed_stratified_extended_cps.h5
@@ -859,6 +869,10 @@ stages:
     label: calibration_package.pkl
     node_type: artifact
     description: X_sparse CSR matrix, targets_df, initial_weights, metadata
+  - id: out_contract
+    label: calibration_package_contract.json
+    node_type: artifact
+    description: Stage 2 package handoff contract written next to calibration_package.pkl
   - id: util_sql
     label: sqlalchemy
     node_type: utility
@@ -884,12 +898,25 @@ stages:
     edge_type: external_source
     label: SQL targets
   - source: in_config_s5
-    target: target_resolve
+    target: stage2_target_config_identity
+    edge_type: data_flow
+    label: config file
+  - source: stage2_target_config_identity
+    target: stage2_target_catalog_load
+    edge_type: data_flow
+    label: resolved path and checksum
+  - source: stage2_target_catalog_load
+    target: stage2_target_config_apply
     edge_type: data_flow
-    label: include list
+    label: include/exclude rules
   - source: target_resolve
+    target: stage2_target_config_apply
+    edge_type: data_flow
+    label: candidate targets
+  - source: stage2_target_config_apply
     target: target_uprate
     edge_type: data_flow
+    label: selected targets
   - source: target_uprate
     target: geo_build
     edge_type: data_flow
@@ -917,8 +944,48 @@ stages:
     target: sparse_build
     edge_type: data_flow
   - source: sparse_build
+    target: build_matrix
+    edge_type: uses_library
+    label: non-chunked path
+  - source: sparse_build
+    target: build_matrix_chunked
+    edge_type: uses_library
+    label: chunked path
+  - source: build_matrix
+    target: stage2_calibration_package_writer
+    edge_type: data_flow
+  - source: build_matrix_chunked
+    target: stage2_calibration_package_writer
+    edge_type: data_flow
+  - 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: out_pkg
+    target: stage2_calibration_package_contract_writer
+    edge_type: data_flow
+  - source: stage2_artifact_specs
+    target: stage2_calibration_package_contract_writer
+    edge_type: uses_utility
+    label: contract path
+  - source: stage2_calibration_package_contract_writer
+    target: out_contract
+    edge_type: produces_artifact
+  - source: out_pkg
+    target: stage2_calibration_package_contract_validator
+    edge_type: validates
+  - source: out_contract
+    target: stage2_calibration_package_contract_validator
+    edge_type: validates
+  - source: in_cps_s5
+    target: stage2_calibration_package_contract_validator
+    edge_type: validates
+  - source: in_db_s5
+    target: stage2_calibration_package_contract_validator
+    edge_type: validates
   - source: util_sql
     target: target_resolve
     edge_type: uses_utility

modal_app/pipeline.py

@@ -92,6 +92,10 @@
     write_run_meta,
 )
 from policyengine_us_data.utils.run_context import RunContext, resolve_run_id  # noqa: E402
+from policyengine_us_data.calibration_package.specs import (  # noqa: E402
+    calibration_package_artifact_paths,
+    resolve_target_config_identity,
+)
 from policyengine_us_data.utils.error_redaction import (  # noqa: E402
     redacted_bounded_error_text,
     redact_error_text,
@@ -162,18 +166,25 @@ def _calibration_package_parameters(
     workers: int,
     n_clones: int,
     target_config: str | None,
+    all_active_targets: bool = False,
     skip_county: bool,
     chunked_matrix: bool,
     chunk_size: int,
     parallel_matrix: bool,
     num_matrix_workers: int,
 ) -> dict:
     """Return manifest parameters that affect package construction."""
+    target_config_identity = resolve_target_config_identity(
+        target_config,
+        all_active_targets=all_active_targets,
+    )
     effective_parallel = bool(chunked_matrix and parallel_matrix)
     params = {
         "workers": workers if not chunked_matrix else None,
         "n_clones": n_clones,
-        "target_config": target_config,
+        "target_config": target_config_identity.path,
+        "target_config_sha256": target_config_identity.sha256,
+        "target_config_mode": target_config_identity.mode,
         "skip_county": skip_county,
         "chunked_matrix": bool(chunked_matrix),
         "chunk_size": chunk_size if chunked_matrix else None,
@@ -547,6 +558,7 @@ def verify_runtime_seams() -> dict:
         "modal_app/step_manifests/errors.py",
         "modal_app/step_manifests/status.py",
         "modal_app/fixtures/h5_cases.py",
+        "policyengine_us_data/calibration_package/specs.py",
         "tests/integration/test_fixture_50hh.h5",
         "policyengine_us_data/calibration/target_config.yaml",
         "policyengine_us_data/calibration/target_config_full.yaml",
@@ -1238,6 +1250,7 @@ def run_pipeline(
                 "database": _artifacts_dir(run_id) / "policy_data.db",
             }
         )
+        package_artifacts = calibration_package_artifact_paths(_artifacts_dir(run_id))
         package_parameters = _calibration_package_parameters(
             workers=num_workers,
             n_clones=n_clones,
@@ -1302,8 +1315,7 @@ def run_pipeline(
             completed_package_manifest = _complete_step_manifest(
                 active_step_manifest,
                 outputs=collect_artifacts(
-                    [_artifacts_dir(run_id) / "calibration_package.pkl"],
-                    missing_ok=True,
+                    package_artifacts.manifest_outputs,
                 ),
                 vol=pipeline_volume,
             )

modal_app/remote_calibration_runner.py

@@ -12,6 +12,9 @@
         sys.path.insert(0, _p)
 
 from modal_app.images import gpu_image as image  # noqa: E402
+from policyengine_us_data.calibration_package.specs import (  # noqa: E402
+    calibration_package_artifact_paths,
+)
 
 app = modal.App(
     os.environ.get("US_DATA_FIT_WEIGHTS_APP_NAME") or "policyengine-us-data-fit-weights"
@@ -379,7 +382,8 @@ def _build_package_impl(
                 f"Missing {label} on pipeline volume: {p}. Run data_build first."
             )
 
-    pkg_path = f"{artifacts}/calibration_package.pkl"
+    package_artifacts = calibration_package_artifact_paths(artifacts)
+    pkg_path = str(package_artifacts.package)
     cmd = [
         *_python_cmd("-m", "policyengine_us_data.calibration.unified_calibration"),
         "--device",
@@ -404,7 +408,7 @@ def _build_package_impl(
     if chunked_matrix:
         cmd.extend(["--chunked-matrix", "--chunk-size", str(chunk_size)])
         if parallel_matrix:
-            chunk_dir = f"{artifacts}/matrix_build"
+            chunk_dir = str(package_artifacts.matrix_build_dir)
             cmd.extend(
                 [
                     "--parallel",
@@ -439,14 +443,12 @@ def _build_package_impl(
         raise RuntimeError(f"Package build failed with code {build_rc}")
 
     from policyengine_us_data.stage_contracts.calibration_package import (
-        CALIBRATION_PACKAGE_CONTRACT_FILENAME,
         validate_persisted_calibration_package_contract,
     )
 
-    contract_path = f"{artifacts}/{CALIBRATION_PACKAGE_CONTRACT_FILENAME}"
     validate_persisted_calibration_package_contract(
-        package_path=Path(pkg_path),
-        contract_path=Path(contract_path),
+        package_path=package_artifacts.package,
+        contract_path=package_artifacts.contract,
         dataset_path=Path(dataset_path),
         db_path=Path(db_path),
     )
@@ -525,8 +527,9 @@ def check_volume_package(artifacts_dir: str = "") -> dict:
     import json
 
     base = artifacts_dir if artifacts_dir else f"{PIPELINE_MOUNT}/artifacts"
-    pkg_path = f"{base}/calibration_package.pkl"
-    sidecar_path = f"{base}/calibration_package_meta.json"
+    package_artifacts = calibration_package_artifact_paths(base)
+    pkg_path = str(package_artifacts.package)
+    sidecar_path = str(package_artifacts.metadata)
     if not os.path.exists(pkg_path):
         return {"exists": False}