Merge pull request #424 from PolicyEngine/release-bundle-single-country-fallback

Support compatible UK model specifiers

Author: anth-volk

changelog.d/423.added.md

@@ -0,0 +1 @@
+Add release-manifest support for declaring additional compatible UK model versions.

policyengine_uk_data/tests/test_release_manifest.py

@@ -186,6 +186,37 @@ def test_build_release_manifest_tracks_uk_release_artifacts(tmp_path):
     }
 
 
+def test_build_release_manifest_adds_additional_compatible_specifiers(tmp_path):
+    dataset_path = _write_file(
+        tmp_path / "enhanced_frs_2023_24.h5",
+        b"enhanced-frs",
+    )
+
+    manifest = build_release_manifest(
+        files_with_repo_paths=[(dataset_path, "enhanced_frs_2023_24.h5")],
+        version="1.40.4",
+        repo_id=PRIVATE_REPO,
+        model_package_version="2.74.0",
+        model_package_git_sha="deadbeef",
+        model_package_data_build_fingerprint="sha256:fingerprint",
+        core_package_metadata=EXPECTED_CORE_PACKAGE,
+        data_package_git_sha="cafebabe",
+        additional_compatible_specifiers=("==2.89.0", ">=2.90.0,<3"),
+        created_at="2026-04-10T12:00:00Z",
+    )
+
+    assert manifest["compatible_model_packages"] == [
+        {"name": "policyengine-uk", "specifier": "==2.74.0"},
+        {"name": "policyengine-uk", "specifier": "==2.89.0"},
+        {"name": "policyengine-uk", "specifier": ">=2.90.0,<3"},
+    ]
+    validate_release_manifest(
+        manifest,
+        version="1.40.4",
+        repo_id=PRIVATE_REPO,
+    )
+
+
 def test_build_release_manifest_defaults_to_current_frs_release(tmp_path):
     enhanced_path = _write_file(
         tmp_path / CURRENT_FRS_RELEASE.enhanced_dataset_file,
@@ -590,6 +621,7 @@ def test_upload_files_to_hf_adds_uk_release_manifest_operations(tmp_path):
         upload_files_to_hf(
             files=[dataset_path],
             version="1.40.4",
+            additional_compatible_specifiers=("==2.89.0",),
         )
 
     operations = mock_api.create_commit.call_args.kwargs["operations"]
@@ -612,6 +644,10 @@ def test_upload_files_to_hf_adds_uk_release_manifest_operations(tmp_path):
     payload = release_ops[0].path_or_fileobj.getvalue()
     manifest = json.loads(payload.decode("utf-8"))
     _assert_single_uk_data_release_version(manifest)
+    assert manifest["compatible_model_packages"] == [
+        {"name": "policyengine-uk", "specifier": "==2.74.0"},
+        {"name": "policyengine-uk", "specifier": "==2.89.0"},
+    ]
     assert manifest["compatible_core_packages"] == EXPECTED_COMPATIBLE_CORE_PACKAGES
     assert manifest["build"]["built_with_core_package"] == EXPECTED_CORE_PACKAGE
     assert manifest["build"]["metadata"] == {

policyengine_uk_data/utils/data_upload.py

@@ -1,5 +1,5 @@
 from io import BytesIO
-from typing import Any, Dict, List, Mapping, Optional, Tuple
+from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple
 from huggingface_hub import HfApi, CommitOperationAdd, hf_hub_download
 from huggingface_hub.errors import EntryNotFoundError, RevisionNotFoundError
 from google.cloud import storage
@@ -296,6 +296,7 @@ def create_release_manifest_commit_operations(
     core_package_metadata: Optional[Mapping[str, Any]] = None,
     data_package_git_sha: Optional[str] = None,
     existing_manifest: Optional[Dict] = None,
+    additional_compatible_specifiers: Optional[Sequence[str]] = None,
 ) -> Tuple[Dict, List[CommitOperationAdd]]:
     manifest = build_release_manifest(
         files_with_repo_paths=files_with_repo_paths,
@@ -309,6 +310,7 @@ def create_release_manifest_commit_operations(
         core_package_metadata=core_package_metadata,
         data_package_git_sha=data_package_git_sha,
         existing_manifest=existing_manifest,
+        additional_compatible_specifiers=additional_compatible_specifiers,
     )
     validate_release_manifest(
         manifest,
@@ -337,6 +339,7 @@ def upload_data_files(
     hf_repo_name: str = PUBLIC_REPO,
     hf_repo_type: str = "model",
     version: str = None,
+    additional_compatible_specifiers: Optional[Sequence[str]] = None,
 ):
     if version is None:
         version = metadata.version("policyengine-uk-data")
@@ -346,6 +349,7 @@ def upload_data_files(
         version=version,
         hf_repo_name=hf_repo_name,
         hf_repo_type=hf_repo_type,
+        additional_compatible_specifiers=additional_compatible_specifiers,
     )
 
     upload_files_to_gcs(
@@ -360,6 +364,7 @@ def upload_files_to_hf(
     version: str,
     hf_repo_name: str = PRIVATE_REPO,
     hf_repo_type: str = "model",
+    additional_compatible_specifiers: Optional[Sequence[str]] = None,
 ):
     """
     Upload files to Hugging Face repository and tag the commit with the version.
@@ -414,6 +419,7 @@ def upload_files_to_hf(
         core_package_metadata=core_package_metadata,
         data_package_git_sha=_get_data_package_git_sha(),
         existing_manifest=existing_manifest,
+        additional_compatible_specifiers=additional_compatible_specifiers,
     )
     if finalized_manifest is not None:
         if candidate_manifest == finalized_manifest:

policyengine_uk_data/utils/release_manifest.py

@@ -129,15 +129,19 @@ def _model_package_compatibility(
     *,
     model_package_name: str,
     model_package_version: str | None,
+    additional_compatible_specifiers: Sequence[str] | None = None,
 ) -> list[Dict[str, str]]:
-    if not model_package_version:
-        return []
-    return [
-        {
-            "name": model_package_name,
-            "specifier": f"=={model_package_version}",
-        }
-    ]
+    compatible_packages = []
+    if model_package_version:
+        compatible_packages.append(
+            {
+                "name": model_package_name,
+                "specifier": f"=={model_package_version}",
+            }
+        )
+    for specifier in additional_compatible_specifiers or ():
+        compatible_packages.append({"name": model_package_name, "specifier": specifier})
+    return compatible_packages
 
 
 def _core_package_compatibility(
@@ -506,11 +510,13 @@ def _update_compatibility(
     model_package_name: str,
     model_package_version: str | None,
     core_package_metadata: Mapping[str, Any] | None,
+    additional_compatible_specifiers: Sequence[str] | None = None,
 ) -> None:
     manifest.setdefault("compatible_model_packages", [])
     model_package_compatibility = _model_package_compatibility(
         model_package_name=model_package_name,
         model_package_version=model_package_version,
+        additional_compatible_specifiers=additional_compatible_specifiers,
     )
     if model_package_compatibility:
         manifest["compatible_model_packages"] = model_package_compatibility
@@ -607,6 +613,7 @@ def build_release_manifest(
     existing_manifest: Mapping | None = None,
     default_datasets: Optional[Mapping[str, str]] = None,
     created_at: str | None = None,
+    additional_compatible_specifiers: Sequence[str] | None = None,
 ) -> Dict:
     manifest = _normalize_existing_manifest(
         existing_manifest,
@@ -644,6 +651,7 @@ def build_release_manifest(
         model_package_name=model_package_name,
         model_package_version=model_package_version,
         core_package_metadata=core_package_metadata,
+        additional_compatible_specifiers=additional_compatible_specifiers,
     )
     _update_artifacts(
         manifest,