From 014cb0e1fce4cbb1f4ce6f9f3ccec21b7c674b34 Mon Sep 17 00:00:00 2001 From: Anthony Volk Date: Sat, 30 May 2026 01:05:37 +0200 Subject: [PATCH] Support compatible UK model specifiers --- changelog.d/423.added.md | 1 + .../tests/test_release_manifest.py | 36 +++++++++++++++++++ policyengine_uk_data/utils/data_upload.py | 8 ++++- .../utils/release_manifest.py | 24 ++++++++----- 4 files changed, 60 insertions(+), 9 deletions(-) create mode 100644 changelog.d/423.added.md diff --git a/changelog.d/423.added.md b/changelog.d/423.added.md new file mode 100644 index 000000000..619ff4e99 --- /dev/null +++ b/changelog.d/423.added.md @@ -0,0 +1 @@ +Add release-manifest support for declaring additional compatible UK model versions. diff --git a/policyengine_uk_data/tests/test_release_manifest.py b/policyengine_uk_data/tests/test_release_manifest.py index 3b4fc0d15..e1858bfda 100644 --- a/policyengine_uk_data/tests/test_release_manifest.py +++ b/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"] == { diff --git a/policyengine_uk_data/utils/data_upload.py b/policyengine_uk_data/utils/data_upload.py index e152598e8..7f96defc2 100644 --- a/policyengine_uk_data/utils/data_upload.py +++ b/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: diff --git a/policyengine_uk_data/utils/release_manifest.py b/policyengine_uk_data/utils/release_manifest.py index 08288d106..442d05a70 100644 --- a/policyengine_uk_data/utils/release_manifest.py +++ b/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,