Merge pull request #1043 from PolicyEngine/agent/stage-3/pr-3a-specs-artifact-identity

Centralize Stage 3 fit specs and artifacts

Author: anth-volk

changelog.d/1040.changed

@@ -0,0 +1 @@
+Centralized Stage 3 fitted-weight specs and artifact filenames.

docs/engineering/stages/fit_weights.md

@@ -0,0 +1,23 @@
+# Stage 3: Fit Weights
+
+Stage 3 produces scoped fitted-weight artifacts for regional and national H5
+builds. The public identity boundary lives in `policyengine_us_data.fit_weights`:
+
+- `FitScope` names the durable regional and national scopes.
+- `FittedWeightsSpec` defines the scoped optimization parameters recorded in
+  step manifests for reuse decisions.
+- `ScopedFitArtifacts` defines the artifact filenames written by the Modal fit
+  step and consumed by downstream H5 builders.
+
+The current artifact names remain behavior-compatible:
+
+- regional: `calibration_weights.npy`, `geography_assignment.npz`,
+  `unified_run_config.json`, `unified_diagnostics.csv`, and
+  `calibration_log.csv`;
+- national: `national_calibration_weights.npy`,
+  `national_geography_assignment.npz`, `national_unified_run_config.json`,
+  `national_unified_diagnostics.csv`, and `national_calibration_log.csv`.
+
+When changing Stage 3 fitting parameters, artifact names, or scope behavior,
+update the central specs first and then adapt Modal callers to consume those
+specs. Do not add parallel filename constants in orchestration code.

docs/pipeline_map.yaml

@@ -947,6 +947,8 @@ stages:
     label: run_calibration()
     description: 'Fit phase: initialize weights, optimize sparse calibration weights, and emit artifacts plus diagnostics'
     node_ids:
+    - fit_spec_regional
+    - fit_artifacts_regional
     - init_weights
     - create_model
     - fit_model
@@ -964,6 +966,14 @@ stages:
     label: Modal GPU Container
     node_type: external
     description: T4 / A10 / A100 / H100 - 32GB RAM, 8 CPU
+  - id: fit_spec_regional
+    label: FittedWeightsSpec regional
+    node_type: library
+    description: Regional Stage 3 fit hyperparameters and deterministic reuse identity
+  - id: fit_artifacts_regional
+    label: ScopedFitArtifacts regional
+    node_type: library
+    description: Regional fitted-weight artifact filenames and remote result mapping
   - id: create_model
     label: Create SparseCalibrationWeights
     node_type: process
@@ -977,7 +987,7 @@ stages:
     node_type: artifact
     description: 'Shape: (n_records x n_clones) - most entries zero'
   - id: out_geo_s6
-    label: geography.npz
+    label: geography_assignment.npz
     node_type: artifact
     description: block_geoid, cd_geoid, county_fips, state_fips
   - id: out_diag
@@ -1000,6 +1010,21 @@ stages:
   - source: in_pkg_s6
     target: init_weights
     edge_type: data_flow
+  - source: fit_spec_regional
+    target: fit_model
+    edge_type: uses_library
+  - source: fit_artifacts_regional
+    target: out_weights
+    edge_type: documents
+  - source: fit_artifacts_regional
+    target: out_geo_s6
+    edge_type: documents
+  - source: fit_artifacts_regional
+    target: out_diag
+    edge_type: documents
+  - source: fit_artifacts_regional
+    target: out_config_s6
+    edge_type: documents
   - source: init_weights
     target: create_model
     edge_type: data_flow
@@ -1047,6 +1072,8 @@ stages:
     label: run_calibration() national
     description: 'National fit phase: initialize national weights, optimize sparse calibration weights, and emit national artifacts plus diagnostics'
     node_ids:
+    - fit_spec_national
+    - fit_artifacts_national
     - init_weights
     - create_model_national
     - fit_model
@@ -1064,6 +1091,14 @@ stages:
     label: Modal GPU Container
     node_type: external
     description: T4 / A10 / A100 / H100 for national calibration
+  - id: fit_spec_national
+    label: FittedWeightsSpec national
+    node_type: library
+    description: National Stage 3 fit hyperparameters and deterministic reuse identity
+  - id: fit_artifacts_national
+    label: ScopedFitArtifacts national
+    node_type: library
+    description: National fitted-weight artifact filenames and remote result mapping
   - id: create_model_national
     label: Create National SparseCalibrationWeights
     node_type: process
@@ -1100,6 +1135,21 @@ stages:
   - source: in_pkg_national_s6
     target: init_weights
     edge_type: data_flow
+  - source: fit_spec_national
+    target: fit_model
+    edge_type: uses_library
+  - source: fit_artifacts_national
+    target: out_national_weights
+    edge_type: documents
+  - source: fit_artifacts_national
+    target: out_national_geo_s6
+    edge_type: documents
+  - source: fit_artifacts_national
+    target: out_national_diag
+    edge_type: documents
+  - source: fit_artifacts_national
+    target: out_national_config_s6
+    edge_type: documents
   - source: init_weights
     target: create_model_national
     edge_type: data_flow

modal_app/local_area.py

@@ -55,6 +55,10 @@
 from policyengine_us_data.build_outputs.worker_inputs import (  # noqa: E402
     WorkerCalibrationInputs,
 )
+from policyengine_us_data.fit_weights import (  # noqa: E402
+    FitScope,
+    fit_artifacts_for_scope,
+)
 from policyengine_us_data.pipeline_metadata import pipeline_node  # noqa: E402
 from policyengine_us_data.pipeline_schema import PipelineNode  # noqa: E402
 from policyengine_us_data.utils.run_context import (  # noqa: E402
@@ -1310,11 +1314,12 @@ def coordinate_publish(
     artifacts = (
         Path(f"/pipeline/artifacts/{run_id}") if run_id else Path("/pipeline/artifacts")
     )
-    weights_path = artifacts / "calibration_weights.npy"
-    geography_path = artifacts / "geography_assignment.npz"
+    regional_fit_artifacts = fit_artifacts_for_scope(FitScope.REGIONAL)
+    weights_path = artifacts / regional_fit_artifacts.weights.filename
+    geography_path = artifacts / regional_fit_artifacts.geography.filename
     db_path = artifacts / "policy_data.db"
     dataset_path = artifacts / "source_imputed_stratified_extended_cps.h5"
-    config_json_path = artifacts / "unified_run_config.json"
+    config_json_path = artifacts / regional_fit_artifacts.run_config.filename
     calibration_package_path = artifacts / "calibration_package.pkl"
 
     required = {
@@ -1609,11 +1614,13 @@ def coordinate_national_publish(
     artifacts = (
         Path(f"/pipeline/artifacts/{run_id}") if run_id else Path("/pipeline/artifacts")
     )
-    weights_path = artifacts / "national_calibration_weights.npy"
-    geography_path = artifacts / "national_geography_assignment.npz"
+    regional_fit_artifacts = fit_artifacts_for_scope(FitScope.REGIONAL)
+    national_fit_artifacts = fit_artifacts_for_scope(FitScope.NATIONAL)
+    weights_path = artifacts / national_fit_artifacts.weights.filename
+    geography_path = artifacts / national_fit_artifacts.geography.filename
     db_path = artifacts / "policy_data.db"
     dataset_path = artifacts / "source_imputed_stratified_extended_cps.h5"
-    config_json_path = artifacts / "national_unified_run_config.json"
+    config_json_path = artifacts / national_fit_artifacts.run_config.filename
 
     required = {
         "weights": weights_path,
@@ -1641,8 +1648,8 @@ def coordinate_national_publish(
         config_json_path,
         artifacts,
         filename_remap={
-            "calibration_weights.npy": "national_calibration_weights.npy",
-            "geography_assignment.npz": "national_geography_assignment.npz",
+            regional_fit_artifacts.weights.filename: national_fit_artifacts.weights.filename,
+            regional_fit_artifacts.geography.filename: national_fit_artifacts.geography.filename,
         },
     )
     fingerprint_inputs = _build_publishing_input_bundle(