Cache197

policyengine_core/simulations/sim_macro_cache.py

@@ -0,0 +1,50 @@
+import shutil
+from pathlib import Path
+import h5py
+from numpy.typing import ArrayLike
+
+
+class Singleton(type):
+    _instances = {}
+
+    def __call__(cls, *args, **kwargs):
+        if cls not in cls._instances:
+            cls._instances[cls] = super(Singleton, cls).__call__(*args, **kwargs)
+        return cls._instances[cls]
+
+
+class SimulationMacroCache(metaclass=Singleton):
+    def __init__(self):
+        self.core_version = None
+        self.country_version = None
+        self.cache_file_path = None
+
+    def set_version(self, core_version: str, country_version: {}):
+        self.core_version = core_version
+        self.country_version = country_version
+
+    def set_cache_path(self, parent_path: Path, dataset_name: str):
+        self.cache_file_path = parent_path / f"{dataset_name}_variable_cache"
+
+    def get_cache_value(self, version: str, country_version: {}, cache_file_path: Path):
+        with h5py.File(cache_file_path, "r") as f:
+            if "metadata:core_version" in f and "metadata:country_version" in f:
+                # Validate version is correct, otherwise flush the cache
+                if f["metadata:core_version"][()] != version or f["metadata:country_version"][()] != country_version:
+                    self.clear_cache(cache_file_path)
+                    return None
+            else:
+                self.clear_cache(cache_file_path)
+                return None
+            return f["values"][()]
+
+    def set_cache_value(self, core_version: str, country_version: {}, cache_file_path: Path, value: ArrayLike):
+        self.set_version(core_version, country_version)
+        with h5py.File(cache_file_path, "w") as f:
+            f.create_dataset("values", data=value)
+            f.create_dataset("metadata:core_version", data=self.core_version)
+            f.create_dataset("metadata:country_version", data=self.country_version)
+        return "cache set successfully"
+
+    def clear_cache(self, cache_file_path: Path):
+        shutil.rmtree(cache_file_path)

policyengine_core/simulations/simulation.py

@@ -4,6 +4,7 @@
 import numpy
 import numpy as np
 import pandas as pd
+import importlib.metadata
 from numpy.typing import ArrayLike
 
 from policyengine_core import commons, periods
@@ -35,6 +36,7 @@
 from policyengine_core.variables import Variable, QuantityType
 from policyengine_core.reforms.reform import Reform
 from policyengine_core.parameters import get_parameter
+from policyengine_core.simulations.sim_macro_cache import SimulationMacroCache
 
 
 class Simulation:
@@ -75,6 +77,12 @@ class Simulation:
     macro_cache_write: bool = True
     """Whether to write to the macro cache."""
 
+    core_version: str = None
+    """Version number of the simulation core."""
+
+    country_version: Dict[str, Any] = None
+    """Version number of the country packages."""
+
     def __init__(
         self,
         tax_benefit_system: "TaxBenefitSystem" = None,
@@ -198,6 +206,24 @@ def __init__(
             self.baseline = None
 
         self.parent_branch = None
+        self.core_version = importlib.metadata.version("policyengine-core")
+
+        # TODO(SylviaDu99)
+        COUNTRIES = ("uk", "us", "ca", "ng", "il")
+        COUNTRY_PACKAGE_NAMES = (
+            "policyengine_uk",
+            "policyengine_us",
+            "policyengine_canada",
+            "policyengine_ng",
+            "policyengine_il",
+        )
+        try:
+            self.country_version = {
+                country: importlib.metadata.version(package_name)
+                for country, package_name in zip(COUNTRIES, COUNTRY_PACKAGE_NAMES)
+            }
+        except:
+            self.country_version = {country: "0.0.0" for country in COUNTRIES}
 
     def apply_reform(self, reform: Union[tuple, Reform]):
         if isinstance(reform, tuple):
@@ -568,7 +594,8 @@ def _calculate(
 
         if alternate_period_handling:
             if cache_path is not None:
-                self._set_macro_cache_value(cache_path, values)
+                message = self._set_macro_cache_value(cache_path, values)
+                print(message is not None)
             return values
 
         self._check_period_consistency(period, variable)
@@ -1379,8 +1406,7 @@ def _get_macro_cache_value(
         """
         if not self.macro_cache_read or self.tax_benefit_system.data_modified:
             return None
-        with h5py.File(cache_file_path, "r") as f:
-            return f["values"][()]
+        return SimulationMacroCache().get_cache_value(self.version, cache_file_path)
 
     def _set_macro_cache_value(
         self,
@@ -1392,8 +1418,8 @@ def _set_macro_cache_value(
         """
         if not self.macro_cache_write or self.tax_benefit_system.data_modified:
             return None
-        with h5py.File(cache_file_path, "w") as f:
-            f.create_dataset("values", data=value)
+        message = SimulationMacroCache().set_cache_value(self.version, cache_file_path, value)
+        return message
 
 
 class NpEncoder(json.JSONEncoder):

tests/core/test_simulations.py

@@ -1,6 +1,6 @@
 from policyengine_core.country_template.situation_examples import single
 from policyengine_core.simulations import SimulationBuilder
-
+from policyengine_core.country_template.data.datasets import country_template_dataset
 
 def test_calculate_full_tracer(tax_benefit_system):
     simulation = SimulationBuilder().build_default_simulation(
@@ -61,3 +61,18 @@ def test_get_memory_usage(tax_benefit_system):
     memory_usage = simulation.get_memory_usage(variables=["salary"])
     assert memory_usage["total_nb_bytes"] > 0
     assert len(memory_usage["by_variable"]) == 1
+
+
+def test_version(tax_benefit_system):
+    simulation = SimulationBuilder().build_from_entities(
+        tax_benefit_system, single
+    )
+    assert simulation.core_version is not None
+    assert simulation.country_version is not None
+
+
+def test_macro_cache(tax_benefit_system):
+    simulation = SimulationBuilder().build_from_entities(
+        tax_benefit_system, single,
+    )
+    simulation.calculate("disposable_income", "2017-01")