Merge pull request #986 from PolicyEngine/codex/h5-migration-pr9-coordinator

Coordinate local H5 workers with catalog requests

Author: anth-volk

changelog.d/985.added

@@ -0,0 +1 @@
+Coordinate local H5 publishing through typed catalog requests and normalize worker responses.

docs/engineering/stages/build_outputs.md

@@ -34,6 +34,19 @@ a postprocessor.
 
 ## Worker Chunk Execution
 
+The Modal coordinator builds canonical typed area requests before spawning
+workers. Regional publish reads the target congressional district universe from
+the staged target database through `TargetUniverseReader`, reads only
+coordinator-needed CD and county geography fields through
+`CalibrationGeographyLoader.load_index()`, then asks `USAreaCatalog` to define
+the regional release shape: every configured state, every target congressional
+district, and the explicitly supported city outputs such as NYC. The
+coordinator wraps those requests in `WeightedAreaRequest`, partitions them with
+`partition_weighted_area_requests()`, and sends workers typed
+`--requests-json` payloads. Completion is measured against the explicit request
+keys, not just a raw file count, so stale or unrelated H5 files cannot satisfy a
+missing expected area.
+
 `LocalH5WorkerService` is the reusable Stage 4 boundary for executing one
 prepared local-H5 worker chunk. It consumes a `WorkerSession`, typed
 `AreaBuildRequest` objects, and a `WorkerExecutionConfig`, then returns a
@@ -51,11 +64,18 @@ service. It may parse legacy `--work-items` and typed `--requests-json`, prepare
 the worker session, and print the legacy coordinator JSON shape, but it should
 not regain build-loop, write-loop, or validation-loop logic.
 
+The legacy `--work-items` input path remains compatibility-only while older
+tests and explicit override callers are retired. New coordinator work should
+prefer typed `AreaBuildRequest` objects and typed worker payloads.
+
 For now, `WorkerResult.to_legacy_dict()` preserves the existing coordinator
 contract with `completed`, `failed`, `errors`, `validation_rows`, and
 `validation_summary`. New code should prefer the structured `results` and
-`issues` fields. Validation exceptions remain visible in legacy `errors` so the
-current coordinator does not drop them before it migrates to structured results.
+`issues` fields. Validation exceptions remain visible as structured `issues`
+with `severity="validation"` by default; they enter legacy `errors` only when a
+worker is configured to fail on validation errors or the area result is otherwise
+failed. This keeps nonfatal validation diagnostics visible without treating them
+as coordinator-fatal worker failures.
 Removing the legacy shape and moving the coordinator off worker subprocess JSON
 is a later migration step.
 

modal_app/fixtures/h5_cases.py

@@ -23,6 +23,7 @@
 COUNTY_FIPS = "37183"
 STATE_FIPS = 37
 N_CLONES = 1
+DEFAULT_CATALOG_N_CLONES = 2
 SEED = 42
 VERSION = "0.0.0"
 
@@ -36,6 +37,8 @@ class SeededCase:
     expected_district_name: str = DISTRICT_NAME
     n_clones: int = N_CLONES
     seed: int = SEED
+    expected_output_count: int | None = None
+    expected_output_paths: tuple[str, ...] = ()
 
 
 def _sha256(path: Path) -> str:
@@ -68,20 +71,25 @@ def _write_dataset(artifact_dir: Path) -> Path:
     return dataset_path
 
 
-def _write_weights(artifact_dir: Path, *, n_records: int) -> Path:
+def _write_weights(
+    artifact_dir: Path,
+    *,
+    n_records: int,
+    n_clones: int = N_CLONES,
+) -> Path:
     import numpy as np
 
     weights_path = artifact_dir / "calibration_weights.npy"
-    np.save(weights_path, np.ones(n_records * N_CLONES, dtype=np.float32))
+    np.save(weights_path, np.ones(n_records * n_clones, dtype=np.float32))
     return weights_path
 
 
-def _base_geography(*, n_records: int):
+def _base_geography(*, n_records: int, n_clones: int = N_CLONES):
     import numpy as np
 
     from policyengine_us_data.calibration.clone_and_assign import GeographyAssignment
 
-    total_rows = n_records * N_CLONES
+    total_rows = n_records * n_clones
     block_geoids = np.array(
         [f"{COUNTY_FIPS}{i:06d}{i:04d}"[:15] for i in range(total_rows)],
         dtype="U15",
@@ -92,21 +100,93 @@ def _base_geography(*, n_records: int):
         county_fips=np.full(total_rows, COUNTY_FIPS, dtype="U5"),
         state_fips=np.full(total_rows, STATE_FIPS, dtype=np.int32),
         n_records=n_records,
-        n_clones=N_CLONES,
+        n_clones=n_clones,
+    )
+
+
+def _default_catalog_geography(*, n_records: int, n_clones: int):
+    import numpy as np
+
+    from policyengine_us_data.calibration.calibration_utils import STATE_CODES
+    from policyengine_us_data.calibration.clone_and_assign import GeographyAssignment
+
+    total_rows = n_records * n_clones
+    state_universe = tuple(sorted(STATE_CODES))
+    if total_rows < len(state_universe):
+        raise ValueError(
+            f"Need at least {len(state_universe)} clone rows to cover default states"
+        )
+
+    state_fips = np.array(
+        [state_universe[index % len(state_universe)] for index in range(total_rows)],
+        dtype=np.int32,
+    )
+    county_fips = np.fromiter(
+        (
+            "37183"
+            if state == STATE_FIPS
+            else "36061"
+            if state == 36
+            else f"{state:02d}001"
+            for state in state_fips
+        ),
+        dtype="U5",
+        count=total_rows,
+    )
+    cd_geoid = np.fromiter(
+        (
+            DISTRICT_GEOID
+            if state == STATE_FIPS
+            else "3601"
+            if state == 36
+            else f"{state}01"
+            for state in state_fips
+        ),
+        dtype="U4",
+        count=total_rows,
+    )
+    block_geoid = np.array(
+        [
+            f"{county}{index:06d}{index:04d}"[:15]
+            for index, county in enumerate(county_fips)
+        ],
+        dtype="U15",
+    )
+    return GeographyAssignment(
+        block_geoid=block_geoid,
+        cd_geoid=cd_geoid,
+        county_fips=county_fips,
+        state_fips=state_fips,
+        n_records=n_records,
+        n_clones=n_clones,
     )
 
 
-def _write_saved_geography(artifact_dir: Path, *, n_records: int) -> Path:
+def _write_saved_geography(
+    artifact_dir: Path,
+    *,
+    n_records: int,
+    n_clones: int = N_CLONES,
+    geography=None,
+) -> Path:
     from policyengine_us_data.calibration.clone_and_assign import save_geography
 
     geography_path = artifact_dir / "geography_assignment.npz"
-    save_geography(_base_geography(n_records=n_records), geography_path)
+    save_geography(
+        geography or _base_geography(n_records=n_records, n_clones=n_clones),
+        geography_path,
+    )
     return geography_path
 
 
-def _write_calibration_package(artifact_dir: Path, *, n_records: int) -> Path:
+def _write_calibration_package(
+    artifact_dir: Path,
+    *,
+    n_records: int,
+    n_clones: int = N_CLONES,
+) -> Path:
     package_path = artifact_dir / "calibration_package.pkl"
-    geography = _base_geography(n_records=n_records)
+    geography = _base_geography(n_records=n_records, n_clones=n_clones)
     payload = {
         "block_geoid": geography.block_geoid,
         "cd_geoid": geography.cd_geoid,
@@ -122,9 +202,14 @@ def _write_calibration_package(artifact_dir: Path, *, n_records: int) -> Path:
     return package_path
 
 
-def _write_misnamed_package(artifact_dir: Path, *, n_records: int) -> Path:
+def _write_misnamed_package(
+    artifact_dir: Path,
+    *,
+    n_records: int,
+    n_clones: int = N_CLONES,
+) -> Path:
     wrong_path = artifact_dir / "calibration_package_typo.pkl"
-    geography = _base_geography(n_records=n_records)
+    geography = _base_geography(n_records=n_records, n_clones=n_clones)
     payload = {
         "block_geoid": geography.block_geoid,
         "cd_geoid": geography.cd_geoid,
@@ -192,25 +277,71 @@ def seed_case(
     _reset_dir(staging_dir)
 
     n_records = _fixture_n_households()
+    n_clones = (
+        DEFAULT_CATALOG_N_CLONES if case_name == "default_catalog_success" else N_CLONES
+    )
     dataset_path = _write_dataset(artifact_dir)
-    weights_path = _write_weights(artifact_dir, n_records=n_records)
+    weights_path = _write_weights(
+        artifact_dir,
+        n_records=n_records,
+        n_clones=n_clones,
+    )
     db_path = _write_db(artifact_dir)
 
     geography_path = None
     package_path = None
+    expected_output_count = None
+    expected_output_paths: tuple[str, ...] = ()
 
     if case_name == "saved_geography_success":
-        geography_path = _write_saved_geography(artifact_dir, n_records=n_records)
+        geography_path = _write_saved_geography(
+            artifact_dir,
+            n_records=n_records,
+            n_clones=n_clones,
+        )
+        _write_run_config(
+            artifact_dir,
+            weights_path=weights_path,
+            geography_path=geography_path,
+        )
+    elif case_name == "default_catalog_success":
+        from policyengine_us_data.calibration.calibration_utils import STATE_CODES
+
+        geography_path = _write_saved_geography(
+            artifact_dir,
+            n_records=n_records,
+            n_clones=n_clones,
+            geography=_default_catalog_geography(
+                n_records=n_records,
+                n_clones=n_clones,
+            ),
+        )
+        expected_output_count = len(STATE_CODES) + 2
+        expected_output_paths = (
+            "states/AL.h5",
+            "states/NC.h5",
+            "states/NY.h5",
+            "districts/NC-01.h5",
+            "cities/NYC.h5",
+        )
         _write_run_config(
             artifact_dir,
             weights_path=weights_path,
             geography_path=geography_path,
         )
     elif case_name == "package_fallback_success":
-        package_path = _write_calibration_package(artifact_dir, n_records=n_records)
+        package_path = _write_calibration_package(
+            artifact_dir,
+            n_records=n_records,
+            n_clones=n_clones,
+        )
         _write_run_config(artifact_dir, weights_path=weights_path)
     elif case_name == "misnamed_package":
-        _write_misnamed_package(artifact_dir, n_records=n_records)
+        _write_misnamed_package(
+            artifact_dir,
+            n_records=n_records,
+            n_clones=n_clones,
+        )
         _write_run_config(artifact_dir, weights_path=weights_path)
     else:
         raise ValueError(f"Unknown H5 test case: {case_name}")
@@ -223,7 +354,10 @@ def seed_case(
             database_path=db_path,
             geography_path=geography_path,
             calibration_package_path=package_path,
-            n_clones=N_CLONES,
+            n_clones=n_clones,
             seed=SEED,
         ).to_wire_dict(),
+        n_clones=n_clones,
+        expected_output_count=expected_output_count,
+        expected_output_paths=expected_output_paths,
     )

modal_app/h5_test_harness.py

@@ -117,6 +117,8 @@ def seed_h5_case(run_id: str, case_name: str) -> dict:
         "expected_district_name": seeded.expected_district_name,
         "n_clones": seeded.n_clones,
         "seed": seeded.seed,
+        "expected_output_count": seeded.expected_output_count,
+        "expected_output_paths": list(seeded.expected_output_paths),
     }