Self-heal publication version sync

.github/scripts/check_data_release_version.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import argparse
+from dataclasses import dataclass
 import json
 import os
 from pathlib import Path
@@ -19,6 +20,17 @@
 )
 VERSION_RE = re.compile(r'^version\s*=\s*"([^"]+)"', re.MULTILINE)
 SEMVER_RE = re.compile(r"^(\d+)\.(\d+)\.(\d+)(?:rc\d+)?$")
+CURRENT = "current"
+BEHIND = "behind"
+AHEAD = "ahead"
+UNKNOWN = "unknown"
+
+
+@dataclass(frozen=True)
+class ReleaseVersionState:
+    package_version: str
+    finalized_release_version: str
+    release_version_relation: str
 
 
 def stable_version_tuple(version: str) -> tuple[int, int, int]:
@@ -28,7 +40,22 @@ def stable_version_tuple(version: str) -> tuple[int, int, int]:
     return tuple(int(part) for part in match.groups())
 
 
-def pyproject_version(root: Path = REPO_ROOT) -> str:
+def release_version_relation(
+    *,
+    package_version: str,
+    finalized_release_version: str,
+) -> str:
+    package_tuple = stable_version_tuple(package_version)
+    finalized_tuple = stable_version_tuple(finalized_release_version)
+    if package_tuple < finalized_tuple:
+        return BEHIND
+    if package_tuple > finalized_tuple:
+        return AHEAD
+    return CURRENT
+
+
+def pyproject_version(root: Path | None = None) -> str:
+    root = root or REPO_ROOT
     text = (root / "pyproject.toml").read_text()
     match = VERSION_RE.search(text)
     if not match:
@@ -58,8 +85,12 @@ def version_violations(
     package_version: str,
     finalized_release_version: str,
 ) -> list[str]:
-    if stable_version_tuple(package_version) >= stable_version_tuple(
-        finalized_release_version
+    if (
+        release_version_relation(
+            package_version=package_version,
+            finalized_release_version=finalized_release_version,
+        )
+        != BEHIND
     ):
         return []
     return [
@@ -70,22 +101,55 @@ def version_violations(
     ]
 
 
-def check_repository(
-    root: Path = REPO_ROOT,
+def check_repository_state(
+    root: Path | None = None,
     *,
     finalized_release_version: str | None = None,
     version_manifest_url: str = DEFAULT_VERSION_MANIFEST_URL,
-) -> list[str]:
+) -> ReleaseVersionState:
+    root = root or REPO_ROOT
     package_version = pyproject_version(root)
     finalized_release_version = finalized_release_version or latest_hf_release_version(
         version_manifest_url
     )
-    return version_violations(
+    relation = release_version_relation(
+        package_version=package_version,
+        finalized_release_version=finalized_release_version,
+    )
+    return ReleaseVersionState(
         package_version=package_version,
         finalized_release_version=finalized_release_version,
+        release_version_relation=relation,
     )
 
 
+def check_repository(
+    root: Path | None = None,
+    *,
+    finalized_release_version: str | None = None,
+    version_manifest_url: str = DEFAULT_VERSION_MANIFEST_URL,
+) -> list[str]:
+    state = check_repository_state(
+        root,
+        finalized_release_version=finalized_release_version,
+        version_manifest_url=version_manifest_url,
+    )
+    return version_violations(
+        package_version=state.package_version,
+        finalized_release_version=state.finalized_release_version,
+    )
+
+
+def write_github_outputs(state: ReleaseVersionState) -> None:
+    output_path = os.environ.get("GITHUB_OUTPUT")
+    if not output_path:
+        return
+    with Path(output_path).open("a") as output:
+        output.write(f"package_version={state.package_version}\n")
+        output.write(f"finalized_release_version={state.finalized_release_version}\n")
+        output.write(f"release_version_relation={state.release_version_relation}\n")
+
+
 def main(argv: list[str] | None = None) -> int:
     parser = argparse.ArgumentParser(description=__doc__)
     parser.add_argument(
@@ -103,17 +167,58 @@ def main(argv: list[str] | None = None) -> int:
     args = parser.parse_args(argv)
 
     try:
-        violations = check_repository(
-            version_manifest_url=args.version_manifest_url,
+        package_version = pyproject_version()
+        stable_version_tuple(package_version)
+    except (OSError, ValueError) as exc:
+        write_github_outputs(
+            ReleaseVersionState(
+                package_version="",
+                finalized_release_version="",
+                release_version_relation=UNKNOWN,
+            )
+        )
+        print(f"Could not read data package version: {exc}", file=sys.stderr)
+        return 1
+
+    try:
+        finalized_release_version = latest_hf_release_version(args.version_manifest_url)
+        state = ReleaseVersionState(
+            package_version=package_version,
+            finalized_release_version=finalized_release_version,
+            release_version_relation=release_version_relation(
+                package_version=package_version,
+                finalized_release_version=finalized_release_version,
+            ),
         )
     except (URLError, OSError, ValueError) as exc:
+        write_github_outputs(
+            ReleaseVersionState(
+                package_version=package_version,
+                finalized_release_version="",
+                release_version_relation=UNKNOWN,
+            )
+        )
         print(
             f"Could not check finalized HF data release version: {exc}", file=sys.stderr
         )
         return 1 if args.mode == "fail" else 0
 
+    write_github_outputs(state)
+    violations = version_violations(
+        package_version=state.package_version,
+        finalized_release_version=state.finalized_release_version,
+    )
     if not violations:
-        print("Data package version is current with the latest finalized HF release.")
+        if state.release_version_relation == AHEAD:
+            print(
+                "Data package version "
+                f"{state.package_version} is ahead of finalized HF release "
+                f"{state.finalized_release_version}."
+            )
+        else:
+            print(
+                "Data package version is current with the latest finalized HF release."
+            )
         return 0
 
     for violation in violations:

.github/scripts/promote_publication_pipeline.py

@@ -5,7 +5,6 @@
 import json
 import os
 import sys
-import tomllib
 from pathlib import Path
 
 import modal
@@ -17,15 +16,9 @@
 from policyengine_us_data.utils.run_context import (  # noqa: E402
     RunContext,
     release_version_from_bump,
-    stable_release_version,
 )
 
 
-def _current_package_version() -> str:
-    with (_REPO_ROOT / "pyproject.toml").open("rb") as file:
-        return stable_release_version(tomllib.load(file)["project"]["version"])
-
-
 def _modal_function(app_name: str, function_name: str, environment_name: str):
     if environment_name:
         return modal.Function.from_name(
@@ -61,8 +54,13 @@ def _promotion_context_from_status(context: RunContext, status: dict) -> RunCont
         raise RuntimeError("Run manifest is missing release_bump.")
     release_version = _manifest_field(manifest, "release_version")
     if not release_version:
+        if not base_release_version:
+            raise RuntimeError(
+                "Run manifest is missing base_release_version, so promotion "
+                "cannot reconstruct release_version from release_bump."
+            )
         release_version = release_version_from_bump(
-            _current_package_version(),
+            base_release_version,
             release_bump,
         )
     return RunContext.from_mapping(

.github/scripts/sync_finalized_data_release_version.py

@@ -0,0 +1,104 @@
+"""Synchronize pyproject.toml with the latest finalized HF data release."""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+import re
+import sys
+from urllib.error import URLError
+
+_SCRIPT_DIR = Path(__file__).resolve().parent
+if str(_SCRIPT_DIR) not in sys.path:
+    sys.path.insert(0, str(_SCRIPT_DIR))
+
+from check_data_release_version import (
+    BEHIND,
+    DEFAULT_VERSION_MANIFEST_URL,
+    REPO_ROOT,
+    check_repository_state,
+)
+
+
+VERSION_RE = re.compile(r'^(version\s*=\s*)"([^"]+)"', re.MULTILINE)
+
+
+def update_pyproject_version(pyproject: Path, release_version: str) -> str:
+    text = pyproject.read_text()
+    match = VERSION_RE.search(text)
+    if not match:
+        raise ValueError("Could not find project version in pyproject.toml")
+
+    current_version = match.group(2)
+    if current_version == release_version:
+        return current_version
+
+    updated = VERSION_RE.sub(rf'\1"{release_version}"', text, count=1)
+    pyproject.write_text(updated)
+    return current_version
+
+
+def sync_finalized_data_release_version(
+    root: Path | None = None,
+    *,
+    finalized_release_version: str | None = None,
+    version_manifest_url: str = DEFAULT_VERSION_MANIFEST_URL,
+) -> bool:
+    root = root or REPO_ROOT
+    state = check_repository_state(
+        root,
+        finalized_release_version=finalized_release_version,
+        version_manifest_url=version_manifest_url,
+    )
+    if state.release_version_relation != BEHIND:
+        print(
+            "No finalized data release version sync needed: "
+            f"package={state.package_version}, "
+            f"finalized={state.finalized_release_version}, "
+            f"relation={state.release_version_relation}."
+        )
+        return False
+
+    previous_version = update_pyproject_version(
+        root / "pyproject.toml",
+        state.finalized_release_version,
+    )
+    print(
+        "Synchronized pyproject.toml with finalized HF data release: "
+        f"{previous_version} -> {state.finalized_release_version}."
+    )
+    return True
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--version-manifest-url",
+        default=os.environ.get(
+            "US_DATA_VERSION_MANIFEST_URL", DEFAULT_VERSION_MANIFEST_URL
+        ),
+    )
+    parser.add_argument(
+        "--finalized-release-version",
+        default=os.environ.get("US_DATA_FINALIZED_RELEASE_VERSION"),
+        help="Already-resolved finalized HF release version to sync to.",
+    )
+    args = parser.parse_args(argv)
+
+    try:
+        sync_finalized_data_release_version(
+            finalized_release_version=args.finalized_release_version,
+            version_manifest_url=args.version_manifest_url,
+        )
+    except (URLError, OSError, ValueError) as exc:
+        print(
+            f"Could not synchronize finalized HF data release version: {exc}",
+            file=sys.stderr,
+        )
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

.github/workflows/pipeline.yaml

@@ -67,11 +67,6 @@ on:
         description: "Number of Modal workers for parallel matrix build"
         default: "50"
         type: string
-      allow_stale_policyengine_us:
-        description: "Allow production build when policyengine-us lags the latest PyPI release"
-        default: false
-        type: boolean
-
 concurrency:
   group: pipeline-${{ github.run_id }}-${{ github.run_attempt }}
   cancel-in-progress: false
@@ -103,11 +98,6 @@ jobs:
           RELEASE_BUMP: ${{ inputs.release_bump || '' }}
         run: python .github/scripts/resolve_run_context.py
 
-      - name: Require current PolicyEngine US dependency
-        env:
-          POLICYENGINE_US_ALLOW_STALE: ${{ inputs.allow_stale_policyengine_us }}
-        run: python .github/scripts/check_policyengine_us_dependency.py --mode fail
-
       - name: Require pyproject.toml to match finalized HF release base
         run: python .github/scripts/check_data_release_version.py --mode fail