Expose compatible model versions in release manifest backfill

Author: Max Ghenis

.github/workflows/backfill_release_manifest.yaml

@@ -16,6 +16,11 @@ on:
         description: "Exact policyengine-us version used with this data release"
         required: true
         type: string
+      compatible_model_package_versions:
+        description: "Optional newline-separated additional exact policyengine-us runtime versions to certify"
+        required: false
+        default: ""
+        type: string
       core_package_version:
         description: "Exact policyengine-core version used with this data release"
         required: true
@@ -107,6 +112,7 @@ jobs:
           VERSION: ${{ inputs.version }}
           ARTIFACTS: ${{ inputs.artifacts }}
           MODEL_PACKAGE_VERSION: ${{ inputs.model_package_version }}
+          COMPATIBLE_MODEL_PACKAGE_VERSIONS: ${{ inputs.compatible_model_package_versions }}
           CORE_PACKAGE_VERSION: ${{ inputs.core_package_version }}
           COMPATIBLE_CORE_PACKAGE_VERSIONS: ${{ inputs.compatible_core_package_versions }}
           MODEL_PACKAGE_GIT_SHA: ${{ inputs.model_package_git_sha }}
@@ -154,6 +160,17 @@ jobs:
             "${artifact_args[@]}"
           )
 
+          while IFS= read -r model_version; do
+            if [ -z "${model_version}" ]; then
+              continue
+            fi
+            if [[ ! "${model_version}" =~ ^[0-9]+[.][0-9]+[.][0-9]+$ ]]; then
+              echo "::error::Compatible model version must use x.y.z form, got '${model_version}'."
+              exit 1
+            fi
+            cmd+=(--compatible-model-package-version "${model_version}")
+          done <<< "${COMPATIBLE_MODEL_PACKAGE_VERSIONS}"
+
           while IFS= read -r core_version; do
             if [ -z "${core_version}" ]; then
               continue

policyengine_us_data/storage/backfill_release_manifest.py

@@ -111,6 +111,7 @@ def build_backfilled_release_manifest(
     artifacts: Sequence[HfArtifactMetadata],
     model_package_version: str,
     core_package_version: str,
+    compatible_model_package_versions: Sequence[str] | None = None,
     compatible_core_package_versions: Sequence[str] | None = None,
     hf_repo_name: str = DEFAULT_HF_REPO_NAME,
     model_package_name: str = DEFAULT_MODEL_PACKAGE_NAME,
@@ -132,6 +133,9 @@ def build_backfilled_release_manifest(
             "name": DEFAULT_CORE_PACKAGE_NAME,
             "version": core_package_version,
         },
+        additional_compatible_specifiers=[
+            f"=={version}" for version in compatible_model_package_versions or ()
+        ],
         additional_core_compatible_specifiers=[
             f"=={version}" for version in compatible_core_package_versions or ()
         ],
@@ -220,6 +224,17 @@ def _parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--model-package-version", required=True)
     parser.add_argument("--core-package-version", required=True)
+    parser.add_argument(
+        "--compatible-model-package-version",
+        action="append",
+        dest="compatible_model_package_versions",
+        default=[],
+        help=(
+            "Additional exact policyengine-us runtime version to certify as "
+            "compatible. Repeat for multiple versions. The build model remains "
+            "--model-package-version."
+        ),
+    )
     parser.add_argument(
         "--compatible-core-package-version",
         action="append",
@@ -278,6 +293,7 @@ def main() -> int:
         artifacts=artifacts,
         model_package_version=args.model_package_version,
         core_package_version=args.core_package_version,
+        compatible_model_package_versions=args.compatible_model_package_versions,
         compatible_core_package_versions=args.compatible_core_package_versions,
         hf_repo_name=args.hf_repo_name,
         model_package_git_sha=args.model_package_git_sha,

tests/unit/test_backfill_release_manifest.py

@@ -146,6 +146,28 @@ def test_build_backfilled_release_manifest_records_additional_core_compatibility
     ]
 
 
+def test_build_backfilled_release_manifest_records_additional_model_compatibility():
+    manifest = build_backfilled_release_manifest(
+        version="1.73.0",
+        artifacts=[
+            HfArtifactMetadata(
+                path="enhanced_cps_2024.h5",
+                sha256="a" * 64,
+                size_bytes=1,
+            )
+        ],
+        model_package_version="1.653.3",
+        compatible_model_package_versions=["1.722.4"],
+        core_package_version="3.26.0",
+    )
+
+    assert manifest["build"]["built_with_model_package"]["version"] == "1.653.3"
+    assert manifest["compatible_model_packages"] == [
+        {"name": "policyengine-us", "specifier": "==1.653.3"},
+        {"name": "policyengine-us", "specifier": "==1.722.4"},
+    ]
+
+
 def test_upload_backfilled_release_manifest_uploads_manifest_and_trace(monkeypatch):
     api = MagicMock()
     api.create_commit.return_value = SimpleNamespace(oid="commit-sha")