Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions changelog.d/423.added.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Add release-manifest support for declaring additional compatible UK model versions.
36 changes: 36 additions & 0 deletions policyengine_uk_data/tests/test_release_manifest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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"]
Expand All @@ -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"] == {
Expand Down
8 changes: 7 additions & 1 deletion policyengine_uk_data/utils/data_upload.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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")
Expand All @@ -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(
Expand All @@ -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.
Expand Down Expand Up @@ -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:
Expand Down
24 changes: 16 additions & 8 deletions policyengine_uk_data/utils/release_manifest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down