From 72729394a9560756c9b6d762315de79864f231cd Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 25 Jun 2026 19:41:45 +0200 Subject: [PATCH 1/3] Certify US state datasets in bundle manifest --- ...1-703bd81a565c-20260620T201958Z.changed.md | 1 + scripts/bundle.py | 21 ++ scripts/certify_data_release.py | 21 ++ src/policyengine/data/bundle/manifest.json | 317 +++++++++++++++++- .../data/bundle/us.trace.tro.jsonld | 9 +- src/policyengine/provenance/certification.py | 161 +++++++++ tests/test_certify_data_release.py | 179 +++++++++- tests/test_us_regions.py | 77 ++++- 8 files changed, 772 insertions(+), 14 deletions(-) create mode 100644 changelog.d/certify-us-populace-us-2024-f0af251-703bd81a565c-20260620T201958Z.changed.md diff --git a/changelog.d/certify-us-populace-us-2024-f0af251-703bd81a565c-20260620T201958Z.changed.md b/changelog.d/certify-us-populace-us-2024-f0af251-703bd81a565c-20260620T201958Z.changed.md new file mode 100644 index 00000000..77f5dcc3 --- /dev/null +++ b/changelog.d/certify-us-populace-us-2024-f0af251-703bd81a565c-20260620T201958Z.changed.md @@ -0,0 +1 @@ +Certify the US populace data release `populace-us-2024-f0af251-703bd81a565c-20260620T201958Z` (populace_us_2024, policyengine-us 1.729.0) into the PolicyEngine bundle manifest, including inherited state datasets from policyengine-us-data 1.115.5. diff --git a/scripts/bundle.py b/scripts/bundle.py index a94b26fb..6d402cd4 100644 --- a/scripts/bundle.py +++ b/scripts/bundle.py @@ -38,6 +38,12 @@ def _certify_data(args: argparse.Namespace) -> int: argv.extend(["--data-producer", args.data_producer]) if args.model_version: argv.extend(["--model-version", args.model_version]) + if args.regional_manifest_uri: + argv.extend(["--regional-manifest-uri", args.regional_manifest_uri]) + if args.regional_artifact_prefix: + argv.extend(["--regional-artifact-prefix", args.regional_artifact_prefix]) + if args.regional_path_template: + argv.extend(["--regional-path-template", args.regional_path_template]) if args.no_generate: argv.append("--no-generate") if args.no_changelog: @@ -141,6 +147,21 @@ def _parser() -> argparse.ArgumentParser: "--model-version", help="Model package version to certify for. Defaults to installed metadata.", ) + certify.add_argument( + "--regional-manifest-uri", + help=( + "Optional regional release_manifest.json URI to merge into US " + "Populace certification." + ), + ) + certify.add_argument( + "--regional-artifact-prefix", + help="Regional artifact prefix to import. Defaults to states/.", + ) + certify.add_argument( + "--regional-path-template", + help="Region dataset path template to certify.", + ) certify.add_argument( "--no-generate", action="store_true", diff --git a/scripts/certify_data_release.py b/scripts/certify_data_release.py index 5da26f83..8da7771a 100644 --- a/scripts/certify_data_release.py +++ b/scripts/certify_data_release.py @@ -49,6 +49,24 @@ def main(argv=None) -> int: required=True, help="hf://dataset/@/", ) + parser.add_argument( + "--regional-manifest-uri", + default=None, + help=( + "Optional regional data release manifest to merge into US Populace " + "certification." + ), + ) + parser.add_argument( + "--regional-artifact-prefix", + default="states/", + help="Regional artifact path prefix to import. Defaults to states/.", + ) + parser.add_argument( + "--regional-path-template", + default="states/{state_code}.h5", + help="Region dataset path template to certify into the bundle manifest.", + ) parser.add_argument( "--model-version", default=None, @@ -78,6 +96,9 @@ def main(argv=None) -> int: / "bundle" / "manifest.json", check_artifacts=not args.skip_artifact_check, + regional_manifest_uri=args.regional_manifest_uri, + regional_artifact_prefix=args.regional_artifact_prefix, + regional_path_template=args.regional_path_template, ) print(result.summary()) diff --git a/src/policyengine/data/bundle/manifest.json b/src/policyengine/data/bundle/manifest.json index 28c7ee32..ef28de63 100644 --- a/src/policyengine/data/bundle/manifest.json +++ b/src/policyengine/data/bundle/manifest.json @@ -157,6 +157,312 @@ "revision": "populace-us-2024-f0af251-703bd81a565c-20260620T201958Z", "sha256": "c577e0b1419f9c12de0cad7a4b3bf32507e718f5ea9d79bc3a32d0c48160cc89" }, + "states/AK": { + "path": "states/AK.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "758b84f75d167ebf35b529c7344f6b6154a86252a68780624d4542c436bf3903" + }, + "states/AL": { + "path": "states/AL.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "7d5555a154e0f4f4bd7c9677cb6473fec69f260c9d1ddae2e2cd2ee2febbcd8a" + }, + "states/AR": { + "path": "states/AR.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "a18d787140d33df51f8f31fe6893892e7228ad2679b1b16011f06d46eb34aedd" + }, + "states/AZ": { + "path": "states/AZ.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "dc047de094fda3f6c61b9ab0b0f73f1e4d481b2b6d5a773abd513653514e0dfd" + }, + "states/CA": { + "path": "states/CA.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "9d7f271cc1a3c84222e284b14cdd749f5364d16a598e1e3693053ce6c580e954" + }, + "states/CO": { + "path": "states/CO.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "7eae1034e34eacd91dcc439a951d8777606fb13b97093a5bf1553bf805445dc0" + }, + "states/CT": { + "path": "states/CT.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "5e5ab0aa48b4dba8fe4879829717461ea6e668637faee9eaae5d3114f45ba2f3" + }, + "states/DC": { + "path": "states/DC.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "51b8cdeecab13d45588206139ae4d106cc972b83d846a1a0f5becb6876707d93" + }, + "states/DE": { + "path": "states/DE.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "186e9b90a11413ae0459de26a799029eb74d4d7ea2bd7625031fd3e7a1f0bf98" + }, + "states/FL": { + "path": "states/FL.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "84bfcecffc59a7c892b4929bf5b4a150e122470a09ebc7643d374726d17057e9" + }, + "states/GA": { + "path": "states/GA.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "f2a3edbb813a43ec17d189e86bb8b087b51f44be2de66cdcbbd933f91a21eed0" + }, + "states/HI": { + "path": "states/HI.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "5ae626eac48714bbdf18ded2485b767e37a9d842bc631035400b455e015dd218" + }, + "states/IA": { + "path": "states/IA.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "6eda342bff2a146af371d1e05f9b47bbc083100a2987a20ef5effcd282017cd2" + }, + "states/ID": { + "path": "states/ID.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "bd43c3dce2c602abf71ab064dce76db0d99846ca4db20d78a78373377aab4201" + }, + "states/IL": { + "path": "states/IL.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "f3840cc2833f3c8f7975631d6bda07a9c27a81dd21c986abf6c831066d6880eb" + }, + "states/IN": { + "path": "states/IN.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "8319b0232e8883e3dc486e98888339a4eff84d22f7b12cae62e54079bc4857d3" + }, + "states/KS": { + "path": "states/KS.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "e9256e79bc0daaa6cb2965ef65768ec336fd9d8c09b449242516832a5bf245df" + }, + "states/KY": { + "path": "states/KY.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "ab6c15f006c0c2f9f66dd9925fc887bfe261a0cff690d40cb09665f4983e89ee" + }, + "states/LA": { + "path": "states/LA.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "d159f6c358019a0ddeeec16072efb7720c2c5efd0aca381f8e0d48ac6aa8ecd9" + }, + "states/MA": { + "path": "states/MA.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "e5b6c0fff3c638185d1c02adbca8aab74359bb92c9969a98af8046c926faf91d" + }, + "states/MD": { + "path": "states/MD.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "80b3a88c36f441dd9d5af24ce1649448eb74f11fa9c68be14da2169742babd7f" + }, + "states/ME": { + "path": "states/ME.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "683c320f238b7e99cecf660194f09c60486beb5b1bc919405b7fffd2a7d19314" + }, + "states/MI": { + "path": "states/MI.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "f1e8220bf6420402b1ae0efe0f745c5b997a3c7a009f4a4f47cd49caa3fa1208" + }, + "states/MN": { + "path": "states/MN.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "481db29ecc2128b59f4190b302dc284afdcbca95d221a559153d565cf2919a11" + }, + "states/MO": { + "path": "states/MO.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "6b61f2dc508fbf9c4b22d0d3054a263c37f31b797b482fd25122ce46a97286d5" + }, + "states/MS": { + "path": "states/MS.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "52c940f80de66ab143df2d9259c36140eb71ce93e4fede8139596be5ad6cf5ff" + }, + "states/MT": { + "path": "states/MT.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "14a208de0e3d97ad95f0ff979ff7fd594b4495f12e6cbbe432a16bb3a1e0cede" + }, + "states/NC": { + "path": "states/NC.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "c4418f5396fc2201f0100d1253f321ebad2d808c174d96b01c36230b43e31d54" + }, + "states/ND": { + "path": "states/ND.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "c5761dc56e30460b50176f088364d95d221eb305d7e9128f028e8f17eb36ca83" + }, + "states/NE": { + "path": "states/NE.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "e187ca6652931ba0efa402bd0570bd714bd789774b7d16e7215fe27c16246132" + }, + "states/NH": { + "path": "states/NH.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "695323c9c07b4cff5f49c7a76e40cb05476e3f464bef0e700f79f46cb6334326" + }, + "states/NJ": { + "path": "states/NJ.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "d1818cac60caed75ce7515715ec3b52cff886ecbd13123bfaeb909241c37ac16" + }, + "states/NM": { + "path": "states/NM.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "c4f561151751f4bf11189c1cabe7988e22bc15e5eb1a0dea059de5685e817989" + }, + "states/NV": { + "path": "states/NV.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "4db82bda4fb9c0a37304c9e0424e30341389f8addb7bbbdb01316066b8332cf9" + }, + "states/NY": { + "path": "states/NY.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "4a1476e298c552a673b88a29c11f12210511ba188291dc15001c5a71d83f437c" + }, + "states/OH": { + "path": "states/OH.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "f6984d9295d00e0e9a7b84c72a0ccb1a231598a5e9f5ff744e8326e4119cca77" + }, + "states/OK": { + "path": "states/OK.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "a92e61e7d445e9d757a0c052afcb31868621882c6b60b3cbd4ef35354bdcf04c" + }, + "states/OR": { + "path": "states/OR.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "5758866cdf930f8312f51c656b7e6ce88cd2877f81e73a27348828dc152948ce" + }, + "states/PA": { + "path": "states/PA.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "c1a7056b6f424cc4c9e847c1cf20395a0ec202a2e0f6d17c46c3b42bb6b5a6d2" + }, + "states/RI": { + "path": "states/RI.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "d1e0bde70b9a760e1963d3481397920a8b73114b22b6b1493afbfc04cb9a7c09" + }, + "states/SC": { + "path": "states/SC.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "8a7bb8d513d73cbddbfab5325d02ca94e43157d0b9b14758c74e0efe58253a17" + }, + "states/SD": { + "path": "states/SD.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "ccd5c65a96ed73e1ddd840556bf6fdf10713d796a4e99f09e228490917ffebad" + }, + "states/TN": { + "path": "states/TN.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "84addf0af5364750b22734246e838fffecdf9a6ed08b1bdbd1d6ca4a76e3be3d" + }, + "states/TX": { + "path": "states/TX.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "27578258f5998b9f3dceefea04d21c5fde9fd5b8e05d80b7b5198fb5b9db924e" + }, + "states/UT": { + "path": "states/UT.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "e165266d283105549d5395409ad6bdf02c600821ce464f95bc4f0ef8d5365b38" + }, + "states/VA": { + "path": "states/VA.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "fb1c3098639521ef744c926327c1ed2ed9f6dd17a3bee637667fb9fb6d5d9b53" + }, + "states/VT": { + "path": "states/VT.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "16c80a27f8a3bf2e832219e9471731a6de1f4da7bb08e770dab07def180dd8bb" + }, + "states/WA": { + "path": "states/WA.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "75723f11aebf83c9867312b422097a273e3f180e033ba2ec23cc4102bf3dc1a6" + }, + "states/WI": { + "path": "states/WI.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "fd1171e43cfcd510b4dcd675849e6d4a6afeb0057956123087a083db83425e5d" + }, + "states/WV": { + "path": "states/WV.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "8b0e718d4a91a3acceedc5a630ea6fd845dad1896b5dfea6baf26786f7b531fe" + }, + "states/WY": { + "path": "states/WY.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + "sha256": "731d83ae37863ff994df2f953740ddb10b36910f43af01b38d36ffb55a88d4b5" + }, "us_source_coverage": { "path": "releases/populace-us-2024-f0af251-703bd81a565c-20260620T201958Z/us_source_coverage.json", "repo_id": "policyengine/populace-us", @@ -173,7 +479,16 @@ "wheel_url": "https://files.pythonhosted.org/packages/b9/7d/778f92ae94997b00c3c9ac34b345f6c9333435f905670ee4eeb2f5e19809/policyengine_us-1.729.0-py3-none-any.whl" }, "policyengine_version": "4.18.2", - "region_datasets": {}, + "region_datasets": { + "national": { + "path_template": "populace_us_2024.h5" + }, + "state": { + "path_template": "states/{state_code}.h5" + } + }, + "regional_release_manifest_uri": "https://huggingface.co/policyengine/policyengine-us-data/resolve/1.115.5/releases/1.115.5/release_manifest.json", + "regional_source_manifest_uri": "hf://model/policyengine/policyengine-us-data@1.115.5/releases/1.115.5/release_manifest.json", "release_manifest_uri": "https://huggingface.co/datasets/policyengine/populace-us/resolve/populace-us-2024-f0af251-703bd81a565c-20260620T201958Z/releases/populace-us-2024-f0af251-703bd81a565c-20260620T201958Z/release_manifest.json", "schema_version": 1, "source_manifest_uri": "hf://dataset/policyengine/populace-us@populace-us-2024-f0af251-703bd81a565c-20260620T201958Z/releases/populace-us-2024-f0af251-703bd81a565c-20260620T201958Z/release_manifest.json", diff --git a/src/policyengine/data/bundle/us.trace.tro.jsonld b/src/policyengine/data/bundle/us.trace.tro.jsonld index ee62b97e..727ef601 100644 --- a/src/policyengine/data/bundle/us.trace.tro.jsonld +++ b/src/policyengine/data/bundle/us.trace.tro.jsonld @@ -75,7 +75,7 @@ "@type": "trov:ResearchArtifact", "schema:name": "policyengine.py bundle manifest for us", "trov:mimeType": "application/json", - "trov:sha256": "8c8687af30e154d29160608e179dfba9892afaff478fea1568ff0e5dafca2bea" + "trov:sha256": "974e5bcd2e8f28ba79b4069c0b6b539ed872a49fb46b44a1d2e013a20755600d" }, { "@id": "composition/1/artifact/data_release_manifest", @@ -102,7 +102,7 @@ "trov:hasFingerprint": { "@id": "composition/1/fingerprint", "@type": "trov:CompositionFingerprint", - "trov:sha256": "fa9d2d6f8676054e67c6faa96d83e2361fbf24789a61c16b7555cd081073224b" + "trov:sha256": "690a3ec10636004a6801f019bf6fa32795a5030a265484d9e99054a712780f9d" } }, "trov:hasPerformance": { @@ -111,12 +111,9 @@ "pe:builtWithModelVersion": "1.729.0", "pe:certifiedBy": "policyengine.py bundle certification", "pe:certifiedForModelVersion": "1.729.0", - "pe:ciGitRef": "refs/heads/main", - "pe:ciGitSha": "819d18a1a665cc23e3f50cab878d1855579e6233", - "pe:ciRunUrl": "https://github.com/PolicyEngine/policyengine.py/actions/runs/28057655642", "pe:compatibilityBasis": "built_with_model_package", "pe:dataBuildId": "populace-us-2024-f0af251-703bd81a565c-20260620T201958Z", - "pe:emittedIn": "github-actions", + "pe:emittedIn": "local", "rdfs:comment": "Certification of build populace-us-2024-f0af251-703bd81a565c-20260620T201958Z for policyengine-us 1.729.0.", "trov:accessedArrangement": { "@id": "arrangement/1" diff --git a/src/policyengine/provenance/certification.py b/src/policyengine/provenance/certification.py index 07cc8fee..6ff573bf 100644 --- a/src/policyengine/provenance/certification.py +++ b/src/policyengine/provenance/certification.py @@ -66,6 +66,59 @@ BASIS_BUILT_WITH = "built_with_model_package" BASIS_PUBLISHER_CLAIM = "compatible_model_packages" POPULACE_US_SOURCE_COVERAGE_FILE = "us_source_coverage.json" +US_STATE_CODES = ( + "AL", + "AK", + "AZ", + "AR", + "CA", + "CO", + "CT", + "DE", + "DC", + "FL", + "GA", + "HI", + "ID", + "IL", + "IN", + "IA", + "KS", + "KY", + "LA", + "ME", + "MD", + "MA", + "MI", + "MN", + "MS", + "MO", + "MT", + "NE", + "NV", + "NH", + "NJ", + "NM", + "NY", + "NC", + "ND", + "OH", + "OK", + "OR", + "PA", + "RI", + "SC", + "SD", + "TN", + "TX", + "UT", + "VT", + "VA", + "WA", + "WV", + "WI", + "WY", +) class CertificationError(ValueError): @@ -332,6 +385,73 @@ def should_validate_vendored_artifacts( ) +def _state_code_from_path(path: str, artifact_prefix: str) -> Optional[str]: + prefix = artifact_prefix.rstrip("/") + "/" + if not path.startswith(prefix): + return None + remainder = path.removeprefix(prefix) + if "/" in remainder or not remainder.endswith(".h5"): + return None + state_code = remainder.removesuffix(".h5") + if state_code not in US_STATE_CODES: + return None + return state_code + + +def merge_us_state_release_manifest( + primary_manifest: DataReleaseManifest, + state_manifest: DataReleaseManifest, + *, + artifact_prefix: str = "states/", + path_template: str = "states/{state_code}.h5", +) -> DataReleaseManifest: + """Add legacy US state artifacts to a primary Populace release manifest.""" + + primary_payload = primary_manifest.model_dump(mode="json", exclude_none=True) + merged_artifacts = primary_payload.setdefault("artifacts", {}) + found_states: dict[str, str] = {} + prefix = artifact_prefix.rstrip("/") + "/" + + for name, artifact in state_manifest.artifacts.items(): + if not name.startswith(prefix) and not artifact.path.startswith(prefix): + continue + state_code = _state_code_from_path(artifact.path, artifact_prefix) + if state_code is None: + raise CertificationError( + "US state artifact path must match " + f"{artifact_prefix}.h5; got {artifact.path!r}." + ) + if not artifact.sha256: + raise CertificationError(f"US state artifact {name!r} has no sha256.") + if name in merged_artifacts: + raise CertificationError( + f"Regional artifact {name!r} conflicts with the primary manifest." + ) + merged_artifacts[name] = artifact.model_dump( + mode="json", + exclude_none=True, + ) + found_states[state_code] = name + + missing_states = sorted(set(US_STATE_CODES) - set(found_states)) + if missing_states: + raise CertificationError( + "Missing US state artifacts: " + ", ".join(missing_states) + ) + + metadata = primary_payload.setdefault("metadata", {}) + region_datasets = metadata.setdefault("region_datasets", {}) + default_dataset = primary_manifest.default_datasets["national"] + default_artifact = primary_manifest.artifacts[default_dataset] + region_datasets.setdefault( + "national", + {"path_template": default_artifact.path}, + ) + region_datasets["state"] = {"path_template": path_template} + + return DataReleaseManifest.model_validate(primary_payload) + + def build_country_manifest_payload( *, country: str, @@ -434,6 +554,8 @@ def build_bundle_data_release_payload( data_producer: str, manifest_uri: str, uri_parts: dict, + regional_manifest_uri: Optional[str] = None, + regional_uri_parts: Optional[dict] = None, ) -> dict: """Map the country certification payload into bundle ``data_releases``.""" payload = copy.deepcopy(country_payload) @@ -445,6 +567,12 @@ def build_bundle_data_release_payload( payload["build_id"] = build_id payload["source_manifest_uri"] = manifest_uri payload["release_manifest_uri"] = https_manifest_url(uri_parts) + if regional_manifest_uri is not None: + payload["regional_source_manifest_uri"] = regional_manifest_uri + if regional_uri_parts is not None: + payload["regional_release_manifest_uri"] = https_manifest_url( + regional_uri_parts + ) payload["default_dataset_uri"] = certified_artifact.get("uri") return payload @@ -461,6 +589,9 @@ def certify( model_version: str, token: Optional[str], check_artifacts: bool, + regional_manifest_uri: Optional[str] = None, + regional_artifact_prefix: str = "states/", + regional_path_template: str = "states/{state_code}.h5", ) -> tuple[dict, str, list[str]]: raise NotImplementedError @@ -477,6 +608,9 @@ def certify( model_version: str, token: Optional[str], check_artifacts: bool, + regional_manifest_uri: Optional[str] = None, + regional_artifact_prefix: str = "states/", + regional_path_template: str = "states/{state_code}.h5", ) -> tuple[dict, str, list[str]]: raise CertificationError( "Legacy data-producer certification updates are not implemented in " @@ -496,10 +630,29 @@ def certify( model_version: str, token: Optional[str], check_artifacts: bool, + regional_manifest_uri: Optional[str] = None, + regional_artifact_prefix: str = "states/", + regional_path_template: str = "states/{state_code}.h5", ) -> tuple[dict, str, list[str]]: manifest, manifest_sha256, uri_parts = fetch_release_manifest( manifest_uri, token=token ) + regional_uri_parts = None + if regional_manifest_uri is not None: + if country != "us": + raise CertificationError( + "Regional data release overlays are only supported for US " + "Populace certification." + ) + state_manifest, _, regional_uri_parts = fetch_release_manifest( + regional_manifest_uri, token=token + ) + manifest = merge_us_state_release_manifest( + manifest, + state_manifest, + artifact_prefix=regional_artifact_prefix, + path_template=regional_path_template, + ) compatibility_basis, warnings = validate_release_manifest( manifest, model_package, model_version ) @@ -546,6 +699,8 @@ def certify( data_producer=self.data_producer, manifest_uri=manifest_uri, uri_parts=uri_parts, + regional_manifest_uri=regional_manifest_uri, + regional_uri_parts=regional_uri_parts, ), manifest_sha256, warnings, @@ -584,6 +739,9 @@ def certify_data_release( token: Optional[str] = None, bundle_path: Optional[Path] = None, check_artifacts: bool = True, + regional_manifest_uri: Optional[str] = None, + regional_artifact_prefix: str = "states/", + regional_path_template: str = "states/{state_code}.h5", ) -> CertificationResult: if country not in COUNTRY_MODEL_PACKAGES: raise CertificationError(f"Unknown country {country!r}.") @@ -598,6 +756,9 @@ def certify_data_release( model_version=model_version, token=token, check_artifacts=check_artifacts, + regional_manifest_uri=regional_manifest_uri, + regional_artifact_prefix=regional_artifact_prefix, + regional_path_template=regional_path_template, ) bundle_path = bundle_path or default_bundle_source_path() bundle = json.loads(bundle_path.read_text()) diff --git a/tests/test_certify_data_release.py b/tests/test_certify_data_release.py index 5ca207f9..ae9838a2 100644 --- a/tests/test_certify_data_release.py +++ b/tests/test_certify_data_release.py @@ -6,9 +6,11 @@ import pytest from policyengine.provenance.certification import ( + US_STATE_CODES, CertificationError, build_country_manifest_payload, certify_data_release, + merge_us_state_release_manifest, parse_manifest_uri, validate_release_manifest, ) @@ -23,6 +25,11 @@ "hf://dataset/policyengine/populace-uk-private" f"@{UK_TAG}/releases/{UK_TAG}/release_manifest.json" ) +US_DATA_VERSION = "1.115.5" +US_DATA_MANIFEST_URI = ( + "hf://model/policyengine/policyengine-us-data" + f"@{US_DATA_VERSION}/releases/{US_DATA_VERSION}/release_manifest.json" +) def _release_manifest_payload() -> dict: @@ -96,6 +103,45 @@ def _release_manifest_payload() -> dict: } +def _populace_manifest_payload_without_regions() -> dict: + payload = _release_manifest_payload() + payload["metadata"] = {} + payload["artifacts"].pop("states/AK") + return payload + + +def _state_release_manifest_payload( + *, + missing: set[str] | None = None, + malformed: dict[str, str] | None = None, + without_hash: set[str] | None = None, +) -> dict: + missing = missing or set() + malformed = malformed or {} + without_hash = without_hash or set() + artifacts = {} + for index, state_code in enumerate(US_STATE_CODES): + if state_code in missing: + continue + path = malformed.get(state_code, f"states/{state_code}.h5") + artifact = { + "kind": "microdata", + "path": path, + "repo_id": "policyengine/policyengine-us-data", + "revision": US_DATA_VERSION, + "sha256": f"{index + 1:064x}", + "size_bytes": 1, + } + if state_code in without_hash: + artifact.pop("sha256") + artifacts[f"states/{state_code}"] = artifact + return { + "schema_version": 1, + "data_package": {"name": "policyengine-us-data", "version": US_DATA_VERSION}, + "artifacts": artifacts, + } + + def _manifest() -> DataReleaseManifest: return DataReleaseManifest.model_validate(_release_manifest_payload()) @@ -127,9 +173,13 @@ def _bundle_source_payload() -> dict: "packages": { "policyengine": {"name": "policyengine", "version": "9.9.9"}, "policyengine-uk": {"name": "policyengine-uk", "version": "2.0.0"}, + "policyengine-us": {"name": "policyengine-us", "version": "1.0.0"}, }, "extras": {}, - "countries": {"uk": {"model_package": "policyengine-uk"}}, + "countries": { + "uk": {"model_package": "policyengine-uk"}, + "us": {"model_package": "policyengine-us"}, + }, "data_releases": {}, } @@ -270,6 +320,69 @@ def test__given_build_provenance__then_certification_carries_it(self): ) +class TestMergeUSStateReleaseManifest: + def test__given_state_manifest__then_adds_state_region_artifacts(self): + primary = DataReleaseManifest.model_validate( + _populace_manifest_payload_without_regions() + ) + states = DataReleaseManifest.model_validate(_state_release_manifest_payload()) + + merged = merge_us_state_release_manifest(primary, states) + payload = build_country_manifest_payload( + country="us", + manifest=merged, + uri_parts=parse_manifest_uri(MANIFEST_URI), + policyengine_version="9.9.9", + model_package="policyengine-us", + model_version="1.723.0", + model_wheel={}, + ) + + assert payload["datasets"]["states/CA"] == { + "path": "states/CA.h5", + "revision": US_DATA_VERSION, + "sha256": f"{US_STATE_CODES.index('CA') + 1:064x}", + "repo_id": "policyengine/policyengine-us-data", + } + assert payload["region_datasets"] == { + "national": {"path_template": "populace_us_2024.h5"}, + "state": {"path_template": "states/{state_code}.h5"}, + } + + def test__given_missing_state_artifact__then_raises(self): + primary = DataReleaseManifest.model_validate( + _populace_manifest_payload_without_regions() + ) + states = DataReleaseManifest.model_validate( + _state_release_manifest_payload(missing={"CA"}) + ) + + with pytest.raises(CertificationError, match="Missing US state artifacts: CA"): + merge_us_state_release_manifest(primary, states) + + def test__given_malformed_state_path__then_raises(self): + primary = DataReleaseManifest.model_validate( + _populace_manifest_payload_without_regions() + ) + states = DataReleaseManifest.model_validate( + _state_release_manifest_payload(malformed={"CA": "state/CA.h5"}) + ) + + with pytest.raises(CertificationError, match="states/.h5"): + merge_us_state_release_manifest(primary, states) + + def test__given_state_without_hash__then_raises(self): + primary = DataReleaseManifest.model_validate( + _populace_manifest_payload_without_regions() + ) + states = DataReleaseManifest.model_validate( + _state_release_manifest_payload(without_hash={"CA"}) + ) + + with pytest.raises(CertificationError, match="states/CA"): + merge_us_state_release_manifest(primary, states) + + class TestCertifyDataRelease: def test__given_fetched_populace_manifest__then_updates_bundle_manifest( self, tmp_path @@ -327,6 +440,70 @@ def test__given_fetched_populace_manifest__then_updates_bundle_manifest( assert result.build_id == UK_TAG assert result.bundle_path == bundle_path + def test__given_us_regional_manifest__then_certifies_state_artifacts( + self, tmp_path + ): + bundle_path = tmp_path / "manifest.json" + bundle_path.write_text(json.dumps(_bundle_source_payload()) + "\n") + primary_response = MagicMock() + primary_response.status_code = 200 + primary_response.content = json.dumps( + _populace_manifest_payload_without_regions() + ).encode() + regional_response = MagicMock() + regional_response.status_code = 200 + regional_response.content = json.dumps( + _state_release_manifest_payload() + ).encode() + + with ( + patch( + "policyengine.provenance.certification.requests.get", + side_effect=[primary_response, regional_response], + ), + patch( + "policyengine.provenance.certification.head_artifact", + return_value=True, + ), + patch( + "policyengine.provenance.certification.head_release_file", + return_value=True, + ), + patch( + "policyengine.provenance.certification.head_artifact_reference", + return_value=True, + ), + patch( + "policyengine.provenance.certification.fetch_pypi_wheel_metadata", + return_value={"sha256": "d" * 64, "url": "https://example"}, + ), + patch( + "policyengine.provenance.certification.policyengine_version", + return_value="9.9.9", + ), + ): + result = certify_data_release( + country="us", + data_producer="populace", + manifest_uri=MANIFEST_URI, + regional_manifest_uri=US_DATA_MANIFEST_URI, + model_version="1.723.0", + bundle_path=bundle_path, + ) + + written = json.loads(bundle_path.read_text()) + release = written["data_releases"]["us"] + assert release["source_manifest_uri"] == MANIFEST_URI + assert release["regional_source_manifest_uri"] == US_DATA_MANIFEST_URI + assert release["region_datasets"]["state"] == { + "path_template": "states/{state_code}.h5" + } + assert release["datasets"]["states/CA"]["repo_id"] == ( + "policyengine/policyengine-us-data" + ) + assert release["datasets"]["states/CA"]["revision"] == US_DATA_VERSION + assert result.dataset_count == 4 + len(US_STATE_CODES) + def test__given_us_without_data_producer__then_legacy_update_is_explicitly_unsupported( self, tmp_path ): diff --git a/tests/test_us_regions.py b/tests/test_us_regions.py index a1250e93..d4022341 100644 --- a/tests/test_us_regions.py +++ b/tests/test_us_regions.py @@ -2,8 +2,10 @@ from policyengine.countries.us.data import DISTRICT_COUNTS, US_STATES from policyengine.countries.us.regions import ( + build_us_region_registry, us_region_registry, ) +from policyengine.provenance.manifest import CountryReleaseManifest class TestUSStates: @@ -93,7 +95,7 @@ def test__given_us_registry__then_country_id_is_us(self): def test__given_us_registry__then_has_national_region(self): """Given: US region registry When: Getting national region - Then: Returns US without a dedicated region dataset path + Then: Returns US with the certified national dataset path """ # When national = us_region_registry.get_national() @@ -103,7 +105,10 @@ def test__given_us_registry__then_has_national_region(self): assert national.code == "us" assert national.label == "United States" assert national.region_type == "national" - assert national.dataset_path is None + assert national.dataset_path == ( + "hf://policyengine/populace-us/populace_us_2024.h5" + "@populace-us-2024-f0af251-703bd81a565c-20260620T201958Z" + ) def test__given_us_registry__then_has_51_states(self): """Given: US region registry @@ -129,7 +134,9 @@ def test__given_california_region__then_has_correct_format(self): assert ca.label == "California" assert ca.region_type == "state" assert ca.parent_code == "us" - assert ca.dataset_path is None + assert ca.dataset_path == ( + "hf://policyengine/policyengine-us-data/states/CA.h5@1.115.5" + ) assert ca.state_code == "CA" assert ca.state_name == "California" assert not ca.requires_filter @@ -225,16 +232,74 @@ def test__given_california__then_children_include_districts_and_places( assert len(district_children) == DISTRICT_COUNTS["CA"] assert len(place_children) >= 10 # CA has many large cities - def test__given_us_registry__then_dataset_regions_is_empty(self): + def test__given_us_registry__then_dataset_regions_are_national_and_states(self): """Given: US region registry When: Getting regions with datasets - Then: Current Populace-certified bundle has no subnational datasets + Then: Current certified bundle has national and state datasets """ # When dataset_regions = us_region_registry.get_dataset_regions() # Then - assert len(dataset_regions) == 0 + assert len(dataset_regions) == 52 + assert {region.region_type for region in dataset_regions} == { + "national", + "state", + } + + def test__given_certified_state_template__then_states_have_dataset_paths( + self, monkeypatch + ): + """Given: US bundle manifest with a certified state template + When: Building the region registry + Then: State regions resolve to pinned state dataset artifacts + """ + manifest = CountryReleaseManifest.model_validate( + { + "country_id": "us", + "policyengine_version": "9.9.9", + "model_package": { + "name": "policyengine-us", + "version": "1.723.0", + }, + "data_package": { + "name": "populace-data", + "version": "0.1.0", + "repo_id": "policyengine/populace-us", + "repo_type": "dataset", + "release_manifest_revision": "populace-us-tag", + }, + "default_dataset": "populace_us_2024", + "datasets": { + "populace_us_2024": { + "path": "populace_us_2024.h5", + "repo_id": "policyengine/populace-us", + "revision": "populace-us-tag", + }, + "states/CA": { + "path": "states/CA.h5", + "repo_id": "policyengine/policyengine-us-data", + "revision": "1.115.5", + }, + }, + "region_datasets": { + "national": {"path_template": "populace_us_2024.h5"}, + "state": {"path_template": "states/{state_code}.h5"}, + }, + } + ) + monkeypatch.setattr( + "policyengine.provenance.manifest.get_release_manifest", + lambda country_id: manifest, + ) + + registry = build_us_region_registry() + ca = registry.get("state/ca") + + assert ca is not None + assert ca.dataset_path == ( + "hf://policyengine/policyengine-us-data/states/CA.h5@1.115.5" + ) def test__given_us_registry__then_filter_regions_are_all_places(self): """Given: US region registry From 7cde69be0d645a717ad996dd7dffeb3373ae1b93 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 25 Jun 2026 19:52:38 +0200 Subject: [PATCH 2/3] Document US state data certification overlay --- docs/bundles.md | 27 ++++++++++++++-- docs/engineering/skills/data-certification.md | 31 ++++++++++++++++++- docs/release-bundles.md | 22 ++++++++++++- 3 files changed, 76 insertions(+), 4 deletions(-) diff --git a/docs/bundles.md b/docs/bundles.md index 4afe5455..0d1c0e45 100644 --- a/docs/bundles.md +++ b/docs/bundles.md @@ -28,12 +28,18 @@ uvx --from policyengine policyengine bundle install When run from `uvx` or `pipx`, the installer creates or reuses `./.venv`. Inside an existing virtualenv or conda environment, it installs into that active environment. The installer then installs the -exact bundled package scaffold with pip, downloads certified US and UK datasets -into `./data`, moves replaced dataset files into +exact bundled package scaffold with pip, downloads certified default US and UK +datasets into `./data`, moves replaced dataset files into `./data/.policyengine-bundle-backups//`, and writes a `./data/.policyengine-bundle-receipt.json` receipt that records the target Python. +The bundle manifest can certify additional regional datasets, such as US state +datasets. Those artifacts are part of the citable bundle manifest, but +`policyengine bundle install` does not eagerly download every regional file. +Runtime callers should use the manifest's regional dataset URI when a regional +simulation needs one. + Country-specific and package-only installs are supported: ```bash @@ -82,6 +88,23 @@ python scripts/bundle.py certify-data \ --manifest-uri hf://dataset/policyengine/populace-uk-private@/releases//release_manifest.json ``` +For US Populace releases, include the inherited state datasets from +`policyengine-us-data`: + +```bash +python scripts/bundle.py certify-data \ + --country us \ + --data-producer populace \ + --manifest-uri hf://dataset/policyengine/populace-us@/releases//release_manifest.json \ + --regional-manifest-uri hf://model/policyengine/policyengine-us-data@/releases//release_manifest.json \ + --model-version +``` + +The regional manifest must include all 51 `states/{STATE}.h5` artifacts with +their original repo, revision, and sha256 pins. The resulting bundle manifest +certifies Populace as the US national default dataset and +`policyengine-us-data` as the state dataset source. + Use `python scripts/bundle.py generate` to regenerate derived bundle metadata, and `python scripts/bundle.py generate --include-tros` when TRACE TRO sidecars should also be regenerated. Private data releases require `HUGGING_FACE_TOKEN` diff --git a/docs/engineering/skills/data-certification.md b/docs/engineering/skills/data-certification.md index ff522638..b70f55c0 100644 --- a/docs/engineering/skills/data-certification.md +++ b/docs/engineering/skills/data-certification.md @@ -26,6 +26,34 @@ python scripts/bundle.py certify-data --country uk --data-producer populace \ --manifest-uri "hf://dataset/policyengine/populace-uk-private@/releases//release_manifest.json" ``` +For US Populace certification, include the inherited state datasets from the +certified `policyengine-us-data` release manifest: + +```bash +python scripts/bundle.py certify-data --country us --data-producer populace \ + --manifest-uri "hf://dataset/policyengine/populace-us@/releases//release_manifest.json" \ + --regional-manifest-uri "hf://model/policyengine/policyengine-us-data@/releases//release_manifest.json" \ + --model-version "" +``` + +The regional manifest is required for US while the stack still serves +state-level datasets from `policyengine-us-data`. It must contain all 51 +`states/{STATE}.h5` artifacts, including DC, and each state artifact must carry +its original `repo_id`, `revision`, and `sha256`. Certification preserves those +per-artifact pins in `data_releases.us.datasets` and writes: + +```json +"region_datasets": { + "national": {"path_template": "populace_us_2024.h5"}, + "state": {"path_template": "states/{state_code}.h5"} +} +``` + +Do not move or rewrite state artifacts into the Populace repo. The certified +bundle is intentionally hybrid: Populace owns the national default dataset, and +`policyengine-us-data` owns the inherited state datasets until that path is +migrated. + The script fetches and validates the manifest (every artifact must carry a revision pin; the certified dataset must be reachable), writes the canonical bundle manifest, exact-pins the country model package in that same manifest, @@ -53,7 +81,8 @@ A certification PR should normally change only: Hard failures (certification refuses): missing national default dataset, default dataset absent from artifacts, any artifact without a revision pin, unreachable certified dataset, missing required supplemental release files -(for example Populace-US `us_source_coverage.json`), unknown country. +(for example Populace-US `us_source_coverage.json`), missing or malformed US +state overlay artifacts when `--regional-manifest-uri` is used, unknown country. Certification gate: the model version must either exactly match the build-time model (`compatibility_basis: built_with_model_package`) or be diff --git a/docs/release-bundles.md b/docs/release-bundles.md index 20ada90f..7c71697a 100644 --- a/docs/release-bundles.md +++ b/docs/release-bundles.md @@ -38,12 +38,17 @@ uvx --from policyengine==4.19.1 policyengine bundle install 4.19.1 When run from `uvx` or `pipx`, the command creates or reuses `.venv`. Inside an existing virtualenv or conda environment, it installs into that active environment. It installs the bundled Python packages with pip, downloads the -certified US and UK datasets into `./data`, and writes a +certified default US and UK datasets into `./data`, and writes a `./data/.policyengine-bundle-receipt.json` receipt that records the target Python. Existing dataset files with the same filename are moved to `./data/.policyengine-bundle-backups//`. +Regional datasets may also be certified in the bundle manifest. They are not +eagerly downloaded by `policyengine bundle install`; callers should materialize +the certified regional URI from the manifest when they run a regional +simulation. + Useful variants: ```bash @@ -123,6 +128,21 @@ python scripts/bundle.py certify-data --country us \ --manifest-uri "hf://dataset/policyengine/populace-us@/releases//release_manifest.json" ``` +US Populace certification currently also needs the inherited state-level +datasets from the certified `policyengine-us-data` release manifest: + +```bash +python scripts/bundle.py certify-data --country us --data-producer populace \ + --manifest-uri "hf://dataset/policyengine/populace-us@/releases//release_manifest.json" \ + --regional-manifest-uri "hf://model/policyengine/policyengine-us-data@/releases//release_manifest.json" \ + --model-version "" +``` + +That produces one US bundle manifest entry containing the Populace national +default dataset plus all 51 `states/{STATE}.h5` artifacts pinned to +`policyengine-us-data`. The resulting `region_datasets.state` template lets +runtime code resolve a state region to the exact certified state artifact. + Earlier releases (policyengine 4.15.x–4.16.x) were certified through the `PolicyEngine/policyengine-bundles` archive flow; those bundles remain the historical record of their certifications. From 3617911e27a1b5c8a11b38fbb08be987604edb31 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Thu, 25 Jun 2026 20:13:02 +0200 Subject: [PATCH 3/3] Harden US regional data certification --- docs/bundles.md | 4 + docs/engineering/skills/data-certification.md | 4 + docs/release-bundles.md | 4 + scripts/bundle.py | 5 +- scripts/certify_data_release.py | 5 +- src/policyengine/provenance/certification.py | 10 ++ tests/test_certify_data_release.py | 93 +++++++++++++++++++ 7 files changed, 123 insertions(+), 2 deletions(-) diff --git a/docs/bundles.md b/docs/bundles.md index 0d1c0e45..df28a2a0 100644 --- a/docs/bundles.md +++ b/docs/bundles.md @@ -104,6 +104,10 @@ The regional manifest must include all 51 `states/{STATE}.h5` artifacts with their original repo, revision, and sha256 pins. The resulting bundle manifest certifies Populace as the US national default dataset and `policyengine-us-data` as the state dataset source. +The regional manifest URI is recorded for traceability; the bundle does not +currently record the regional manifest's own sha256. The citable pins are the +artifact-level repo, revision, and sha256 values copied into +`data_releases.us.datasets`. Use `python scripts/bundle.py generate` to regenerate derived bundle metadata, and `python scripts/bundle.py generate --include-tros` when TRACE TRO sidecars diff --git a/docs/engineering/skills/data-certification.md b/docs/engineering/skills/data-certification.md index b70f55c0..890a0fd4 100644 --- a/docs/engineering/skills/data-certification.md +++ b/docs/engineering/skills/data-certification.md @@ -53,6 +53,10 @@ Do not move or rewrite state artifacts into the Populace repo. The certified bundle is intentionally hybrid: Populace owns the national default dataset, and `policyengine-us-data` owns the inherited state datasets until that path is migrated. +The regional manifest URI is recorded for traceability, but the bundle does not +currently record the regional manifest's own sha256. Treat the copied +artifact-level repo, revision, and sha256 pins in `data_releases.us.datasets` +as the citable state dataset certification. The script fetches and validates the manifest (every artifact must carry a revision pin; the certified dataset must be reachable), writes the canonical diff --git a/docs/release-bundles.md b/docs/release-bundles.md index 7c71697a..c16a938f 100644 --- a/docs/release-bundles.md +++ b/docs/release-bundles.md @@ -142,6 +142,10 @@ That produces one US bundle manifest entry containing the Populace national default dataset plus all 51 `states/{STATE}.h5` artifacts pinned to `policyengine-us-data`. The resulting `region_datasets.state` template lets runtime code resolve a state region to the exact certified state artifact. +The regional manifest URI is retained for traceability, but the bundle does not +currently store the regional manifest's own sha256. For inherited state data, +the citable pins are the copied artifact-level repo, revision, and sha256 +values in `data_releases.us.datasets`. Earlier releases (policyengine 4.15.x–4.16.x) were certified through the `PolicyEngine/policyengine-bundles` archive flow; those bundles remain the diff --git a/scripts/bundle.py b/scripts/bundle.py index 6d402cd4..8c2aa2ba 100644 --- a/scripts/bundle.py +++ b/scripts/bundle.py @@ -175,7 +175,10 @@ def _parser() -> argparse.ArgumentParser: certify.add_argument( "--skip-artifact-check", action="store_true", - help="Skip the certified dataset reachability check.", + help=( + "Skip reachability checks for the certified dataset and any " + "vendored/regional artifacts." + ), ) certify.set_defaults(func=_certify_data) diff --git a/scripts/certify_data_release.py b/scripts/certify_data_release.py index 8da7771a..c4c1ec54 100644 --- a/scripts/certify_data_release.py +++ b/scripts/certify_data_release.py @@ -77,7 +77,10 @@ def main(argv=None) -> int: parser.add_argument( "--skip-artifact-check", action="store_true", - help="Skip the reachability HEAD on the certified dataset.", + help=( + "Skip reachability HEAD checks for the certified dataset and any " + "vendored/regional artifacts." + ), ) args = parser.parse_args(argv) diff --git a/src/policyengine/provenance/certification.py b/src/policyengine/provenance/certification.py index 6ff573bf..a0ade420 100644 --- a/src/policyengine/provenance/certification.py +++ b/src/policyengine/provenance/certification.py @@ -39,6 +39,7 @@ import hashlib import json import re +import warnings from dataclasses import dataclass, field from pathlib import Path from typing import Optional @@ -427,6 +428,15 @@ def merge_us_state_release_manifest( raise CertificationError( f"Regional artifact {name!r} conflicts with the primary manifest." ) + if state_code in found_states: + warnings.warn( + "Duplicate US state artifact for " + f"{state_code}: ignoring {name!r}; already using " + f"{found_states[state_code]!r}.", + RuntimeWarning, + stacklevel=2, + ) + continue merged_artifacts[name] = artifact.model_dump( mode="json", exclude_none=True, diff --git a/tests/test_certify_data_release.py b/tests/test_certify_data_release.py index ae9838a2..c8d9119f 100644 --- a/tests/test_certify_data_release.py +++ b/tests/test_certify_data_release.py @@ -1,6 +1,10 @@ """Tests for direct data-release certification.""" +import importlib.util import json +import sys +from pathlib import Path +from types import ModuleType from unittest.mock import MagicMock, patch import pytest @@ -146,6 +150,19 @@ def _manifest() -> DataReleaseManifest: return DataReleaseManifest.model_validate(_release_manifest_payload()) +def _load_bundle_script(monkeypatch): + scripts_dir = Path(__file__).resolve().parents[1] / "scripts" + monkeypatch.syspath_prepend(str(scripts_dir)) + spec = importlib.util.spec_from_file_location( + "bundle_script", + scripts_dir / "bundle.py", + ) + module = importlib.util.module_from_spec(spec) + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + def _uk_release_manifest_payload() -> dict: payload = _release_manifest_payload() payload["compatible_model_packages"] = [ @@ -382,6 +399,26 @@ def test__given_state_without_hash__then_raises(self): with pytest.raises(CertificationError, match="states/CA"): merge_us_state_release_manifest(primary, states) + def test__given_duplicate_state_artifact__then_warns_and_ignores_duplicate(self): + primary = DataReleaseManifest.model_validate( + _populace_manifest_payload_without_regions() + ) + state_payload = _state_release_manifest_payload() + state_payload["artifacts"]["states/CA_duplicate"] = { + **state_payload["artifacts"]["states/CA"], + "sha256": "f" * 64, + } + states = DataReleaseManifest.model_validate(state_payload) + + with pytest.warns(RuntimeWarning, match="Duplicate US state artifact for CA"): + merged = merge_us_state_release_manifest(primary, states) + + assert "states/CA_duplicate" not in merged.artifacts + assert ( + merged.artifacts["states/CA"].sha256 + == (state_payload["artifacts"]["states/CA"]["sha256"]) + ) + class TestCertifyDataRelease: def test__given_fetched_populace_manifest__then_updates_bundle_manifest( @@ -518,6 +555,62 @@ def test__given_us_without_data_producer__then_legacy_update_is_explicitly_unsup bundle_path=bundle_path, ) + def test__given_bundle_wrapper_regional_args__then_forwards_to_certifier( + self, monkeypatch + ): + bundle_script = _load_bundle_script(monkeypatch) + captured: dict[str, list[str]] = {} + fake_certifier = ModuleType("certify_data_release") + + def fake_main(argv: list[str]) -> int: + captured["argv"] = argv + return 0 + + fake_certifier.main = fake_main + monkeypatch.setitem(sys.modules, "certify_data_release", fake_certifier) + + result = bundle_script.main( + [ + "certify-data", + "--country", + "us", + "--data-producer", + "populace", + "--manifest-uri", + MANIFEST_URI, + "--regional-manifest-uri", + US_DATA_MANIFEST_URI, + "--regional-artifact-prefix", + "states/", + "--regional-path-template", + "states/{state_code}.h5", + "--model-version", + "1.723.0", + "--no-generate", + "--skip-artifact-check", + ] + ) + + assert result == 0 + assert captured["argv"] == [ + "--country", + "us", + "--manifest-uri", + MANIFEST_URI, + "--data-producer", + "populace", + "--model-version", + "1.723.0", + "--regional-manifest-uri", + US_DATA_MANIFEST_URI, + "--regional-artifact-prefix", + "states/", + "--regional-path-template", + "states/{state_code}.h5", + "--no-generate", + "--skip-artifact-check", + ] + def test__given_missing_populace_us_source_coverage__then_raises(self, tmp_path): response = MagicMock() response.status_code = 200