Add Stage 2 input artifact bundles

Author: anth-volk

changelog.d/1065.changed

@@ -0,0 +1 @@
+Stage 2 calibration package construction now resolves its inputs and outputs through run-scoped artifact bundles.

docs/pipeline_map.yaml

@@ -805,6 +805,8 @@ stages:
     label: run_calibration()
     description: 'Build phase: resolve targets and constraints, assemble clone values, and package the sparse calibration matrix'
     node_ids:
+    - stage2_input_bundle
+    - stage2_build_context
     - stage2_target_config_identity
     - stage2_target_catalog_load
     - target_resolve
@@ -825,6 +827,10 @@ stages:
     - out_contract
     - stage2_calibration_package_contract_validator
   extra_nodes:
+  - id: in_stage1_contract_s2
+    label: dataset_build_output.json
+    node_type: artifact
+    description: Stage 1 handoff contract preferred for Stage 2 input resolution
   - id: in_cps_s5
     label: source_imputed_stratified_extended_cps.h5
     node_type: artifact
@@ -890,9 +896,34 @@ stages:
     node_type: utility
     description: CSR/COO matrix construction
   edges:
+  - source: in_stage1_contract_s2
+    target: stage2_input_bundle
+    edge_type: data_flow
+    label: preferred input contract
   - source: in_cps_s5
+    target: stage2_input_bundle
+    edge_type: data_flow
+    label: compatibility fallback
+  - source: in_db_s5
+    target: stage2_input_bundle
+    edge_type: external_source
+    label: compatibility fallback
+  - source: stage2_input_bundle
+    target: stage2_build_context
+    edge_type: data_flow
+    label: validated inputs
+  - source: stage2_artifact_specs
+    target: stage2_build_context
+    edge_type: uses_utility
+    label: output bundle paths
+  - source: stage2_build_context
     target: target_resolve
     edge_type: data_flow
+    label: dataset and database paths
+  - source: stage2_build_context
+    target: stage2_calibration_package_writer
+    edge_type: uses_utility
+    label: package output bundle
   - source: in_db_s5
     target: target_resolve
     edge_type: external_source

modal_app/pipeline.py

@@ -93,8 +93,9 @@
 )
 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,
+    Stage2InputBundleError,
     resolve_target_config_identity,
+    stage2_build_context_for_run,
 )
 from policyengine_us_data.utils.error_redaction import (  # noqa: E402
     redacted_bounded_error_text,
@@ -1243,14 +1244,13 @@ def run_pipeline(
             print(f"  Completed in {completed_build_manifest.duration_s}s")
 
         # ── Step 2: Build calibration package ──
+        package_context = stage2_build_context_for_run(PIPELINE_MOUNT, run_id)
+        package_input_validation = package_context.input_bundle.validation_report()
         package_inputs = _artifact_identities(
-            {
-                "dataset": _artifacts_dir(run_id)
-                / "source_imputed_stratified_extended_cps.h5",
-                "database": _artifacts_dir(run_id) / "policy_data.db",
-            }
+            package_context.input_bundle.manifest_inputs
         )
-        package_artifacts = calibration_package_artifact_paths(_artifacts_dir(run_id))
+        package_inputs["input_validation"] = package_input_validation.to_dict()
+        package_artifacts = package_context.output_bundle
         package_parameters = _calibration_package_parameters(
             workers=num_workers,
             n_clones=n_clones,
@@ -1261,6 +1261,18 @@ def run_pipeline(
             parallel_matrix=parallel_matrix,
             num_matrix_workers=num_matrix_workers,
         )
+        if package_input_validation.status != "pass":
+            active_step_manifest = _start_step_manifest(
+                meta,
+                BUILD_CALIBRATION_PACKAGE,
+                parameters=package_parameters,
+                input_identities=package_inputs,
+                vol=pipeline_volume,
+            )
+            raise Stage2InputBundleError(
+                package_context.input_bundle,
+                package_input_validation,
+            )
         package_reuse = _step_reusable(
             meta,
             BUILD_CALIBRATION_PACKAGE,

modal_app/remote_calibration_runner.py

@@ -14,6 +14,7 @@
 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,
+    stage2_build_context_for_run,
 )
 
 app = modal.App(
@@ -371,18 +372,13 @@ def _build_package_impl(
     _ensure_geography_prerequisites()
 
     pipeline_vol.reload()
-    artifacts = f"{PIPELINE_MOUNT}/artifacts"
-    if run_id:
-        artifacts = f"{artifacts}/{run_id}"
-    db_path = f"{artifacts}/policy_data.db"
-    dataset_path = f"{artifacts}/source_imputed_stratified_extended_cps.h5"
-    for label, p in [("database", db_path), ("dataset", dataset_path)]:
-        if not os.path.exists(p):
-            raise RuntimeError(
-                f"Missing {label} on pipeline volume: {p}. Run data_build first."
-            )
-
-    package_artifacts = calibration_package_artifact_paths(artifacts)
+    build_context = stage2_build_context_for_run(
+        PIPELINE_MOUNT, run_id
+    ).require_inputs()
+    input_bundle = build_context.input_bundle
+    package_artifacts = build_context.output_bundle
+    db_path = str(input_bundle.target_database)
+    dataset_path = str(input_bundle.source_dataset)
     pkg_path = str(package_artifacts.package)
     cmd = [
         *_python_cmd("-m", "policyengine_us_data.calibration.unified_calibration"),

policyengine_us_data/calibration_package/__init__.py

@@ -5,25 +5,51 @@
     CALIBRATION_PACKAGE_FILENAME,
     CALIBRATION_PACKAGE_METADATA_FILENAME,
     CALIBRATION_PACKAGE_SUBSTAGE_ID,
+    CALIBRATION_REPORTS_DIRNAME,
+    DATASET_BUILD_OUTPUT_CONTRACT_FILENAME,
     DEFAULT_TARGET_CONFIG_PATH,
     MATRIX_BUILD_DIRNAME,
+    SOURCE_DATASET_FILENAME,
     TARGET_CONFIG_IDENTITY_MODES,
+    TARGET_DATABASE_FILENAME,
     CalibrationPackageArtifactPaths,
+    CalibrationPackageOutputBundle,
+    Stage2BuildContext,
+    Stage2InputBundle,
+    Stage2InputBundleError,
+    Stage2InputSource,
     TargetConfigIdentity,
     calibration_package_artifact_paths,
     resolve_target_config_identity,
+    stage2_build_context_for_run,
+    stage2_input_bundle_from_artifacts_dir,
+    stage2_input_bundle_from_stage1_contract,
+    stage2_input_bundle_from_stage1_contract_path,
 )
 
 __all__ = [
     "CALIBRATION_PACKAGE_CONTRACT_FILENAME",
     "CALIBRATION_PACKAGE_FILENAME",
     "CALIBRATION_PACKAGE_METADATA_FILENAME",
     "CALIBRATION_PACKAGE_SUBSTAGE_ID",
+    "CALIBRATION_REPORTS_DIRNAME",
+    "DATASET_BUILD_OUTPUT_CONTRACT_FILENAME",
     "DEFAULT_TARGET_CONFIG_PATH",
     "MATRIX_BUILD_DIRNAME",
+    "SOURCE_DATASET_FILENAME",
     "TARGET_CONFIG_IDENTITY_MODES",
+    "TARGET_DATABASE_FILENAME",
     "CalibrationPackageArtifactPaths",
+    "CalibrationPackageOutputBundle",
+    "Stage2BuildContext",
+    "Stage2InputBundle",
+    "Stage2InputBundleError",
+    "Stage2InputSource",
     "TargetConfigIdentity",
     "calibration_package_artifact_paths",
     "resolve_target_config_identity",
+    "stage2_build_context_for_run",
+    "stage2_input_bundle_from_artifacts_dir",
+    "stage2_input_bundle_from_stage1_contract",
+    "stage2_input_bundle_from_stage1_contract_path",
 ]