From dc90992be7a6d293e7a96b154a8f0260e64f35ae Mon Sep 17 00:00:00 2001 From: Nikhil Woodruff Date: Mon, 2 Jun 2025 12:30:49 +0100 Subject: [PATCH 1/7] Fix Pass data and country package versions to APIv2 #2500 --- changelog_entry.yaml | 4 ++ .../jobs/calculate_economy_simulation_job.py | 52 +++++++++++++++++-- 2 files changed, 51 insertions(+), 5 deletions(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index e69de29bb..9ab62b197 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -0,0 +1,4 @@ +- bump: patch + changes: + fixed: + - Pass country and data versions to APIv2. diff --git a/policyengine_api/jobs/calculate_economy_simulation_job.py b/policyengine_api/jobs/calculate_economy_simulation_job.py index f74a1619a..13fdb0ff0 100644 --- a/policyengine_api/jobs/calculate_economy_simulation_job.py +++ b/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, + 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, + }, +} + +uk_dataset_version = get_latest_commit_tag( + repo_id="policyengine/policyengine-us-data", + repo_type="model", +) +us_dataset_version = get_latest_commit_tag( + repo_id="policyengine/policyengine-uk-data", + 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 @@ def run( 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 @@ def _create_simulation_uk( 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 @@ def _create_simulation_us( 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 @@ def _create_simulation_us( 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 @@ def _create_simulation_us( 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 @@ def _setup_sim_options( 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 @@ def _setup_sim_options( "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: From 6a9ffe4ed707b6e7308eec2bee92f6c8ba93ea1d Mon Sep 17 00:00:00 2001 From: Nikhil Woodruff Date: Mon, 2 Jun 2025 12:41:48 +0100 Subject: [PATCH 2/7] Add missing file --- .../jobs/calculate_economy_simulation_job.py | 2 +- policyengine_api/utils/hugging_face.py | 108 ++++++++++++++++++ 2 files changed, 109 insertions(+), 1 deletion(-) create mode 100644 policyengine_api/utils/hugging_face.py diff --git a/policyengine_api/jobs/calculate_economy_simulation_job.py b/policyengine_api/jobs/calculate_economy_simulation_job.py index 13fdb0ff0..6de8c526f 100644 --- a/policyengine_api/jobs/calculate_economy_simulation_job.py +++ b/policyengine_api/jobs/calculate_economy_simulation_job.py @@ -28,8 +28,8 @@ from policyengine_core.simulations import Microsimulation from policyengine_core.tools.hugging_face import ( download_huggingface_dataset, - get_latest_commit_tag, ) +from policyengine_api.utils.hugging_face import get_latest_commit_tag import h5py from policyengine_us import Microsimulation diff --git a/policyengine_api/utils/hugging_face.py b/policyengine_api/utils/hugging_face.py new file mode 100644 index 000000000..e3f86fb5c --- /dev/null +++ b/policyengine_api/utils/hugging_face.py @@ -0,0 +1,108 @@ +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() + + # Get list of commits + commits = api.list_repo_commits(repo_id=repo_id, repo_type=repo_type) + + 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).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 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. + """ + # Attempt connection to Hugging Face model_info endpoint + # (https://huggingface.co/docs/huggingface_hub/v0.26.5/en/package_reference/hf_api#huggingface_hub.HfApi.model_info) + # Attempt to fetch model info to determine if repo is private + # A RepositoryNotFoundError & 401 likely means the repo is private, + # but this error will also surface for public repos with malformed URL, etc. + try: + fetched_model_info: ModelInfo = model_info(repo) + is_repo_private: bool = fetched_model_info.private + except RepositoryNotFoundError as e: + # If this error type arises, it's likely the repo is private; see docs above + is_repo_private = True + pass + except Exception as e: + # Otherwise, there probably is just a download error + raise Exception( + f"Unable to download dataset {repo_filename} from Hugging Face. This may be because the repo " + + f"is private, the URL is malformed, or the dataset does not exist. The full error is {traceback.format_exc()}" + ) + + 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 From c3859ba32a3652c9b380436c86810676497ecffc Mon Sep 17 00:00:00 2001 From: Nikhil Woodruff Date: Mon, 2 Jun 2025 13:05:25 +0100 Subject: [PATCH 3/7] Adjust UK address --- policyengine_api/jobs/calculate_economy_simulation_job.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/policyengine_api/jobs/calculate_economy_simulation_job.py b/policyengine_api/jobs/calculate_economy_simulation_job.py index 6de8c526f..89495ad47 100644 --- a/policyengine_api/jobs/calculate_economy_simulation_job.py +++ b/policyengine_api/jobs/calculate_economy_simulation_job.py @@ -65,8 +65,8 @@ repo_type="model", ) us_dataset_version = get_latest_commit_tag( - repo_id="policyengine/policyengine-uk-data", - repo_type="model", + repo_id="policyengine/policyengine-uk-data-private", + repo_type="dataset", ) for dataset in datasets["uk"]: From a581f32b7b7dfbc5c92aa1b411b6232fe36582ac Mon Sep 17 00:00:00 2001 From: Nikhil Woodruff Date: Mon, 2 Jun 2025 14:10:04 +0100 Subject: [PATCH 4/7] Typo --- policyengine_api/jobs/calculate_economy_simulation_job.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/policyengine_api/jobs/calculate_economy_simulation_job.py b/policyengine_api/jobs/calculate_economy_simulation_job.py index 89495ad47..9eab7d623 100644 --- a/policyengine_api/jobs/calculate_economy_simulation_job.py +++ b/policyengine_api/jobs/calculate_economy_simulation_job.py @@ -66,7 +66,7 @@ ) us_dataset_version = get_latest_commit_tag( repo_id="policyengine/policyengine-uk-data-private", - repo_type="dataset", + repo_type="model", ) for dataset in datasets["uk"]: From af84b97dc1ff81cdfe112655c50e66a4bca2faae Mon Sep 17 00:00:00 2001 From: Nikhil Woodruff Date: Mon, 2 Jun 2025 14:12:34 +0100 Subject: [PATCH 5/7] Use token --- policyengine_api/utils/hugging_face.py | 54 ++++++++++++++++---------- 1 file changed, 34 insertions(+), 20 deletions(-) diff --git a/policyengine_api/utils/hugging_face.py b/policyengine_api/utils/hugging_face.py index e3f86fb5c..7fc5cb105 100644 --- a/policyengine_api/utils/hugging_face.py +++ b/policyengine_api/utils/hugging_face.py @@ -21,8 +21,16 @@ def get_latest_commit_tag(repo_id, repo_type="model"): """ api = HfApi() + is_repo_private = 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) + commits = api.list_repo_commits( + repo_id=repo_id, repo_type=repo_type, token=authentication_token + ) if not commits: return None @@ -30,7 +38,9 @@ def get_latest_commit_tag(repo_id, repo_type="model"): 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).tags + 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: @@ -40,6 +50,27 @@ def get_latest_commit_tag(repo_id, repo_type="model"): return None +def 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, @@ -55,24 +86,7 @@ def download_huggingface_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. """ - # Attempt connection to Hugging Face model_info endpoint - # (https://huggingface.co/docs/huggingface_hub/v0.26.5/en/package_reference/hf_api#huggingface_hub.HfApi.model_info) - # Attempt to fetch model info to determine if repo is private - # A RepositoryNotFoundError & 401 likely means the repo is private, - # but this error will also surface for public repos with malformed URL, etc. - try: - fetched_model_info: ModelInfo = model_info(repo) - is_repo_private: bool = fetched_model_info.private - except RepositoryNotFoundError as e: - # If this error type arises, it's likely the repo is private; see docs above - is_repo_private = True - pass - except Exception as e: - # Otherwise, there probably is just a download error - raise Exception( - f"Unable to download dataset {repo_filename} from Hugging Face. This may be because the repo " - + f"is private, the URL is malformed, or the dataset does not exist. The full error is {traceback.format_exc()}" - ) + is_repo_private = is_repo_private(repo) authentication_token: str = None if is_repo_private: From de0784ebd6a5b8edff993a0f93191ce609504702 Mon Sep 17 00:00:00 2001 From: Nikhil Woodruff Date: Mon, 2 Jun 2025 14:25:40 +0100 Subject: [PATCH 6/7] Rename function --- policyengine_api/utils/hugging_face.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/policyengine_api/utils/hugging_face.py b/policyengine_api/utils/hugging_face.py index 7fc5cb105..e882eb300 100644 --- a/policyengine_api/utils/hugging_face.py +++ b/policyengine_api/utils/hugging_face.py @@ -21,7 +21,7 @@ def get_latest_commit_tag(repo_id, repo_type="model"): """ api = HfApi() - is_repo_private = is_repo_private(repo_id) + is_repo_private = check_is_repo_private(repo_id) authentication_token: str = None if is_repo_private: @@ -50,7 +50,7 @@ def get_latest_commit_tag(repo_id, repo_type="model"): return None -def is_repo_private(repo: str) -> bool: +def check_is_repo_private(repo: str) -> bool: """ Check if a Hugging Face repository is private. @@ -86,7 +86,7 @@ def download_huggingface_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 = is_repo_private(repo) + is_repo_private = check_is_repo_private(repo) authentication_token: str = None if is_repo_private: From 2653504b2f979c07c68f379cd117cab112d5bf40 Mon Sep 17 00:00:00 2001 From: Nikhil Woodruff Date: Mon, 2 Jun 2025 15:15:06 +0100 Subject: [PATCH 7/7] Fix accidental uk/us mixup --- policyengine_api/jobs/calculate_economy_simulation_job.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/policyengine_api/jobs/calculate_economy_simulation_job.py b/policyengine_api/jobs/calculate_economy_simulation_job.py index 9eab7d623..22ae25686 100644 --- a/policyengine_api/jobs/calculate_economy_simulation_job.py +++ b/policyengine_api/jobs/calculate_economy_simulation_job.py @@ -60,11 +60,11 @@ }, } -uk_dataset_version = get_latest_commit_tag( +us_dataset_version = get_latest_commit_tag( repo_id="policyengine/policyengine-us-data", repo_type="model", ) -us_dataset_version = get_latest_commit_tag( +uk_dataset_version = get_latest_commit_tag( repo_id="policyengine/policyengine-uk-data-private", repo_type="model", )