Add versions argument to Simulation

.github/workflows/code_changes.yaml

@@ -21,6 +21,9 @@ jobs:
               args: ". -l 79 --check"
   Test:
       runs-on: ubuntu-latest
+      permissions:
+        contents: "read"
+        id-token: "write"
       steps:
           - name: Checkout repo
             uses: actions/checkout@v2
@@ -31,6 +34,10 @@ jobs:
             uses: actions/setup-python@v2
             with:
                 python-version: '3.11'
+          - uses: "google-github-actions/auth@v2"
+            with:
+              workload_identity_provider: "projects/322898545428/locations/global/workloadIdentityPools/policyengine-research-id-pool/providers/prod-github-provider"
+              service_account: "policyengine-research@policyengine-research.iam.gserviceaccount.com"
 
           - name: Install package
             run: uv pip install .[dev] --system

.github/workflows/publish_documentation.yaml

@@ -7,6 +7,9 @@ on:
 
 jobs:
   Publish:
+    permissions:
+      contents: "read"
+      id-token: "write"
     runs-on: ubuntu-latest
     steps:
       - name: Checkout repo
@@ -15,6 +18,10 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: 3.12
+      - uses: "google-github-actions/auth@v2"
+        with:
+          workload_identity_provider: "projects/322898545428/locations/global/workloadIdentityPools/policyengine-research-id-pool/providers/prod-github-provider"
+          service_account: "policyengine-research@policyengine-research.iam.gserviceaccount.com"
       - name: Publish a git tag
         run: ".github/publish-git-tag.sh || true"
       - name: Install package

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: minor
+  changes:
+    added:
+    - Error handling for data and package version mismatches.

policyengine/constants.py

@@ -2,40 +2,40 @@
 
 from policyengine_core.data import Dataset
 from policyengine.utils.data_download import download
+from typing import Tuple
 
-# Datasets
-ENHANCED_FRS = "hf://policyengine/policyengine-uk-data/enhanced_frs_2022_23.h5"
-FRS = "hf://policyengine/policyengine-uk-data/frs_2022_23.h5"
-ENHANCED_CPS = "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5"
-CPS = "hf://policyengine/policyengine-us-data/cps_2023.h5"
-POOLED_CPS = "hf://policyengine/policyengine-us-data/pooled_3_year_cps_2023.h5"
 
-
-def get_default_dataset(country: str, region: str):
+def get_default_dataset(country: str, region: str) -> Tuple[Dataset, str]:
     if country == "uk":
-        data_file = download(
+        data_file, version = download(
             filepath="enhanced_frs_2022_23.h5",
             huggingface_repo="policyengine-uk-data",
             gcs_bucket="policyengine-uk-data-private",
+            return_version=True,
         )
         time_period = None
     elif country == "us":
         if region is not None and region != "us":
-            data_file = download(
+            data_file, version = download(
                 filepath="pooled_3_year_cps_2023.h5",
                 huggingface_repo="policyengine-us-data",
                 gcs_bucket="policyengine-us-data",
+                return_version=True,
             )
             time_period = 2023
         else:
-            data_file = download(
+            data_file, version = download(
                 filepath="cps_2023.h5",
                 huggingface_repo="policyengine-us-data",
                 gcs_bucket="policyengine-us-data",
+                return_version=True,
             )
             time_period = 2023
 
-    return Dataset.from_file(
-        file_path=data_file,
-        time_period=time_period,
+    return (
+        Dataset.from_file(
+            file_path=data_file,
+            time_period=time_period,
+        ),
+        version,
     )

policyengine/simulation.py

@@ -18,6 +18,7 @@
     Simulation as UKSimulation,
     Microsimulation as UKMicrosimulation,
 )
+from importlib import metadata
 import h5py
 from pathlib import Path
 import pandas as pd
@@ -62,6 +63,14 @@ class SimulationOptions(BaseModel):
         False,
         description="Whether to include tax-benefit cliffs in the simulation analyses. If True, cliffs will be included.",
     )
+    package_versions: Dict[str, str] | None = Field(
+        None,
+        description="The versions of the packages used in the simulation. If not provided, the current package versions will be used. If provided, this package will throw an error if the package versions do not match.",
+    )
+    data_versions: Dict[str, str] | None = Field(
+        None,
+        description="The versions of the data used in the simulation. If not provided, the current data versions will be used. If provided, this package will throw an error if the data versions do not match.",
+    )
 
 
 class Simulation:
@@ -73,6 +82,7 @@ class Simulation:
     """The baseline tax-benefit simulation."""
     reform_simulation: CountrySimulation | None = None
     """The reform tax-benefit simulation."""
+    data_versions: Dict[str, str] | None = None
 
     def __init__(self, **options: SimulationOptions):
         self.options = SimulationOptions(**options)
@@ -113,8 +123,9 @@ def _add_output_functions(self):
                             )
 
     def _set_data(self):
+        self.data_versions = {}
         if self.options.data is None:
-            self.options.data = get_default_dataset(
+            self.options.data, version = get_default_dataset(
                 country=self.options.country,
                 region=self.options.region,
             )
@@ -135,13 +146,17 @@ def _set_data(self):
                         -1
                     ].split("/", 2)
 
-                file_path = download(
+                file_path, version = download(
                     filepath=filename,
                     huggingface_org=hf_org,
                     huggingface_repo=hf_repo,
                     gcs_bucket=bucket,
+                    return_version=True,
                 )
                 filename = str(Path(file_path))
+            else:
+                # If it's a local file, we can't infer the version.
+                version = None
             if "cps_2023" in filename:
                 time_period = 2023
             else:
@@ -150,6 +165,8 @@ def _set_data(self):
                 filename, time_period=time_period
             )
 
+        self.data_versions[self.options.data.file_path.name] = version
+
     def _initialise_simulations(self):
         self.baseline_simulation = self._initialise_simulation(
             scope=self.options.scope,
@@ -327,3 +344,33 @@ def _apply_region_to_simulation(
                 )
 
         return simulation
+
+    def check_package_versions(self) -> None:
+        """
+        Check the package versions of the simulation against the current package versions.
+        """
+        if self.options.package_versions is not None:
+            for package, version in self.options.package_versions.items():
+                try:
+                    installed_version = metadata.version(package)
+                except metadata.PackageNotFoundError:
+                    raise ValueError(f"Package {package} not found.")
+                if installed_version != version:
+                    raise ValueError(
+                        f"Package {package} version {installed_version} does not match expected version {version}."
+                    )
+
+    def check_data_versions(self) -> None:
+        """
+        Check the data versions of the simulation against the current data versions.
+        """
+        if self.options.data_versions is not None:
+            for file, version in self.options.data_versions.items():
+                if file not in self.data_versions:
+                    raise ValueError(
+                        f"Data file {file} not found in simulation."
+                    )
+                if self.data_versions[file] != version:
+                    raise ValueError(
+                        f"Data file {file} version {self.data_versions[file]} does not match expected version {version}."
+                    )

policyengine/utils/data_download.py

@@ -4,6 +4,8 @@
 from policyengine.utils.huggingface import download_from_hf
 from policyengine.utils.google_cloud_bucket import download_file_from_gcs
 from pydantic import BaseModel
+import json
+from typing import Tuple
 
 
 class DataFile(BaseModel):
@@ -18,7 +20,8 @@ def download(
     huggingface_repo: str = None,
     gcs_bucket: str = None,
     huggingface_org: str = "policyengine",
-):
+    return_version: bool = False,
+) -> str | Tuple[str, str]:
     data_file = DataFile(
         filepath=filepath,
         huggingface_org=huggingface_org,
@@ -31,12 +34,24 @@ def download(
     if data_file.huggingface_repo is not None:
         logging.info("Using Hugging Face for download.")
         try:
-            return download_from_hf(
+            data = download_from_hf(
                 repo=data_file.huggingface_org
                 + "/"
                 + data_file.huggingface_repo,
                 repo_filename=data_file.filepath,
             )
+            if return_version:
+                version_file = download_from_hf(
+                    repo=data_file.huggingface_org
+                    + "/"
+                    + data_file.huggingface_repo,
+                    repo_filename="version.json",
+                    return_version=True,
+                )
+                with open(version_file, "r") as f:
+                    version = json.load(f).get("version")
+                return data, version
+            return data
         except:
             logging.info("Failed to download from Hugging Face.")
 
@@ -47,6 +62,15 @@ def download(
             file_name=filepath,
             destination_path=filepath,
         )
+        if return_version:
+            download_file_from_gcs(
+                bucket_name=data_file.gcs_bucket,
+                file_name="version.json",
+                destination_path="version.json",
+            )
+            with open("version.json", "r") as f:
+                version = json.load(f).get("version")
+            return filepath, version
         return filepath
 
     raise ValueError(

tests/country/test_uk.py

@@ -21,3 +21,45 @@ def test_uk_macro_comparison():
     )
 
     sim.calculate_economy_comparison()
+
+
+def test_uk_macro_bad_package_versions_fail():
+    from policyengine import Simulation
+
+    try:
+        sim = Simulation(
+            scope="macro",
+            country="uk",
+            reform={
+                "gov.hmrc.income_tax.allowances.personal_allowance.amount": 15_000,
+            },
+            package_versions={
+                "policyengine-uk": "-1.0.0",
+            },
+        )
+        raise ValueError(
+            "Simulation should have failed with a bad package version."
+        )
+    except:
+        pass
+
+
+def test_uk_macro_bad_data_versions_fail():
+    from policyengine import Simulation
+
+    try:
+        sim = Simulation(
+            scope="macro",
+            country="uk",
+            reform={
+                "gov.hmrc.income_tax.allowances.personal_allowance.amount": 15_000,
+            },
+            data_versions={
+                "enhanced_frs_2022_23.h5": "-1.0.0",
+            },
+        )
+        raise ValueError(
+            "Simulation should have failed with a bad data version."
+        )
+    except:
+        pass