Invalidate SPI income model cache by release

Author: MaxGhenis

policyengine_uk_data/datasets/imputations/income.py

@@ -125,7 +125,20 @@ def generate_spi_table(
 IMPUTATIONS = INCOME_COMPONENTS + ["gift_aid", "charitable_investment_gifts"]
 
 
-INCOME_MODEL_PATH = STORAGE_FOLDER / "income.pkl"
+INCOME_MODEL_METADATA = {
+    "spi_release_name": SPI_RELEASE_NAME,
+    "spi_tab_filename": SPI_TAB_FILENAME,
+    "imputations": tuple(IMPUTATIONS),
+}
+INCOME_MODEL_PATH = STORAGE_FOLDER / f"income_{SPI_RELEASE_NAME}.pkl"
+
+
+def _income_model_matches_current_release(model) -> bool:
+    if getattr(model, "metadata", {}) != INCOME_MODEL_METADATA:
+        return False
+
+    cached_outputs = set(getattr(model.model, "imputed_variables", []))
+    return cached_outputs == set(IMPUTATIONS)
 
 
 def save_imputation_models():
@@ -138,6 +151,7 @@ def save_imputation_models():
     from policyengine_uk_data.utils import QRF
 
     income = QRF()
+    income.metadata = INCOME_MODEL_METADATA
     spi = pd.read_csv(SPI_TAB_FOLDER / SPI_TAB_FILENAME, delimiter="\t")
     spi = generate_spi_table(spi)
     spi = spi[PREDICTORS + IMPUTATIONS]
@@ -150,10 +164,9 @@ def create_income_model(overwrite_existing: bool = False):
     """
     Create or load income imputation model.
 
-    If a cached model exists and its trained output columns don't match the
-    current ``IMPUTATIONS`` list, the cache is discarded and the model is
-    retrained. This handles the case where ``IMPUTATIONS`` is extended in
-    code but an older pickle is still on disk.
+    If a cached model exists and its training metadata or output columns don't
+    match the current SPI release and ``IMPUTATIONS`` list, the cache is
+    discarded and the model is retrained.
 
     Args:
         overwrite_existing: Whether to retrain model if it exists.
@@ -165,10 +178,9 @@ def create_income_model(overwrite_existing: bool = False):
 
     if INCOME_MODEL_PATH.exists() and not overwrite_existing:
         cached = QRF(file_path=INCOME_MODEL_PATH)
-        cached_outputs = set(getattr(cached.model, "imputed_variables", []))
-        if cached_outputs == set(IMPUTATIONS):
+        if _income_model_matches_current_release(cached):
             return cached
-        # Cached model was trained against a different output set; retrain.
+        # Cached model was trained against a different SPI release or output set.
     return save_imputation_models()
 
 

policyengine_uk_data/tests/test_spi_build.py

@@ -17,6 +17,8 @@
 
 import importlib.util
 import inspect
+import pickle
+from types import SimpleNamespace
 
 import numpy as np
 import pandas as pd
@@ -206,3 +208,68 @@ def test_income_projection_uses_current_spi_release():
 
     assert incomes_projection.SPI_DATASET.endswith(SPI_H5_FILENAME)
     assert incomes_projection.SPI_FISCAL_YEAR == SPI_FISCAL_YEAR
+
+
+def test_income_model_cache_is_release_scoped():
+    from policyengine_uk_data.datasets.imputations.income import (
+        INCOME_MODEL_PATH,
+    )
+    from policyengine_uk_data.datasets.spi import SPI_RELEASE_NAME
+
+    assert INCOME_MODEL_PATH.name == f"income_{SPI_RELEASE_NAME}.pkl"
+
+
+def test_income_model_cache_rejects_stale_spi_release(tmp_path, monkeypatch):
+    from policyengine_uk_data.datasets.imputations import income as income_module
+
+    cache = tmp_path / "income_spi_2022_23.pkl"
+    stale_metadata = {
+        **income_module.INCOME_MODEL_METADATA,
+        "spi_release_name": "spi_2020_21",
+        "spi_tab_filename": "put2021uk.tab",
+    }
+    with cache.open("wb") as f:
+        pickle.dump(
+            {
+                "model": SimpleNamespace(
+                    imputed_variables=list(income_module.IMPUTATIONS)
+                ),
+                "input_columns": income_module.PREDICTORS,
+                "metadata": stale_metadata,
+            },
+            f,
+        )
+
+    sentinel = object()
+    monkeypatch.setattr(income_module, "INCOME_MODEL_PATH", cache)
+    monkeypatch.setattr(income_module, "save_imputation_models", lambda: sentinel)
+
+    assert income_module.create_income_model() is sentinel
+
+
+def test_income_model_cache_accepts_current_spi_release(tmp_path, monkeypatch):
+    from policyengine_uk_data.datasets.imputations import income as income_module
+
+    cache = tmp_path / "income_spi_2022_23.pkl"
+    with cache.open("wb") as f:
+        pickle.dump(
+            {
+                "model": SimpleNamespace(
+                    imputed_variables=list(income_module.IMPUTATIONS)
+                ),
+                "input_columns": income_module.PREDICTORS,
+                "metadata": income_module.INCOME_MODEL_METADATA,
+            },
+            f,
+        )
+
+    monkeypatch.setattr(income_module, "INCOME_MODEL_PATH", cache)
+    monkeypatch.setattr(
+        income_module,
+        "save_imputation_models",
+        lambda: pytest.fail("current SPI release cache should be reused"),
+    )
+
+    assert income_module.create_income_model().metadata == (
+        income_module.INCOME_MODEL_METADATA
+    )

policyengine_uk_data/utils/qrf.py

@@ -39,6 +39,7 @@ def __init__(self, file_path: str = None):
                 data = pickle.load(f)
                 self.model = data["model"]
                 self.input_columns = data["input_columns"]
+                self.metadata = data.get("metadata", {})
 
     def fit(self, X, y):
         """
@@ -74,4 +75,11 @@ def save(self, file_path: str):
             file_path: Path where model should be saved.
         """
         with open(file_path, "wb") as f:
-            pickle.dump({"model": self.model, "input_columns": self.input_columns}, f)
+            pickle.dump(
+                {
+                    "model": self.model,
+                    "input_columns": self.input_columns,
+                    "metadata": getattr(self, "metadata", {}),
+                },
+                f,
+            )