Fix Pass data and country package versions to APIv2 #2500

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: patch
+  changes:
+    fixed:
+    - Pass country and data versions to APIv2.

policyengine_api/jobs/calculate_economy_simulation_job.py

@@ -26,12 +26,16 @@
     compute_difference,
 )
 from policyengine_core.simulations import Microsimulation
-from policyengine_core.tools.hugging_face import download_huggingface_dataset
+from policyengine_core.tools.hugging_face import (
+    download_huggingface_dataset,
+)
+from policyengine_api.utils.hugging_face import get_latest_commit_tag
 import h5py
 
 from policyengine_us import Microsimulation
 from policyengine_uk import Microsimulation
 import logging
+import huggingface_hub
 
 load_dotenv()
 
@@ -44,6 +48,34 @@
 CPS = "hf://policyengine/policyengine-us-data/cps_2023.h5"
 POOLED_CPS = "hf://policyengine/policyengine-us-data/pooled_3_year_cps_2023.h5"
 
+datasets = {
+    "uk": {
+        "enhanced_frs": ENHANCED_FRS,
+        "frs": FRS,
+    },
+    "us": {
+        "enhanced_cps": ENHANCED_CPS,
+        "cps": CPS,
+        "pooled_cps": POOLED_CPS,
+    },
+}
+
+us_dataset_version = get_latest_commit_tag(
+    repo_id="policyengine/policyengine-us-data",
+    repo_type="model",
+)
+uk_dataset_version = get_latest_commit_tag(
+    repo_id="policyengine/policyengine-uk-data-private",
+    repo_type="model",
+)
+
+for dataset in datasets["uk"]:
+    datasets["uk"][dataset] = f"{datasets['uk'][dataset]}@{uk_dataset_version}"
+
+for dataset in datasets["us"]:
+    datasets["us"][dataset] = f"{datasets['us'][dataset]}@{us_dataset_version}"
+
+
 check_against_api_v2 = (
     os.environ.get("GOOGLE_APPLICATION_CREDENTIALS") is not None
 )
@@ -189,6 +221,12 @@
                     time_period=time_period,
                     region=region,
                     dataset=dataset,
+                    model_version=COUNTRY_PACKAGE_VERSIONS[country_id],
+                    data_version=(
+                        uk_dataset_version
+                        if country_id == "uk"
+                        else us_dataset_version
+                    ),
                 )
 
                 try:
@@ -454,7 +492,7 @@
 
         simulation = CountryMicrosimulation(
             reform=reform,
-            dataset=ENHANCED_FRS,
+            dataset=datasets["uk"]["enhanced_frs"],
         )
         simulation.default_calculation_period = time_period
         if region != "uk":
@@ -514,7 +552,7 @@
         if dataset in DATASETS:
             print(f"Running simulation using {dataset} dataset")
 
-            sim_options["dataset"] = ENHANCED_CPS
+            sim_options["dataset"] = datasets["us"]["enhanced_cps"]
 
         # Handle region settings
         if region != "us":
@@ -526,7 +564,7 @@
             if "dataset" in sim_options:
                 filter_dataset = sim_options["dataset"]
             else:
-                filter_dataset = POOLED_CPS
+                filter_dataset = datasets["us"]["pooled_cps"]
 
             # Run sim to filter by region
             region_sim = Microsimulation(
@@ -547,7 +585,7 @@
                 sim_options["dataset"] = df[state_code == region.upper()]
 
         if dataset == "default" and region == "us":
-            sim_options["dataset"] = CPS
+            sim_options["dataset"] = datasets["us"]["cps"]
 
         # Return completed simulation
         return Microsimulation(**sim_options)
@@ -723,6 +761,8 @@
         dataset: str,
         time_period: str,
         scope: Literal["macro", "household"] = "macro",
+        model_version: str | None = None,
+        data_version: str | None = None,
     ) -> dict[str, Any]:
         """
         Set up the simulation options for the APIv2 job.
@@ -738,6 +778,8 @@
             "data": self._setup_data(
                 dataset=dataset, country_id=country_id, region=region
             ),
+            "model_version": model_version,
+            "data_version": data_version,
         }
 
     def _setup_region(self, country_id: str, region: str) -> str:

policyengine_api/utils/hugging_face.py

@@ -0,0 +1,122 @@
+from huggingface_hub import (
+    hf_hub_download,
+    model_info,
+    ModelInfo,
+    HfApi,
+)
+from huggingface_hub.errors import RepositoryNotFoundError
+from getpass import getpass
+import os
+import warnings
+import traceback
+
+with warnings.catch_warnings():
+    warnings.simplefilter("ignore")
+
+
+def get_latest_commit_tag(repo_id, repo_type="model"):
+    """
+    Get the tag associated with the latest commit in a HF repo.
+    Returns the tag name or None if no tag is associated.
+    """
+    api = HfApi()
+
+    is_repo_private = check_is_repo_private(repo_id)
+
+    authentication_token: str = None
+    if is_repo_private:
+        authentication_token: str = get_or_prompt_hf_token()
+
+    # Get list of commits
+    commits = api.list_repo_commits(
+        repo_id=repo_id, repo_type=repo_type, token=authentication_token
+    )
+
+    if not commits:
+        return None
+
+    latest_commit = commits[0]  # Most recent commit is first
+
+    # Get all tags in the repository
+    tags = api.list_repo_refs(
+        repo_id=repo_id, repo_type=repo_type, token=authentication_token
+    ).tags
+
+    # Find tag that points to the latest commit
+    for tag in tags:
+        if tag.target_commit == latest_commit.commit_id:
+            return tag.ref.replace("refs/tags/", "")
+
+    return None
+
+
+def check_is_repo_private(repo: str) -> bool:
+    """
+    Check if a Hugging Face repository is private.
+
+    Args:
+        repo (str): The Hugging Face repo name, in format "{org}/{repo}".
+
+    Returns:
+        bool: True if the repo is private, False otherwise.
+    """
+    try:
+        fetched_model_info: ModelInfo = model_info(repo)
+        return fetched_model_info.private
+    except RepositoryNotFoundError:
+        return True  # If repo not found, assume it's private
+    except Exception as e:
+        raise Exception(
+            f"Unable to check if repo {repo} is private. The full error is {traceback.format_exc()}"
+        )
+
+
+def download_huggingface_dataset(
+    repo: str,
+    repo_filename: str,
+    version: str = None,
+    local_dir: str | None = None,
+):
+    """
+    Download a dataset from the Hugging Face Hub.
+
+    Args:
+        repo (str): The Hugging Face repo name, in format "{org}/{repo}".
+        repo_filename (str): The filename of the dataset.
+        version (str, optional): The version of the dataset. Defaults to None.
+        local_dir (str, optional): The local directory to save the dataset to. Defaults to None.
+    """
+    is_repo_private = check_is_repo_private(repo)
+
+    authentication_token: str = None
+    if is_repo_private:
+        authentication_token: str = get_or_prompt_hf_token()
+
+    return hf_hub_download(
+        repo_id=repo,
+        repo_type="model",
+        filename=repo_filename,
+        revision=version,
+        token=authentication_token,
+        local_dir=local_dir,
+    )
+
+
+def get_or_prompt_hf_token() -> str:
+    """
+    Either get the Hugging Face token from the environment,
+    or prompt the user for it and store it in the environment.
+
+    Returns:
+        str: The Hugging Face token.
+    """
+
+    token = os.environ.get("HUGGING_FACE_TOKEN")
+    if token is None:
+        token = getpass(
+            "Enter your Hugging Face token (or set HUGGING_FACE_TOKEN environment variable): "
+        )
+        # Optionally store in env for subsequent calls in same session
+        os.environ["HUGGING_FACE_TOKEN"] = token
+
+    return token