Merge pull request #741 from PolicyEngine/codex/release-manifest

Publish canonical release manifests for US data releases

Author: MaxGhenis

.github/workflows/pr.yaml

@@ -143,11 +143,27 @@ jobs:
       HUGGING_FACE_TOKEN: ${{ secrets.HUGGING_FACE_TOKEN }}
     steps:
       - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
       - uses: actions/setup-python@v5
         with:
           python-version: "3.14"
       - run: pip install modal
       - name: Build datasets and run integration tests on Modal
         run: |
+          STAGE_ARGS=""
+          if git diff --name-only origin/main...HEAD | grep -qx 'pyproject.toml'; then
+            VERSION=$(python -c "from pathlib import Path; import tomllib; print(tomllib.load(Path('pyproject.toml').open('rb'))['project']['version'])")
+            STAGE_ARGS="--upload --stage-only --run-id=${VERSION}"
+            {
+              echo "## Release Artifact Staging"
+              echo ""
+              echo "- package version: \`${VERSION}\`"
+              echo "- staged HF prefix: \`staging/${VERSION}/\`"
+              echo "- promote with: \`uv run python policyengine_us_data/storage/upload_completed_datasets.py --promote-only --run-id=${VERSION} --version=${VERSION}\`"
+            } >> "$GITHUB_STEP_SUMMARY"
+          fi
+
           modal run modal_app/data_build.py \
-            --branch=${{ github.head_ref || github.ref_name }}
+            --branch=${{ github.head_ref || github.ref_name }} \
+            ${STAGE_ARGS}

Makefile

@@ -126,7 +126,7 @@ validate-package:
 	python -m policyengine_us_data.calibration.validate_package
 
 publish-local-area:
-	python policyengine_us_data/calibration/publish_local_area.py --upload
+	python -m policyengine_us_data.calibration.promote_local_h5s --local-dir local_area_build
 
 build-h5s:
 	python -m policyengine_us_data.calibration.publish_local_area \

changelog.d/release-manifest-preflight.fixed.md

@@ -0,0 +1 @@
+Preflight release-manifest finalization before promoting staged artifacts so failed finalization checks cannot partially publish a release.

modal_app/data_build.py

@@ -238,6 +238,8 @@ def validate_and_maybe_upload_datasets(
     upload: bool,
     skip_enhanced_cps: bool,
     env: dict,
+    stage_only: bool = False,
+    run_id: str = "",
 ) -> None:
     validation_args = ["--validate-only"]
     if skip_enhanced_cps:
@@ -254,6 +256,10 @@ def validate_and_maybe_upload_datasets(
         upload_args = []
         if skip_enhanced_cps:
             upload_args.append("--no-require-enhanced-cps")
+        if stage_only:
+            upload_args.append("--stage-only")
+        if run_id:
+            upload_args.append(f"--run-id={run_id}")
         run_script(
             "policyengine_us_data/storage/upload_completed_datasets.py",
             args=upload_args,
@@ -375,6 +381,7 @@ def build_datasets(
     clear_checkpoints: bool = False,
     skip_tests: bool = False,
     skip_enhanced_cps: bool = False,
+    stage_only: bool = False,
     run_id: str = "",
 ):
     """Build all datasets with preemption-resilient checkpointing.
@@ -387,6 +394,7 @@ def build_datasets(
         skip_tests: Skip running the test suite (useful for calibration runs).
         skip_enhanced_cps: Skip enhanced_cps.py and small_enhanced_cps.py
             (useful for calibration runs that only need source_imputed H5).
+        stage_only: Upload to HF staging only, without promoting a release.
     """
     setup_gcp_credentials()
 
@@ -673,6 +681,8 @@ def build_datasets(
         upload=upload,
         skip_enhanced_cps=skip_enhanced_cps,
         env=env,
+        stage_only=stage_only,
+        run_id=run_id,
     )
 
     # Clean up checkpoints after successful completion
@@ -689,6 +699,8 @@ def main(
     clear_checkpoints: bool = False,
     skip_tests: bool = False,
     skip_enhanced_cps: bool = False,
+    stage_only: bool = False,
+    run_id: str = "",
 ):
     result = build_datasets.remote(
         upload=upload,
@@ -697,5 +709,7 @@ def main(
         clear_checkpoints=clear_checkpoints,
         skip_tests=skip_tests,
         skip_enhanced_cps=skip_enhanced_cps,
+        stage_only=stage_only,
+        run_id=run_id,
     )
     print(result)

modal_app/local_area.py

@@ -28,8 +28,8 @@
     if _p not in sys.path:
         sys.path.insert(0, _p)
 
-from modal_app.images import cpu_image as image
-from modal_app.resilience import reconcile_run_dir_fingerprint
+from modal_app.images import cpu_image as image  # noqa: E402
+from modal_app.resilience import reconcile_run_dir_fingerprint  # noqa: E402
 
 app = modal.App("policyengine-us-data-local-area")
 
@@ -70,6 +70,170 @@ def setup_repo(branch: str):
     os.chdir("/root/policyengine-us-data")
 
 
+def _build_promote_national_publish_script(
+    *,
+    version: str,
+    run_id: str,
+    rel_paths: list[str],
+) -> str:
+    rel_paths_json = json.dumps(rel_paths)
+    return f"""
+import json
+from pathlib import Path
+from policyengine_us_data.utils.data_upload import (
+    promote_staging_to_production_hf,
+    cleanup_staging_hf,
+    upload_local_area_file,
+    publish_release_manifest_to_hf,
+    should_finalize_local_area_release,
+)
+from policyengine_us_data.utils.version_manifest import (
+    HFVersionInfo,
+    build_manifest,
+    upload_manifest,
+)
+
+version = "{version}"
+run_id = "{run_id}"
+rel_paths = json.loads('''{rel_paths_json}''')
+run_dir = Path("{VOLUME_MOUNT}") / run_id
+
+print(f"Promoting national H5 from staging to production (run_id={{run_id!r}})...")
+promoted = promote_staging_to_production_hf(rel_paths, version, run_id=run_id)
+print(f"Promoted {{promoted}} files to HuggingFace production")
+
+national_h5 = run_dir / "national" / "US.h5"
+if national_h5.exists():
+    print("Uploading national H5 to GCS...")
+    upload_local_area_file(
+        str(national_h5), "national", version=version, skip_hf=True
+    )
+    print("Uploaded national H5 to GCS")
+else:
+    raise RuntimeError(f"Expected national H5 at {{national_h5}}")
+
+print("Updating release manifest...")
+should_finalize, missing_prefixes = should_finalize_local_area_release(
+    version=version,
+    new_repo_paths=["national/US.h5"],
+)
+manifest = publish_release_manifest_to_hf(
+    [(national_h5, "national/US.h5")],
+    version=version,
+    create_tag=should_finalize,
+)
+if should_finalize:
+    upload_manifest(
+        build_manifest(
+            version=version,
+            blob_names=sorted(
+                artifact["path"] for artifact in manifest["artifacts"].values()
+            ),
+            hf_info=HFVersionInfo(
+                repo="policyengine/policyengine-us-data",
+                commit=version,
+            ),
+        )
+    )
+    print("Updated release manifest and created tag")
+else:
+    print(
+        "Updated release manifest without creating a tag; "
+        f"missing prefixes: {{', '.join(missing_prefixes)}}"
+    )
+
+print("Cleaning up staging...")
+cleaned = cleanup_staging_hf(rel_paths, version, run_id=run_id)
+print(f"Cleaned up {{cleaned}} files from staging")
+print(f"Successfully promoted national H5 for version {{version}}")
+"""
+
+
+def _build_promote_publish_script(
+    *,
+    version: str,
+    run_id: str,
+    rel_paths: list[str],
+) -> str:
+    rel_paths_json = json.dumps(rel_paths)
+    return f"""
+import json
+from pathlib import Path
+from policyengine_us_data.utils.data_upload import (
+    promote_staging_to_production_hf,
+    cleanup_staging_hf,
+    upload_local_area_file,
+    publish_release_manifest_to_hf,
+    should_finalize_local_area_release,
+)
+from policyengine_us_data.utils.version_manifest import (
+    HFVersionInfo,
+    build_manifest,
+    upload_manifest,
+)
+
+rel_paths = json.loads('''{rel_paths_json}''')
+version = "{version}"
+run_id = "{run_id}"
+run_dir = Path("{VOLUME_MOUNT}") / run_id
+
+print(f"Promoting {{len(rel_paths)}} files from staging/ to production (run_id={{run_id!r}})...")
+promoted = promote_staging_to_production_hf(rel_paths, version, run_id=run_id)
+print(f"Promoted {{promoted}} files to HuggingFace production")
+
+print(f"Uploading {{len(rel_paths)}} files to GCS...")
+gcs_count = 0
+for rel_path in rel_paths:
+    local_path = run_dir / rel_path
+    subdirectory = str(Path(rel_path).parent)
+    upload_local_area_file(
+        str(local_path),
+        subdirectory,
+        version=version,
+        skip_hf=True,
+    )
+    gcs_count += 1
+print(f"Uploaded {{gcs_count}} files to GCS")
+
+print("Updating release manifest...")
+should_finalize, missing_prefixes = should_finalize_local_area_release(
+    version=version,
+    new_repo_paths=rel_paths,
+)
+manifest = publish_release_manifest_to_hf(
+    [(run_dir / rel_path, rel_path) for rel_path in rel_paths],
+    version=version,
+    create_tag=should_finalize,
+)
+if should_finalize:
+    upload_manifest(
+        build_manifest(
+            version=version,
+            blob_names=sorted(
+                artifact["path"] for artifact in manifest["artifacts"].values()
+            ),
+            hf_info=HFVersionInfo(
+                repo="policyengine/policyengine-us-data",
+                commit=version,
+            ),
+        )
+    )
+    print("Updated release manifest and created tag")
+else:
+    print(
+        "Updated release manifest without final tag; missing local-area prefixes: "
+        + ", ".join(missing_prefixes)
+    )
+    print("Deferring version_manifest.json update until release finalization")
+
+print("Cleaning up staging/...")
+cleaned = cleanup_staging_hf(rel_paths, version, run_id=run_id)
+print(f"Cleaned up {{cleaned}} files from staging/")
+
+print(f"Successfully published version {{version}}")
+"""
+
+
 def validate_artifacts(
     config_path: Path,
     artifact_dir: Path,
@@ -556,52 +720,17 @@ def promote_publish(branch: str = "main", version: str = "", run_id: str = "") -
     with open(manifest_path) as f:
         manifest = json.load(f)
 
-    rel_paths_json = json.dumps(list(manifest["files"].keys()))
-
     result = subprocess.run(
         [
             "uv",
             "run",
             "python",
             "-c",
-            f"""
-import json
-from pathlib import Path
-from policyengine_us_data.utils.data_upload import (
-    promote_staging_to_production_hf,
-    cleanup_staging_hf,
-    upload_local_area_file,
-)
-
-rel_paths = json.loads('''{rel_paths_json}''')
-version = "{version}"
-run_id = "{run_id}"
-run_dir = Path("{VOLUME_MOUNT}") / run_id
-
-print(f"Promoting {{len(rel_paths)}} files from staging/ to production (run_id={{run_id!r}})...")
-promoted = promote_staging_to_production_hf(rel_paths, version, run_id=run_id)
-print(f"Promoted {{promoted}} files to HuggingFace production")
-
-print(f"Uploading {{len(rel_paths)}} files to GCS...")
-gcs_count = 0
-for rel_path in rel_paths:
-    local_path = run_dir / rel_path
-    subdirectory = str(Path(rel_path).parent)
-    upload_local_area_file(
-        str(local_path),
-        subdirectory,
-        version=version,
-        skip_hf=True,
-    )
-    gcs_count += 1
-print(f"Uploaded {{gcs_count}} files to GCS")
-
-print("Cleaning up staging/...")
-cleaned = cleanup_staging_hf(rel_paths, version, run_id=run_id)
-print(f"Cleaned up {{cleaned}} files from staging/")
-
-print(f"Successfully published version {{version}}")
-""",
+            _build_promote_publish_script(
+                version=version,
+                run_id=run_id,
+                rel_paths=list(manifest["files"].keys()),
+            ),
         ],
         text=True,
         env=os.environ.copy(),
@@ -1133,39 +1262,11 @@ def promote_national_publish(
             "run",
             "python",
             "-c",
-            f"""
-import json
-from pathlib import Path
-from policyengine_us_data.utils.data_upload import (
-    promote_staging_to_production_hf,
-    cleanup_staging_hf,
-    upload_local_area_file,
-)
-
-version = "{version}"
-run_id = "{run_id}"
-rel_paths = {json.dumps(rel_paths)}
-run_dir = Path("{VOLUME_MOUNT}") / run_id
-
-print(f"Promoting national H5 from staging to production (run_id={{run_id!r}})...")
-promoted = promote_staging_to_production_hf(rel_paths, version, run_id=run_id)
-print(f"Promoted {{promoted}} files to HuggingFace production")
-
-national_h5 = run_dir / "national" / "US.h5"
-if national_h5.exists():
-    print("Uploading national H5 to GCS...")
-    upload_local_area_file(
-        str(national_h5), "national", version=version, skip_hf=True
-    )
-    print("Uploaded national H5 to GCS")
-else:
-    print(f"WARNING: {{national_h5}} not on volume, skipping GCS")
-
-print("Cleaning up staging...")
-cleaned = cleanup_staging_hf(rel_paths, version, run_id=run_id)
-print(f"Cleaned up {{cleaned}} files from staging")
-print(f"Successfully promoted national H5 for version {{version}}")
-""",
+            _build_promote_national_publish_script(
+                version=version,
+                run_id=run_id,
+                rel_paths=rel_paths,
+            ),
         ],
         text=True,
         env=os.environ.copy(),